-- SPDX-License-Identifier: MPL-2.0

{- | Fitted preprocessing and index-based splits. Fit on training rows before
transforming validation/test rows. Unknown categories are rejected unless
explicitly mapped to an all-zero vector. Sample quantiles interpolate between
order statistics; constant features use scale one and transform to zero.
-}
module HMath.MachineLearning.Preprocessing (
    Scaling (..),
    Scaler,
    fitScaler,
    transform,
    inverseTransform,
    QuantileTransformer,
    fitQuantileTransformer,
    quantileTransform,
    inverseQuantileTransform,
    PowerTransform (..),
    powerTransform,
    inversePowerTransform,
    Imputation (..),
    Imputer,
    fitImputer,
    impute,
    Encoder,
    UnknownCategory (..),
    fitEncoder,
    labelEncode,
    labelDecode,
    oneHotEncode,
    Fold (..),
    trainTestSplit,
    kFold,
    stratifiedKFold,
    gridSearch,
    randomSearch,
) where

import Data.List (nub, sort, sortOn, transpose)
import Data.Word (Word64)
import HMath.Error
import HMath.MachineLearning.Internal
import qualified HMath.Random as R
import qualified HMath.Scalar as S

-- | Population z-score, unit-range, or median/IQR feature scaling.
data Scaling = Standardize | MinMax | RobustScale deriving (Scaling -> Scaling -> Bool
(Scaling -> Scaling -> Bool)
-> (Scaling -> Scaling -> Bool) -> Eq Scaling
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Scaling -> Scaling -> Bool
== :: Scaling -> Scaling -> Bool
$c/= :: Scaling -> Scaling -> Bool
/= :: Scaling -> Scaling -> Bool
Eq, Int -> Scaling -> ShowS
[Scaling] -> ShowS
Scaling -> String
(Int -> Scaling -> ShowS)
-> (Scaling -> String) -> ([Scaling] -> ShowS) -> Show Scaling
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Scaling -> ShowS
showsPrec :: Int -> Scaling -> ShowS
$cshow :: Scaling -> String
show :: Scaling -> String
$cshowList :: [Scaling] -> ShowS
showList :: [Scaling] -> ShowS
Show)

-- | Immutable training-feature centers and nonzero scales; constant columns use scale one.
data Scaler = Scaler [Double] [Double] deriving (Scaler -> Scaler -> Bool
(Scaler -> Scaler -> Bool)
-> (Scaler -> Scaler -> Bool) -> Eq Scaler
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Scaler -> Scaler -> Bool
== :: Scaler -> Scaler -> Bool
$c/= :: Scaler -> Scaler -> Bool
/= :: Scaler -> Scaler -> Bool
Eq, Int -> Scaler -> ShowS
[Scaler] -> ShowS
Scaler -> String
(Int -> Scaler -> ShowS)
-> (Scaler -> String) -> ([Scaler] -> ShowS) -> Show Scaler
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Scaler -> ShowS
showsPrec :: Int -> Scaler -> ShowS
$cshow :: Scaler -> String
show :: Scaler -> String
$cshowList :: [Scaler] -> ShowS
showList :: [Scaler] -> ShowS
Show)

quantile :: Double -> [Double] -> Double
quantile :: Double -> [Double] -> Double
quantile Double
q [Double]
xs =
    let s :: [Double]
s = [Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
xs; z :: Double
z = Double
q Double -> Double -> Double
forall a. Num a => a -> a -> a
* 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1); i :: Int
i = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor Double
z; a :: Double
a = Double
z Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
     in (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
a) Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
s [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
a Double -> Double -> Double
forall a. Num a => a -> a -> a
* [Double]
s [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | Estimate one center and scale per training feature. Empty/ragged or nonfinite data fail.
fitScaler :: Scaling -> Dataset -> Either MathError Scaler
fitScaler :: Scaling -> Dataset -> Either MathError Scaler
fitScaler Scaling
method Dataset
xs = do
    _ <- Dataset -> Either MathError (Int, Int)
dataset Dataset
xs
    let fit [Double]
col = case Scaling
method of
            Scaling
Standardize ->
                let m :: Double
m = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
col Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ 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]
col)
                 in (Double
m, Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
v <- [Double]
col] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ 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]
col)))
            Scaling
MinMax -> ([Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Double]
col, [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Double]
col Double -> Double -> Double
forall a. Num a => a -> a -> a
- [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Double]
col)
            Scaling
RobustScale -> (Double -> [Double] -> Double
quantile Double
0.5 [Double]
col, Double -> [Double] -> Double
quantile Double
0.75 [Double]
col Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double -> [Double] -> Double
quantile Double
0.25 [Double]
col)
        pairs = ([Double] -> (Double, Double)) -> Dataset -> [(Double, Double)]
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> (Double, Double)
fit (Dataset -> Dataset
forall a. [[a]] -> [[a]]
transpose Dataset
xs)
        origins = ((Double, Double) -> Double) -> [(Double, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Double) -> Double
forall a b. (a, b) -> a
fst [(Double, Double)]
pairs
        scales = ((Double, Double) -> Double) -> [(Double, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(Double
_, Double
s) -> if Double
s Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
1 else Double
s) [(Double, Double)]
pairs
    checked (origins ++ scales)
    Right (Scaler origins scales)

-- | Apply training centers/scales to rows with the original feature count; no state is updated.
transform :: Scaler -> Dataset -> Either MathError Dataset
transform :: Scaler -> Dataset -> Either MathError Dataset
transform (Scaler [Double]
origins [Double]
scales) Dataset
xs = do
    Int -> Dataset -> Either MathError ()
queries ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
origins) Dataset
xs
    let ys :: Dataset
ys = ([Double] -> [Double]) -> Dataset -> Dataset
forall a b. (a -> b) -> [a] -> [b]
map (\[Double]
r -> (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Fractional a => a -> a -> a
(/) ([Double] -> [Double] -> [Double]
sub [Double]
r [Double]
origins) [Double]
scales) Dataset
xs
    [Double] -> Either MathError ()
checked (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
ys)
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right Dataset
ys

-- | Undo a fitted affine scaler, including its constant-feature convention.
inverseTransform :: Scaler -> Dataset -> Either MathError Dataset
inverseTransform :: Scaler -> Dataset -> Either MathError Dataset
inverseTransform (Scaler [Double]
origins [Double]
scales) Dataset
xs = do
    Int -> Dataset -> Either MathError ()
queries ([Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
origins) Dataset
xs
    let ys :: Dataset
ys = ([Double] -> [Double]) -> Dataset -> Dataset
forall a b. (a -> b) -> [a] -> [b]
map (\[Double]
r -> [Double] -> [Double] -> [Double]
add [Double]
origins ((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]
r [Double]
scales)) Dataset
xs
    [Double] -> Either MathError ()
checked (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
ys)
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right Dataset
ys

-- | Sorted training samples retained separately for each feature.
newtype QuantileTransformer = QT Dataset deriving (QuantileTransformer -> QuantileTransformer -> Bool
(QuantileTransformer -> QuantileTransformer -> Bool)
-> (QuantileTransformer -> QuantileTransformer -> Bool)
-> Eq QuantileTransformer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: QuantileTransformer -> QuantileTransformer -> Bool
== :: QuantileTransformer -> QuantileTransformer -> Bool
$c/= :: QuantileTransformer -> QuantileTransformer -> Bool
/= :: QuantileTransformer -> QuantileTransformer -> Bool
Eq, Int -> QuantileTransformer -> ShowS
[QuantileTransformer] -> ShowS
QuantileTransformer -> String
(Int -> QuantileTransformer -> ShowS)
-> (QuantileTransformer -> String)
-> ([QuantileTransformer] -> ShowS)
-> Show QuantileTransformer
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> QuantileTransformer -> ShowS
showsPrec :: Int -> QuantileTransformer -> ShowS
$cshow :: QuantileTransformer -> String
show :: QuantileTransformer -> String
$cshowList :: [QuantileTransformer] -> ShowS
showList :: [QuantileTransformer] -> ShowS
Show)

-- | Fit empirical feature quantiles using training rows only.
fitQuantileTransformer :: Dataset -> Either MathError QuantileTransformer
fitQuantileTransformer :: Dataset -> Either MathError QuantileTransformer
fitQuantileTransformer Dataset
xs = Dataset -> Either MathError (Int, Int)
dataset Dataset
xs Either MathError (Int, Int)
-> Either MathError QuantileTransformer
-> Either MathError QuantileTransformer
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> QuantileTransformer -> Either MathError QuantileTransformer
forall a b. b -> Either a b
Right (Dataset -> QuantileTransformer
QT (([Double] -> [Double]) -> Dataset -> Dataset
forall a b. (a -> b) -> [a] -> [b]
map [Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort (Dataset -> Dataset
forall a. [[a]] -> [[a]]
transpose Dataset
xs)))

-- | Map values to [0,1]. Ties use their average rank and values outside training extrema clamp.
quantileTransform :: QuantileTransformer -> Dataset -> Either MathError Dataset
quantileTransform :: QuantileTransformer -> Dataset -> Either MathError Dataset
quantileTransform (QT Dataset
cols) Dataset
xs = do
    Int -> Dataset -> Either MathError ()
queries (Dataset -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Dataset
cols) Dataset
xs
    let rank :: [a] -> a -> a
rank [a]
s a
v
            | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = a
0.5
            | a
v a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
0 = a
0
            | a
v a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> [a] -> a
forall a. HasCallStack => [a] -> a
last [a]
s = a
1
            | Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
equal) =
                [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral [Int]
equal) a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
equal) 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]
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            | Bool
otherwise =
                let j :: Int
j = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
v) [a]
s); lo :: a
lo = [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1); hi :: a
hi = [a]
s [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
j
                 in let magnitude :: a
magnitude = a -> a -> a
forall a. Ord a => a -> a -> a
max (a -> a
forall a. Num a => a -> a
abs a
lo) (a -> a
forall a. Num a => a -> a
abs a
hi)
                        fraction :: a
fraction =
                            if a
magnitude a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
0 else (a
v a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
magnitude a -> a -> a
forall a. Num a => a -> a -> a
- a
lo a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
magnitude) a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
hi a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
magnitude a -> a -> a
forall a. Num a => a -> a -> a
- a
lo a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
magnitude)
                     in (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) a -> a -> a
forall a. Num a => a -> a -> a
+ a
fraction) 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]
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
          where
            equal :: [Int]
equal = [Int
i | (Int
i, a
z) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [a]
s, a
z a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v]
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right [[[Double] -> Double -> Double
forall {a}. (Ord a, Fractional a) => [a] -> a -> a
rank [Double]
c Double
v | ([Double]
c, Double
v) <- Dataset -> [Double] -> [([Double], Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip Dataset
cols [Double]
r] | [Double]
r <- Dataset
xs]

-- | Interpolate training order statistics at probabilities in [0,1].
inverseQuantileTransform :: QuantileTransformer -> Dataset -> Either MathError Dataset
inverseQuantileTransform :: QuantileTransformer -> Dataset -> Either MathError Dataset
inverseQuantileTransform (QT Dataset
cols) Dataset
xs = do
    Int -> Dataset -> Either MathError ()
queries (Dataset -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Dataset
cols) Dataset
xs
    Bool -> String -> Either MathError ()
require ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 Bool -> Bool -> Bool
&& Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
1) (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
xs)) String
"quantiles must be in [0,1]"
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right [[Double -> [Double] -> Double
quantile Double
q [Double]
c | ([Double]
c, Double
q) <- Dataset -> [Double] -> [([Double], Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip Dataset
cols [Double]
r] | [Double]
r <- Dataset
xs]

-- | Box-Cox or Yeo-Johnson with a caller-selected finite power lambda; lambda is not estimated.
data PowerTransform = BoxCox Double | YeoJohnson Double deriving (PowerTransform -> PowerTransform -> Bool
(PowerTransform -> PowerTransform -> Bool)
-> (PowerTransform -> PowerTransform -> Bool) -> Eq PowerTransform
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PowerTransform -> PowerTransform -> Bool
== :: PowerTransform -> PowerTransform -> Bool
$c/= :: PowerTransform -> PowerTransform -> Bool
/= :: PowerTransform -> PowerTransform -> Bool
Eq, Int -> PowerTransform -> ShowS
[PowerTransform] -> ShowS
PowerTransform -> String
(Int -> PowerTransform -> ShowS)
-> (PowerTransform -> String)
-> ([PowerTransform] -> ShowS)
-> Show PowerTransform
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PowerTransform -> ShowS
showsPrec :: Int -> PowerTransform -> ShowS
$cshow :: PowerTransform -> String
show :: PowerTransform -> String
$cshowList :: [PowerTransform] -> ShowS
showList :: [PowerTransform] -> ShowS
Show)

-- | Apply the selected power to features. Box-Cox requires positive values; Yeo-Johnson allows signs.
powerTransform :: PowerTransform -> Dataset -> Either MathError Dataset
powerTransform :: PowerTransform -> Dataset -> Either MathError Dataset
powerTransform PowerTransform
method Dataset
xs = do
    _ <- Dataset -> Either MathError (Int, Int)
dataset Dataset
xs
    let l = case PowerTransform
method of BoxCox Double
z -> Double
z; YeoJohnson Double
z -> Double
z
    checked [l]
    case method of
        BoxCox Double
_ -> Bool -> String -> Either MathError ()
require ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0) (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
xs)) String
"Box-Cox requires positive values"
        PowerTransform
_ -> () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let quotient a
power a
logarithm =
            if a
power a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
                then a -> Either MathError a
forall a b. b -> Either a b
Right a
logarithm
                else
                    if a
power a -> a -> a
forall a. Num a => a -> a -> a
* a
logarithm a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
                        then a -> Either MathError a
forall a b. b -> Either a b
Right a
logarithm
                        else (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
power) (a -> a) -> Either MathError a -> Either MathError a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Either MathError a
forall a. RealFloat a => a -> Either MathError a
S.expm1 (a
power a -> a -> a
forall a. Num a => a -> a -> a
* a
logarithm)
        one Double
x = case PowerTransform
method of
            BoxCox Double
_ -> Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
quotient Double
l (Double -> Double
forall a. Floating a => a -> a
log Double
x)
            YeoJohnson Double
_
                | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 -> Double -> Either MathError Double
forall a. RealFloat a => a -> Either MathError a
S.log1p Double
x Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
quotient Double
l
                | Bool
otherwise -> Double -> Double
forall a. Num a => a -> a
negate (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Double -> Either MathError Double
forall a. RealFloat a => a -> Either MathError a
S.log1p (-Double
x) Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
quotient (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
l))
    ys <- traverse (traverse one) xs

    checked (concat ys)
    Right ys

-- | Invert the selected power, rejecting values outside the transform range.
inversePowerTransform :: PowerTransform -> Dataset -> Either MathError Dataset
inversePowerTransform :: PowerTransform -> Dataset -> Either MathError Dataset
inversePowerTransform PowerTransform
method Dataset
xs = do
    _ <- Dataset -> Either MathError (Int, Int)
dataset Dataset
xs
    let l = case PowerTransform
method of BoxCox Double
z -> Double
z; YeoJohnson Double
z -> Double
z
    checked [l]
    let scaledLog a
power a
value =
            if a
power a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 Bool -> Bool -> Bool
|| a
power a -> a -> a
forall a. Num a => a -> a -> a
* a
value a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0
                then a -> Either MathError a
forall a b. b -> Either a b
Right a
value
                else (a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
power) (a -> a) -> Either MathError a -> Either MathError a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Either MathError a
forall a. RealFloat a => a -> Either MathError a
S.log1p (a
power a -> a -> a
forall a. Num a => a -> a -> a
* a
value)
        one Double
x = case PowerTransform
method of
            BoxCox Double
_ -> Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
scaledLog Double
l Double
x
            YeoJohnson Double
_
                | Double
x Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
0 -> Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
scaledLog Double
l Double
x Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Double -> Either MathError Double
forall a. RealFloat a => a -> Either MathError a
S.expm1
                | Bool
otherwise -> Double -> Double
forall a. Num a => a -> a
negate (Double -> Double)
-> Either MathError Double -> Either MathError Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Double -> Double -> Either MathError Double
forall {a}. RealFloat a => a -> a -> Either MathError a
scaledLog (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
l) (-Double
x) Either MathError Double
-> (Double -> Either MathError Double) -> Either MathError Double
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Double -> Either MathError Double
forall a. RealFloat a => a -> Either MathError a
S.expm1)
    ys <- traverse (traverse one) xs
    checked (concat ys)
    Right ys

-- | Training-column mean, median, or a finite constant for NaN-valued observations.
data Imputation = MeanImpute | MedianImpute | ConstantImpute Double deriving (Imputation -> Imputation -> Bool
(Imputation -> Imputation -> Bool)
-> (Imputation -> Imputation -> Bool) -> Eq Imputation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Imputation -> Imputation -> Bool
== :: Imputation -> Imputation -> Bool
$c/= :: Imputation -> Imputation -> Bool
/= :: Imputation -> Imputation -> Bool
Eq, Int -> Imputation -> ShowS
[Imputation] -> ShowS
Imputation -> String
(Int -> Imputation -> ShowS)
-> (Imputation -> String)
-> ([Imputation] -> ShowS)
-> Show Imputation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Imputation -> ShowS
showsPrec :: Int -> Imputation -> ShowS
$cshow :: Imputation -> String
show :: Imputation -> String
$cshowList :: [Imputation] -> ShowS
showList :: [Imputation] -> ShowS
Show)

-- | Fitted replacement values; an imputer never learns from rows passed to impute.
newtype Imputer = Imputer [Double] deriving (Imputer -> Imputer -> Bool
(Imputer -> Imputer -> Bool)
-> (Imputer -> Imputer -> Bool) -> Eq Imputer
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Imputer -> Imputer -> Bool
== :: Imputer -> Imputer -> Bool
$c/= :: Imputer -> Imputer -> Bool
/= :: Imputer -> Imputer -> Bool
Eq, Int -> Imputer -> ShowS
[Imputer] -> ShowS
Imputer -> String
(Int -> Imputer -> ShowS)
-> (Imputer -> String) -> ([Imputer] -> ShowS) -> Show Imputer
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Imputer -> ShowS
showsPrec :: Int -> Imputer -> ShowS
$cshow :: Imputer -> String
show :: Imputer -> String
$cshowList :: [Imputer] -> ShowS
showList :: [Imputer] -> ShowS
Show)

-- | Treat NaN as missing and reject infinities. All-missing columns require ConstantImpute.
fitImputer :: Imputation -> Dataset -> Either MathError Imputer
fitImputer :: Imputation -> Dataset -> Either MathError Imputer
fitImputer Imputation
method Dataset
xs = do
    Bool -> String -> Either MathError ()
require
        (Bool -> Bool
not (Dataset -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Dataset
xs) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (Dataset -> [Double]
forall a. [[a]] -> [a]
firstRow Dataset
xs)) Bool -> Bool -> Bool
&& ([Double] -> Bool) -> Dataset -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Dataset -> [Double]
forall a. [[a]] -> [a]
firstRow Dataset
xs)) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) Dataset
xs)
        String
"imputer requires rectangular data"
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
xs))) String
"imputer rejects infinities"
    vals <- ([Double] -> Either MathError Double)
-> Dataset -> Either MathError [Double]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse [Double] -> Either MathError Double
fit (Dataset -> Dataset
forall a. [[a]] -> [[a]]
transpose Dataset
xs)
    checked vals
    Right (Imputer vals)
  where
    fit :: [Double] -> Either MathError Double
fit [Double]
col =
        let good :: [Double]
good = (Double -> Bool) -> [Double] -> [Double]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Double -> Bool) -> Double -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN) [Double]
col
         in case Imputation
method of
                ConstantImpute Double
c -> [Double] -> Either MathError ()
checked [Double
c] Either MathError ()
-> Either MathError Double -> Either MathError Double
forall a b.
Either MathError a -> Either MathError b -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
c
                Imputation
_ -> do
                    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
good)) String
"entirely missing training feature"
                    Double -> Either MathError Double
forall a b. b -> Either a b
Right
                        (case Imputation
method of Imputation
MeanImpute -> [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 -> Double
forall a. Fractional a => a -> a -> a
/ 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]
good)) [Double]
good); Imputation
_ -> Double -> [Double] -> Double
quantile Double
0.5 [Double]
good)

-- | Replace NaN with fitted training values while preserving observed entries and row order.
impute :: Imputer -> Dataset -> Either MathError Dataset
impute :: Imputer -> Dataset -> Either MathError Dataset
impute (Imputer [Double]
vals) Dataset
xs = do
    Bool -> String -> Either MathError ()
require (([Double] -> Bool) -> Dataset -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
vals) (Int -> Bool) -> ([Double] -> Int) -> [Double] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) Dataset
xs) String
"imputation feature mismatch"
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ((Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite (Dataset -> [Double]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat Dataset
xs))) String
"imputer rejects infinities"
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right [(Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Double
x Double
v -> if Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x then Double
v else Double
x) [Double]
r [Double]
vals | [Double]
r <- Dataset
xs]

-- | Sorted distinct training categories with stable zero-based integer codes.
newtype Encoder a = Encoder [a] deriving (Encoder a -> Encoder a -> Bool
(Encoder a -> Encoder a -> Bool)
-> (Encoder a -> Encoder a -> Bool) -> Eq (Encoder a)
forall a. Eq a => Encoder a -> Encoder a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Encoder a -> Encoder a -> Bool
== :: Encoder a -> Encoder a -> Bool
$c/= :: forall a. Eq a => Encoder a -> Encoder a -> Bool
/= :: Encoder a -> Encoder a -> Bool
Eq, Int -> Encoder a -> ShowS
[Encoder a] -> ShowS
Encoder a -> String
(Int -> Encoder a -> ShowS)
-> (Encoder a -> String)
-> ([Encoder a] -> ShowS)
-> Show (Encoder a)
forall a. Show a => Int -> Encoder a -> ShowS
forall a. Show a => [Encoder a] -> ShowS
forall a. Show a => Encoder a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Encoder a -> ShowS
showsPrec :: Int -> Encoder a -> ShowS
$cshow :: forall a. Show a => Encoder a -> String
show :: Encoder a -> String
$cshowList :: forall a. Show a => [Encoder a] -> ShowS
showList :: [Encoder a] -> ShowS
Show)

-- | Whether a previously unseen category fails or produces an all-zero one-hot row.
data UnknownCategory = RejectUnknown | ZeroUnknown deriving (UnknownCategory -> UnknownCategory -> Bool
(UnknownCategory -> UnknownCategory -> Bool)
-> (UnknownCategory -> UnknownCategory -> Bool)
-> Eq UnknownCategory
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: UnknownCategory -> UnknownCategory -> Bool
== :: UnknownCategory -> UnknownCategory -> Bool
$c/= :: UnknownCategory -> UnknownCategory -> Bool
/= :: UnknownCategory -> UnknownCategory -> Bool
Eq, Int -> UnknownCategory -> ShowS
[UnknownCategory] -> ShowS
UnknownCategory -> String
(Int -> UnknownCategory -> ShowS)
-> (UnknownCategory -> String)
-> ([UnknownCategory] -> ShowS)
-> Show UnknownCategory
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> UnknownCategory -> ShowS
showsPrec :: Int -> UnknownCategory -> ShowS
$cshow :: UnknownCategory -> String
show :: UnknownCategory -> String
$cshowList :: [UnknownCategory] -> ShowS
showList :: [UnknownCategory] -> ShowS
Show)

-- | Learn sorted category levels from a nonempty training label vector.
fitEncoder :: (Ord a) => [a] -> Either MathError (Encoder a)
fitEncoder :: forall a. Ord a => [a] -> Either MathError (Encoder a)
fitEncoder [a]
xs = do Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs)) String
"encoder requires categories"; Encoder a -> Either MathError (Encoder a)
forall a b. b -> Either a b
Right ([a] -> Encoder a
forall a. [a] -> Encoder a
Encoder ([a] -> [a]
forall a. Ord a => [a] -> [a]
sort ([a] -> [a]
forall a. Eq a => [a] -> [a]
nub [a]
xs)))

-- | Map known categories to zero-based codes; unseen categories fail.
labelEncode :: (Eq a) => Encoder a -> [a] -> Either MathError [Int]
labelEncode :: forall a. Eq a => Encoder a -> [a] -> Either MathError [Int]
labelEncode (Encoder [a]
levels) =
    (a -> Either MathError Int) -> [a] -> Either MathError [Int]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse
        ( \a
x -> case [Int
i | (Int
i, a
y) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [a]
levels, a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y] of
            Int
i : [Int]
_ -> Int -> Either MathError Int
forall a b. b -> Either a b
Right Int
i
            [Int]
_ -> MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"ML: unknown category")
        )

-- | Recover categories from checked integer codes.
labelDecode :: Encoder a -> [Int] -> Either MathError [a]
labelDecode :: forall a. Encoder a -> [Int] -> Either MathError [a]
labelDecode (Encoder [a]
levels) [Int]
ix = do
    Bool -> String -> Either MathError ()
require ((Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\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
< [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
levels) [Int]
ix) String
"unknown encoded label"
    [a] -> Either MathError [a]
forall a b. b -> Either a b
Right ((Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ([a]
levels [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
ix)

-- | Emit one column per sorted training category under the explicit unknown-category policy.
oneHotEncode :: (Eq a) => UnknownCategory -> Encoder a -> [a] -> Either MathError Dataset
oneHotEncode :: forall a.
Eq a =>
UnknownCategory -> Encoder a -> [a] -> Either MathError Dataset
oneHotEncode UnknownCategory
policy (Encoder [a]
levels) [a]
xs = do
    Bool -> String -> Either MathError ()
require (UnknownCategory
policy UnknownCategory -> UnknownCategory -> Bool
forall a. Eq a => a -> a -> Bool
== UnknownCategory
ZeroUnknown Bool -> Bool -> Bool
|| (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> [a] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [a]
levels) [a]
xs) String
"unknown category"
    Dataset -> Either MathError Dataset
forall a b. b -> Either a b
Right [[if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
c then Double
1 else Double
0 | a
c <- [a]
levels] | a
x <- [a]
xs]

-- | Disjoint indices into the original dataset. No estimator is fitted here.
data Fold = Fold
    { Fold -> [Int]
trainingIndices :: [Int]
    -- ^ Zero-based original row indices assigned to training in this fold.
    , Fold -> [Int]
validationIndices :: [Int]
    -- ^ Zero-based original row indices held out for validation in this fold.
    }
    deriving (Fold -> Fold -> Bool
(Fold -> Fold -> Bool) -> (Fold -> Fold -> Bool) -> Eq Fold
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Fold -> Fold -> Bool
== :: Fold -> Fold -> Bool
$c/= :: Fold -> Fold -> Bool
/= :: Fold -> Fold -> Bool
Eq, Int -> Fold -> ShowS
[Fold] -> ShowS
Fold -> String
(Int -> Fold -> ShowS)
-> (Fold -> String) -> ([Fold] -> ShowS) -> Show Fold
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Fold -> ShowS
showsPrec :: Int -> Fold -> ShowS
$cshow :: Fold -> String
show :: Fold -> String
$cshowList :: [Fold] -> ShowS
showList :: [Fold] -> ShowS
Show)

permutation :: Word64 -> Int -> [Int]
permutation :: Word64 -> Int -> [Int]
permutation Word64
seed Int
n = ((Word64, Int) -> Int) -> [(Word64, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Word64, Int) -> Int
forall a b. (a, b) -> b
snd (((Word64, Int) -> Word64) -> [(Word64, Int)] -> [(Word64, Int)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Word64, Int) -> Word64
forall a b. (a, b) -> a
fst ([Word64] -> [Int] -> [(Word64, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> [Word64] -> [Word64]
forall a. Int -> [a] -> [a]
take Int
n (RNG -> [Word64]
stream (Generator -> Word64 -> RNG
R.seedRNG Generator
R.PCG32 Word64
seed))) [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]))
  where
    stream :: RNG -> [Word64]
stream RNG
r = let (Word64
v, RNG
r') = RNG -> (Word64, RNG)
R.nextWord64 RNG
r in Word64
v Word64 -> [Word64] -> [Word64]
forall a. a -> [a] -> [a]
: RNG -> [Word64]
stream RNG
r'

-- | Seeded split of n row indices. The fraction denotes test rows, rounded and clamped to [1,n-1].
trainTestSplit :: Word64 -> Double -> Int -> Either MathError Fold
trainTestSplit :: Word64 -> Double -> Int -> Either MathError Fold
trainTestSplit Word64
seed Double
fraction Int
n = do
    Bool -> String -> Either MathError ()
require (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Double -> Bool
finite Double
fraction Bool -> Bool -> Bool
&& Double
fraction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
0 Bool -> Bool -> Bool
&& Double
fraction Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1) String
"invalid train/test split"
    let p :: [Int]
p = Word64 -> Int -> [Int]
permutation Word64
seed Int
n; count :: Int
count = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (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
fraction Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n)))
    Fold -> Either MathError Fold
forall a b. b -> Either a b
Right ([Int] -> [Int] -> Fold
Fold (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
count [Int]
p) (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
count [Int]
p))

-- | Partition n shuffled indices into k nonempty validation folds; each row is held out exactly once.
kFold :: Word64 -> Int -> Int -> Either MathError [Fold]
kFold :: Word64 -> Int -> Int -> Either MathError [Fold]
kFold Word64
seed Int
k Int
n = do
    Bool -> String -> Either MathError ()
require (Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n) String
"fold count must be between two and sample count"
    let p :: [Int]
p = Word64 -> Int -> [Int]
permutation Word64
seed Int
n
    [Fold] -> Either MathError [Fold]
forall a b. b -> Either a b
Right
        [ let held :: [Int]
held = [Int
i | (Int
j, Int
i) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Int]
p, Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
f]
           in [Int] -> [Int] -> Fold
Fold ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
held) [Int]
p) [Int]
held
        | Int
f <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        ]

-- | Round-robin each shuffled class across k folds. Every class needs at least k observations.
stratifiedKFold :: (Ord a) => Word64 -> Int -> [a] -> Either MathError [Fold]
stratifiedKFold :: forall a. Ord a => Word64 -> Int -> [a] -> Either MathError [Fold]
stratifiedKFold Word64
seed Int
k [a]
labels = do
    Bool -> String -> Either MathError ()
require (Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 Bool -> Bool -> Bool
&& Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
labels) String
"invalid stratified fold count"
    let groups :: [[Int]]
groups = [[Int
i | (Int
i, a
y) <- [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [a]
labels, a
y a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
c] | a
c <- [a] -> [a]
forall a. Ord a => [a] -> [a]
sort ([a] -> [a]
forall a. Eq a => [a] -> [a]
nub [a]
labels)]
    Bool -> String -> Either MathError ()
require (([Int] -> Bool) -> [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
k) (Int -> Bool) -> ([Int] -> Int) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) [[Int]]
groups) String
"each class needs at least one member per fold"
    let ordered :: [[Int]]
ordered =
            [(Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
g [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) (Word64 -> Int -> [Int]
permutation (Word64
seed Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+ Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
j) ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
g)) | (Int
j, [Int]
g) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [[Int]]
groups]
        held :: Int -> [Int]
held Int
f = [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Int
i | (Int
j, Int
i) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Int]
g, Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
f] | [Int]
g <- [[Int]]
ordered]
    [Fold] -> Either MathError [Fold]
forall a b. b -> Either a b
Right [[Int] -> [Int] -> Fold
Fold ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` Int -> [Int]
held Int
f) [Int
0 .. [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
labels Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) (Int -> [Int]
held Int
f) | Int
f <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]

{- | Minimize a user-supplied validation loss, evaluated independently for every
candidate/fold. The callback receives only the split indices, so it can fit
preprocessing/model state inside each fold. Ties preserve candidate order.
-}
gridSearch ::
    [p] -> [Fold] -> (p -> Fold -> Either MathError Double) -> Either MathError (p, Double)
gridSearch :: forall p.
[p]
-> [Fold]
-> (p -> Fold -> Either MathError Double)
-> Either MathError (p, Double)
gridSearch [p]
candidates [Fold]
folds p -> Fold -> Either MathError Double
score = do
    Bool -> String -> Either MathError ()
require (Bool -> Bool
not ([p] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [p]
candidates) Bool -> Bool -> Bool
&& Bool -> Bool
not ([Fold] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Fold]
folds)) String
"empty parameter search"
    Bool -> String -> Either MathError ()
require ((Fold -> Bool) -> [Fold] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Fold -> Bool
valid [Fold]
folds) String
"search folds must have disjoint nonempty unique nonnegative indices"
    results <-
        (p -> Either MathError (p, Double))
-> [p] -> Either MathError [(p, Double)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse
            ( \p
p -> do
                costs <- (Fold -> Either MathError Double)
-> [Fold] -> Either MathError [Double]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (p -> Fold -> Either MathError Double
score p
p) [Fold]
folds
                checked costs
                let average = [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 -> Double
forall a. Fractional a => a -> a -> a
/ 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]
costs)) [Double]
costs)
                checked [average]
                Right (p, average)
            )
            [p]
candidates
    Right (results !! argmin (map snd results))
  where
    valid :: Fold -> Bool
valid Fold
f =
        let tr :: [Int]
tr = Fold -> [Int]
trainingIndices Fold
f; va :: [Int]
va = Fold -> [Int]
validationIndices Fold
f
         in Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
tr)
                Bool -> Bool -> Bool
&& Bool -> Bool
not ([Int] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Int]
va)
                Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) ([Int]
tr [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
va)
                Bool -> Bool -> Bool
&& [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
tr) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
tr
                Bool -> Bool -> Bool
&& [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
va) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
va
                Bool -> Bool -> Bool
&& (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
va) [Int]
tr

-- | Sample count distinct candidates from a finite set, then evaluate the supplied folds. Lower loss wins.
randomSearch ::
    Word64 ->
    Int ->
    [p] ->
    [Fold] ->
    (p -> Fold -> Either MathError Double) ->
    Either MathError (p, Double)
randomSearch :: forall p.
Word64
-> Int
-> [p]
-> [Fold]
-> (p -> Fold -> Either MathError Double)
-> Either MathError (p, Double)
randomSearch Word64
seed Int
count [p]
candidates [Fold]
folds p -> Fold -> Either MathError Double
score = do
    Bool -> String -> Either MathError ()
require
        (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= [p] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [p]
candidates)
        String
"random search budget exceeds finite candidate set"
    [p]
-> [Fold]
-> (p -> Fold -> Either MathError Double)
-> Either MathError (p, Double)
forall p.
[p]
-> [Fold]
-> (p -> Fold -> Either MathError Double)
-> Either MathError (p, Double)
gridSearch ((Int -> p) -> [Int] -> [p]
forall a b. (a -> b) -> [a] -> [b]
map ([p]
candidates [p] -> Int -> p
forall a. HasCallStack => [a] -> Int -> a
!!) (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
count (Word64 -> Int -> [Int]
permutation Word64
seed ([p] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [p]
candidates)))) [Fold]
folds p -> Fold -> Either MathError Double
score