{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

{- | Fit a separate base model per categorical value-combination, routing each
row to its segment at predict time; unseen or too-small segments fall back to a
global fit. Optional partial pooling (linear base) shrinks small segments.
-}
module DataFrame.Segmented (
    module DataFrame.Model,
    Segmented (..),
    segmented,
    segmentOn,
    pooled,
    Segment (..),
    SegmentedModel (..),
    SegmentFit (..),
) where

import Control.Exception (throw)
import Data.List (foldl', (\\))
import qualified Data.Map.Strict as M
import Data.Maybe (isJust)
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import DataFrame.Errors (DataFrameException (..))

import DataFrame.Featurize.Internal (featureNames, numericMatrix, targetDoubles)
import DataFrame.Internal.Column (
    Column (..),
    Columnable,
    bitmapTestBit,
    columnBitmap,
    hasElemType,
 )
import DataFrame.Internal.DataFrame (DataFrame, unsafeGetColumn)
import DataFrame.Internal.Expression (Expr (..))
import DataFrame.Internal.Types (SBool (..), sIntegral)
import DataFrame.LinearAlgebra (Matrix, gram, matVec, tMatVec)
import DataFrame.LinearAlgebra.Solve (choleskySolve, qrLeastSquares)
import DataFrame.LinearModel.Logistic (LogisticConfig)
import DataFrame.LinearModel.Regression (
    LinearConfig (..),
    LinearRegressor (..),
 )
import DataFrame.Model
import DataFrame.Operations.Core (nRows)
import DataFrame.Operations.Subset (columnToTextVec, exclude, rowsAtIndices)
import DataFrame.Operators ((.&&.), (.==.))
import DataFrame.SymbolicRegression (SRConfig)

{- | A base estimator @cfg@ wrapped to fit one model per categorical
value-combination. @segOn@ picks the columns ('Nothing' = auto-detect),
@segMinRows@ the smallest own-model segment, @segPool@ the pooling strength @λ@.
-}
data Segmented cfg = Segmented
    { forall cfg. Segmented cfg -> cfg
segBase :: !cfg
    , forall cfg. Segmented cfg -> Maybe [Text]
segOn :: !(Maybe [T.Text])
    , forall cfg. Segmented cfg -> Int
segMaxCard :: !Int
    , forall cfg. Segmented cfg -> Int
segMinRows :: !Int
    , forall cfg. Segmented cfg -> Double
segPool :: !Double
    }
    deriving (Segmented cfg -> Segmented cfg -> Bool
(Segmented cfg -> Segmented cfg -> Bool)
-> (Segmented cfg -> Segmented cfg -> Bool) -> Eq (Segmented cfg)
forall cfg. Eq cfg => Segmented cfg -> Segmented cfg -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall cfg. Eq cfg => Segmented cfg -> Segmented cfg -> Bool
== :: Segmented cfg -> Segmented cfg -> Bool
$c/= :: forall cfg. Eq cfg => Segmented cfg -> Segmented cfg -> Bool
/= :: Segmented cfg -> Segmented cfg -> Bool
Eq, Int -> Segmented cfg -> ShowS
[Segmented cfg] -> ShowS
Segmented cfg -> [Char]
(Int -> Segmented cfg -> ShowS)
-> (Segmented cfg -> [Char])
-> ([Segmented cfg] -> ShowS)
-> Show (Segmented cfg)
forall cfg. Show cfg => Int -> Segmented cfg -> ShowS
forall cfg. Show cfg => [Segmented cfg] -> ShowS
forall cfg. Show cfg => Segmented cfg -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall cfg. Show cfg => Int -> Segmented cfg -> ShowS
showsPrec :: Int -> Segmented cfg -> ShowS
$cshow :: forall cfg. Show cfg => Segmented cfg -> [Char]
show :: Segmented cfg -> [Char]
$cshowList :: forall cfg. Show cfg => [Segmented cfg] -> ShowS
showList :: [Segmented cfg] -> ShowS
Show)

-- | Wrap a base config with the defaults: auto-detect, cap 32, min 30 rows, no pooling.
segmented :: cfg -> Segmented cfg
segmented :: forall cfg. cfg -> Segmented cfg
segmented cfg
base = cfg -> Maybe [Text] -> Int -> Int -> Double -> Segmented cfg
forall cfg.
cfg -> Maybe [Text] -> Int -> Int -> Double -> Segmented cfg
Segmented cfg
base Maybe [Text]
forall a. Maybe a
Nothing Int
32 Int
30 Double
0

-- | Segment only on the named columns (each must be Text), overriding auto-detect.
segmentOn :: Segmented cfg -> [T.Text] -> Segmented cfg
segmentOn :: forall cfg. Segmented cfg -> [Text] -> Segmented cfg
segmentOn Segmented cfg
s [Text]
cols = Segmented cfg
s{segOn = Just cols}

-- | Set the pooling strength @λ@ (shrink segments toward the reference).
pooled :: Segmented cfg -> Double -> Segmented cfg
pooled :: forall cfg. Segmented cfg -> Double -> Segmented cfg
pooled Segmented cfg
s Double
lam = Segmented cfg
s{segPool = lam}

-- | One fitted segment: its categorical key, row count, and base model.
data Segment model = Segment
    { forall model. Segment model -> [Text]
segKey :: ![T.Text]
    , forall model. Segment model -> Int
segN :: !Int
    , forall model. Segment model -> model
segModel :: !model
    }
    deriving (Int -> Segment model -> ShowS
[Segment model] -> ShowS
Segment model -> [Char]
(Int -> Segment model -> ShowS)
-> (Segment model -> [Char])
-> ([Segment model] -> ShowS)
-> Show (Segment model)
forall model. Show model => Int -> Segment model -> ShowS
forall model. Show model => [Segment model] -> ShowS
forall model. Show model => Segment model -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall model. Show model => Int -> Segment model -> ShowS
showsPrec :: Int -> Segment model -> ShowS
$cshow :: forall model. Show model => Segment model -> [Char]
show :: Segment model -> [Char]
$cshowList :: forall model. Show model => [Segment model] -> ShowS
showList :: [Segment model] -> ShowS
Show)

{- | A fitted segmented model: the columns segmented on, the per-segment models
(ascending key order), the observed combinations that fell back (key + row
count), the global fallback model, and the compiled routing expression.
-}
data SegmentedModel a model = SegmentedModel
    { forall a model. SegmentedModel a model -> [Text]
smCatCols :: ![T.Text]
    , forall a model. SegmentedModel a model -> [Segment model]
smSegments :: ![Segment model]
    , forall a model. SegmentedModel a model -> [([Text], Int)]
smFellBack :: ![([T.Text], Int)]
    , forall a model. SegmentedModel a model -> model
smFallback :: !model
    , forall a model. SegmentedModel a model -> Expr a
smExpr :: !(Expr a)
    }
    deriving (Int -> SegmentedModel a model -> ShowS
[SegmentedModel a model] -> ShowS
SegmentedModel a model -> [Char]
(Int -> SegmentedModel a model -> ShowS)
-> (SegmentedModel a model -> [Char])
-> ([SegmentedModel a model] -> ShowS)
-> Show (SegmentedModel a model)
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
forall a model.
(Show model, Show a) =>
Int -> SegmentedModel a model -> ShowS
forall a model.
(Show model, Show a) =>
[SegmentedModel a model] -> ShowS
forall a model.
(Show model, Show a) =>
SegmentedModel a model -> [Char]
$cshowsPrec :: forall a model.
(Show model, Show a) =>
Int -> SegmentedModel a model -> ShowS
showsPrec :: Int -> SegmentedModel a model -> ShowS
$cshow :: forall a model.
(Show model, Show a) =>
SegmentedModel a model -> [Char]
show :: SegmentedModel a model -> [Char]
$cshowList :: forall a model.
(Show model, Show a) =>
[SegmentedModel a model] -> ShowS
showList :: [SegmentedModel a model] -> ShowS
Show)

{- | How a base estimator fits its per-segment models under pooling strength @λ@.
The default fits each segment independently and rejects @λ > 0@; the linear
instance overrides it with closed-form shrinkage. Every base model needs an instance.
-}
class (Fit cfg (Expr a)) => SegmentFit cfg a where
    -- | Fit the qualifying segments (each a numeric-only frame, in order).
    fitSegments :: cfg -> Double -> Expr a -> [DataFrame] -> [ModelOf cfg (Expr a)]
    fitSegments cfg
cfg Double
lam Expr a
target [DataFrame]
dfs
        | Double
lam Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = (DataFrame -> ModelOf cfg (Expr a))
-> [DataFrame] -> [ModelOf cfg (Expr a)]
forall a b. (a -> b) -> [a] -> [b]
map (cfg
-> Expr a
-> FrameFor (Expr a)
-> FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
forall cfg input.
(Fit cfg input,
 CheckFrame (FrameReq cfg input) (FrameFor input)) =>
cfg
-> input
-> FrameFor input
-> FitResult (FrameFor input) (ModelOf cfg input)
fit cfg
cfg Expr a
target) [DataFrame]
dfs
        | Bool
otherwise =
            [Char] -> [ModelOf cfg (Expr a)]
forall a. HasCallStack => [Char] -> a
error
                [Char]
"Segmented: pooling (lambda > 0) is not supported for this base model; use lambda = 0 or a linear base."

-- | Logistic segments support independent fitting (lambda = 0) only, for now.
instance (Columnable a, Ord a) => SegmentFit LogisticConfig a

-- | Symbolic-regression segments support independent fitting (lambda = 0) only.
instance SegmentFit SRConfig Double

instance
    ( Fit cfg (Expr a)
    , SegmentFit cfg a
    , Predict (ModelOf cfg (Expr a))
    , Prediction (ModelOf cfg (Expr a)) ~ Expr a
    , Columnable a
    ) =>
    Fit (Segmented cfg) (Expr a)
    where
    type ModelOf (Segmented cfg) (Expr a) = SegmentedModel a (ModelOf cfg (Expr a))
    type FrameReq (Segmented cfg) (Expr a) = 'AnyFrame
    fit :: CheckFrame
  (FrameReq (Segmented cfg) (Expr a)) (FrameFor (Expr a)) =>
Segmented cfg
-> Expr a
-> FrameFor (Expr a)
-> FitResult (FrameFor (Expr a)) (ModelOf (Segmented cfg) (Expr a))
fit = Segmented cfg
-> Expr a -> DataFrame -> SegmentedModel a (ModelOf cfg (Expr a))
Segmented cfg
-> Expr a
-> FrameFor (Expr a)
-> FitResult (FrameFor (Expr a)) (ModelOf (Segmented cfg) (Expr a))
forall cfg a.
(Fit cfg (Expr a), SegmentFit cfg a,
 Predict (ModelOf cfg (Expr a)),
 Prediction (ModelOf cfg (Expr a)) ~ Expr a, Columnable a) =>
Segmented cfg
-> Expr a -> DataFrame -> SegmentedModel a (ModelOf cfg (Expr a))
fitSegmented

instance Predict (SegmentedModel a model) where
    type Prediction (SegmentedModel a model) = Expr a
    predict :: SegmentedModel a model -> Prediction (SegmentedModel a model)
predict = SegmentedModel a model -> Expr a
SegmentedModel a model -> Prediction (SegmentedModel a model)
forall a model. SegmentedModel a model -> Expr a
smExpr

fitSegmented ::
    forall cfg a.
    ( Fit cfg (Expr a)
    , SegmentFit cfg a
    , Predict (ModelOf cfg (Expr a))
    , Prediction (ModelOf cfg (Expr a)) ~ Expr a
    , Columnable a
    ) =>
    Segmented cfg ->
    Expr a ->
    DataFrame ->
    SegmentedModel a (ModelOf cfg (Expr a))
fitSegmented :: forall cfg a.
(Fit cfg (Expr a), SegmentFit cfg a,
 Predict (ModelOf cfg (Expr a)),
 Prediction (ModelOf cfg (Expr a)) ~ Expr a, Columnable a) =>
Segmented cfg
-> Expr a -> DataFrame -> SegmentedModel a (ModelOf cfg (Expr a))
fitSegmented (Segmented cfg
base Maybe [Text]
mcols Int
maxCard Int
minRows Double
lam) Expr a
target DataFrame
df =
    ()
-> SegmentedModel a (ModelOf cfg (Expr a))
-> SegmentedModel a (ModelOf cfg (Expr a))
forall a b. a -> b -> b
seq (DataFrame -> [Text] -> Expr a -> ()
forall a. DataFrame -> [Text] -> Expr a -> ()
guardNumeric DataFrame
df [Text]
textCols Expr a
target) SegmentedModel a (ModelOf cfg (Expr a))
result
  where
    mTarget :: Maybe Text
mTarget = case Expr a
target of
        Col Text
n -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
n
        Expr a
_ -> Maybe Text
forall a. Maybe a
Nothing
    feats :: [Text]
feats = Expr a -> DataFrame -> [Text]
forall a. Expr a -> DataFrame -> [Text]
featureNames Expr a
target DataFrame
df
    textCols :: [Text]
textCols = [Text
c | Text
c <- [Text]
feats, DataFrame -> Text -> Bool
isTextCol DataFrame
df Text
c]
    catCols :: [Text]
catCols = DataFrame -> Maybe Text -> [Text] -> Maybe [Text] -> Int -> [Text]
resolveCatCols DataFrame
df Maybe Text
mTarget [Text]
textCols Maybe [Text]
mcols Int
maxCard
    numericFrame :: DataFrame -> DataFrame
numericFrame = [Text] -> DataFrame -> DataFrame
exclude [Text]
textCols
    globalM :: FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
globalM = cfg
-> Expr a
-> FrameFor (Expr a)
-> FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
forall cfg input.
(Fit cfg input,
 CheckFrame (FrameReq cfg input) (FrameFor input)) =>
cfg
-> input
-> FrameFor input
-> FitResult (FrameFor input) (ModelOf cfg input)
fit cfg
base Expr a
target (DataFrame -> DataFrame
numericFrame DataFrame
df)
    result :: SegmentedModel a (ModelOf cfg (Expr a))
result
        | [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
catCols =
            [Text]
-> [Segment (ModelOf cfg (Expr a))]
-> [([Text], Int)]
-> ModelOf cfg (Expr a)
-> Expr a
-> SegmentedModel a (ModelOf cfg (Expr a))
forall a model.
[Text]
-> [Segment model]
-> [([Text], Int)]
-> model
-> Expr a
-> SegmentedModel a model
SegmentedModel [] [] [] ModelOf cfg (Expr a)
FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
globalM (ModelOf cfg (Expr a) -> Prediction (ModelOf cfg (Expr a))
forall model. Predict model => model -> Prediction model
predict ModelOf cfg (Expr a)
FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
globalM)
        | Bool
otherwise =
            let d :: Int
d = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text]
feats [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Text]
textCols)
                floor' :: Int
floor' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
minRows (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                grouped :: [([Text], Vector Int)]
grouped = DataFrame -> [Text] -> [([Text], Vector Int)]
groupByKey DataFrame
df [Text]
catCols
                ([([Text], Vector Int)]
qualifying, [([Text], Vector Int)]
undersized) =
                    (([Text], Vector Int) -> Bool)
-> [([Text], Vector Int)]
-> ([([Text], Vector Int)], [([Text], Vector Int)])
forall b. (b -> Bool) -> [b] -> ([b], [b])
span' (\([Text]
_, Vector Int
ixs) -> Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ixs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
floor') [([Text], Vector Int)]
grouped
                qualFrames :: [DataFrame]
qualFrames =
                    [DataFrame -> DataFrame
numericFrame (Vector Int -> DataFrame -> DataFrame
rowsAtIndices Vector Int
ixs DataFrame
df) | ([Text]
_, Vector Int
ixs) <- [([Text], Vector Int)]
qualifying]
                segModels :: [ModelOf cfg (Expr a)]
segModels = cfg -> Double -> Expr a -> [DataFrame] -> [ModelOf cfg (Expr a)]
forall cfg a.
SegmentFit cfg a =>
cfg -> Double -> Expr a -> [DataFrame] -> [ModelOf cfg (Expr a)]
fitSegments cfg
base Double
lam Expr a
target [DataFrame]
qualFrames
                segments :: [Segment (ModelOf cfg (Expr a))]
segments =
                    (([Text], Vector Int)
 -> ModelOf cfg (Expr a) -> Segment (ModelOf cfg (Expr a)))
-> [([Text], Vector Int)]
-> [ModelOf cfg (Expr a)]
-> [Segment (ModelOf cfg (Expr a))]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
                        (\([Text]
k, Vector Int
ixs) ModelOf cfg (Expr a)
m -> [Text]
-> Int -> ModelOf cfg (Expr a) -> Segment (ModelOf cfg (Expr a))
forall model. [Text] -> Int -> model -> Segment model
Segment [Text]
k (Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ixs) ModelOf cfg (Expr a)
m)
                        [([Text], Vector Int)]
qualifying
                        [ModelOf cfg (Expr a)]
segModels
                fellBack :: [([Text], Int)]
fellBack = [([Text]
k, Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ixs) | ([Text]
k, Vector Int
ixs) <- [([Text], Vector Int)]
undersized]
                expr :: Expr a
expr = [Text]
-> [Segment (ModelOf cfg (Expr a))]
-> ModelOf cfg (Expr a)
-> Expr a
forall a model.
(Columnable a, Predict model, Prediction model ~ Expr a) =>
[Text] -> [Segment model] -> model -> Expr a
buildExpr [Text]
catCols [Segment (ModelOf cfg (Expr a))]
segments ModelOf cfg (Expr a)
FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
globalM
             in [Text]
-> [Segment (ModelOf cfg (Expr a))]
-> [([Text], Int)]
-> ModelOf cfg (Expr a)
-> Expr a
-> SegmentedModel a (ModelOf cfg (Expr a))
forall a model.
[Text]
-> [Segment model]
-> [([Text], Int)]
-> model
-> Expr a
-> SegmentedModel a model
SegmentedModel [Text]
catCols [Segment (ModelOf cfg (Expr a))]
segments [([Text], Int)]
fellBack ModelOf cfg (Expr a)
FitResult (FrameFor (Expr a)) (ModelOf cfg (Expr a))
globalM Expr a
expr

-- | 'span' over a predicate that need not hold contiguously (a filter partition).
span' :: (b -> Bool) -> [b] -> ([b], [b])
span' :: forall b. (b -> Bool) -> [b] -> ([b], [b])
span' b -> Bool
p [b]
xs = ((b -> Bool) -> [b] -> [b]
forall a. (a -> Bool) -> [a] -> [a]
filter b -> Bool
p [b]
xs, (b -> Bool) -> [b] -> [b]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (b -> Bool) -> b -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
p) [b]
xs)

{- | Compile the routing: a right-folded @If@ ladder of @key == value@ conjuncts,
ending in the fallback model's prediction. Keys are disjoint, so order is
immaterial.
-}
buildExpr ::
    (Columnable a, Predict model, Prediction model ~ Expr a) =>
    [T.Text] ->
    [Segment model] ->
    model ->
    Expr a
buildExpr :: forall a model.
(Columnable a, Predict model, Prediction model ~ Expr a) =>
[Text] -> [Segment model] -> model -> Expr a
buildExpr [Text]
catCols [Segment model]
segs model
fallback =
    (Segment model -> Expr a -> Expr a)
-> Expr a -> [Segment model] -> Expr a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
        (\(Segment [Text]
key Int
_ model
m) Expr a
acc -> Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If ([Text] -> [Text] -> Expr Bool
keyCond [Text]
catCols [Text]
key) (model -> Prediction model
forall model. Predict model => model -> Prediction model
predict model
m) Expr a
acc)
        (model -> Prediction model
forall model. Predict model => model -> Prediction model
predict model
fallback)
        [Segment model]
segs

-- | @col1 == v1 && col2 == v2 && ...@ for a segment's key.
keyCond :: [T.Text] -> [T.Text] -> Expr Bool
keyCond :: [Text] -> [Text] -> Expr Bool
keyCond [Text]
catCols [Text]
vals =
    (Expr Bool -> Expr Bool -> Expr Bool) -> [Expr Bool] -> Expr Bool
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Expr Bool -> Expr Bool -> Expr Bool
(.&&.) [(Text -> Expr Text
forall a. Columnable a => Text -> Expr a
Col Text
c :: Expr T.Text) Expr Text -> Expr Text -> Expr Bool
forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.==. Text -> Expr Text
forall a. Columnable a => a -> Expr a
Lit Text
v | (Text
c, Text
v) <- [Text] -> [Text] -> [(Text, Text)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
catCols [Text]
vals]

{- | The Text feature columns to segment on, chosen from the frame's Text features
@textCols@. An explicit list is validated to be all-Text; auto-detect keeps Text
columns with at most @maxCard@ distinct values.
-}
resolveCatCols ::
    DataFrame -> Maybe T.Text -> [T.Text] -> Maybe [T.Text] -> Int -> [T.Text]
resolveCatCols :: DataFrame -> Maybe Text -> [Text] -> Maybe [Text] -> Int -> [Text]
resolveCatCols DataFrame
df Maybe Text
mTarget [Text]
textCols Maybe [Text]
mcols Int
maxCard = case Maybe [Text]
mcols of
    Just [Text]
cols ->
        let bad :: [Text]
bad = (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Text
c -> Bool -> Bool
not (DataFrame -> Text -> Bool
isTextCol DataFrame
df Text
c) Bool -> Bool -> Bool
|| Text -> Maybe Text
forall a. a -> Maybe a
Just Text
c Maybe Text -> Maybe Text -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe Text
mTarget) [Text]
cols
         in if [Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
bad
                then [Text]
cols
                else
                    [Char] -> [Text]
forall a. HasCallStack => [Char] -> a
error
                        ( [Char]
"Segmented: segmentOn columns must be Text features (not the target); invalid: "
                            [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Text] -> [Char]
forall a. Show a => a -> [Char]
show [Text]
bad
                        )
    Maybe [Text]
Nothing -> [Text
c | Text
c <- [Text]
textCols, DataFrame -> Text -> Int
distinctCount DataFrame
df Text
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
maxCard]

isTextCol :: DataFrame -> T.Text -> Bool
isTextCol :: DataFrame -> Text -> Bool
isTextCol DataFrame
df Text
c = forall a. Columnable a => Column -> Bool
hasElemType @T.Text (Text -> DataFrame -> Column
unsafeGetColumn Text
c DataFrame
df)

distinctCount :: DataFrame -> T.Text -> Int
distinctCount :: DataFrame -> Text -> Int
distinctCount DataFrame
df Text
c =
    Set Text -> Int
forall a. Set a -> Int
Set.size ([Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList (Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList (Column -> Vector Text
columnToTextVec (Text -> DataFrame -> Column
unsafeGetColumn Text
c DataFrame
df))))

{- | Reject feature columns that are neither Text (dropped\/segmented) nor
non-null 'Double', naming each with its fix — clearer than the base fitter's
raw type mismatch.
-}
guardNumeric :: DataFrame -> [T.Text] -> Expr a -> ()
guardNumeric :: forall a. DataFrame -> [Text] -> Expr a -> ()
guardNumeric DataFrame
df [Text]
textCols Expr a
target =
    case [(Text, [Char])]
problems of
        [] -> ()
        [(Text, [Char])]
ps ->
            [Char] -> ()
forall a. HasCallStack => [Char] -> a
error
                ( [Char]
"Segmented: unusable feature column(s):\n"
                    [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [[Char]] -> [Char]
unlines (((Text, [Char]) -> [Char]) -> [(Text, [Char])] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (Text, [Char]) -> [Char]
fmt [(Text, [Char])]
ps)
                )
  where
    problems :: [(Text, [Char])]
problems =
        [ (Text
c, [Char]
r)
        | Text
c <- Expr a -> DataFrame -> [Text]
forall a. Expr a -> DataFrame -> [Text]
featureNames Expr a
target DataFrame
df [Text] -> [Text] -> [Text]
forall a. Eq a => [a] -> [a] -> [a]
\\ [Text]
textCols
        , Just [Char]
r <- [Column -> Maybe [Char]
forall {a}. IsString a => Column -> Maybe a
reason (Text -> DataFrame -> Column
unsafeGetColumn Text
c DataFrame
df)]
        ]
    fmt :: (Text, [Char]) -> [Char]
fmt (Text
c, [Char]
r) = [Char]
"  " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
T.unpack Text
c [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
": " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
r
    reason :: Column -> Maybe a
reason Column
col
        | Maybe Bitmap -> Bool
forall a. Maybe a -> Bool
isJust (Column -> Maybe Bitmap
columnBitmap Column
col) =
            a -> Maybe a
forall a. a -> Maybe a
Just
                a
"has missing values — drop them (filterJust / filterAllJust) or model missingness explicitly; imputing risks train/inference skew"
        | Column -> Bool
isIntegralCol Column
col =
            a -> Maybe a
forall a. a -> Maybe a
Just a
"is an integer column — cast to Double with F.toDouble"
        | Bool -> Bool
not (forall a. Columnable a => Column -> Bool
hasElemType @Double Column
col) =
            a -> Maybe a
forall a. a -> Maybe a
Just
                a
"is not Double — convert to Double (numeric) or segment on it (categorical)"
        | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
    isIntegralCol :: Column -> Bool
isIntegralCol Column
col = case Column
col of
        UnboxedColumn Maybe Bitmap
_ (Vector a
_ :: VU.Vector b) -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @b of
            SBool (IntegralTypes a)
STrue -> Bool
True
            SBool (IntegralTypes a)
_ -> Bool
False
        Column
_ -> Bool
False

{- | Group row indices by their composite categorical key, dropping rows whose key
has a null in any segmented column (served by the fallback). Ascending key order.
-}
groupByKey :: DataFrame -> [T.Text] -> [([T.Text], VU.Vector Int)]
groupByKey :: DataFrame -> [Text] -> [([Text], Vector Int)]
groupByKey DataFrame
df [Text]
catCols =
    (([Text], [Int]) -> ([Text], Vector Int))
-> [([Text], [Int])] -> [([Text], Vector Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\([Text]
k, [Int]
is) -> ([Text]
k, [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
is))) (Map [Text] [Int] -> [([Text], [Int])]
forall k a. Map k a -> [(k, a)]
M.toAscList Map [Text] [Int]
grouped)
  where
    n :: Int
n = DataFrame -> Int
nRows DataFrame
df
    cols :: [Column]
cols = (Text -> Column) -> [Text] -> [Column]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> DataFrame -> Column
`unsafeGetColumn` DataFrame
df) [Text]
catCols
    textVecs :: [Vector Text]
textVecs = (Column -> Vector Text) -> [Column] -> [Vector Text]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Vector Text
columnToTextVec [Column]
cols
    bitmaps :: [Maybe Bitmap]
bitmaps = (Column -> Maybe Bitmap) -> [Column] -> [Maybe Bitmap]
forall a b. (a -> b) -> [a] -> [b]
map Column -> Maybe Bitmap
columnBitmap [Column]
cols
    validRow :: Int -> Bool
validRow Int
i = (Maybe Bitmap -> Bool) -> [Maybe Bitmap] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> (Bitmap -> Bool) -> Maybe Bitmap -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
True (Bitmap -> Int -> Bool
`bitmapTestBit` Int
i)) [Maybe Bitmap]
bitmaps
    keyOf :: Int -> [Text]
keyOf Int
i = [Vector Text
tv Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
V.! Int
i | Vector Text
tv <- [Vector Text]
textVecs]
    grouped :: Map [Text] [Int]
grouped =
        (Map [Text] [Int] -> Int -> Map [Text] [Int])
-> Map [Text] [Int] -> [Int] -> Map [Text] [Int]
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
            (\Map [Text] [Int]
m Int
i -> if Int -> Bool
validRow Int
i then ([Int] -> [Int] -> [Int])
-> [Text] -> [Int] -> Map [Text] [Int] -> Map [Text] [Int]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
(++) (Int -> [Text]
keyOf Int
i) [Int
i] Map [Text] [Int]
m else Map [Text] [Int]
m)
            Map [Text] [Int]
forall k a. Map k a
M.empty
            [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

{- | Linear segments with exact closed-form pooling. @λ = 0@ is independent OLS;
@λ > 0@ shrinks each segment's coefficients toward the @n_g@-weighted mean of the
per-segment fits.
-}
instance SegmentFit LinearConfig Double where
    fitSegments :: LinearConfig
-> Double
-> Expr Double
-> [DataFrame]
-> [ModelOf LinearConfig (Expr Double)]
fitSegments LinearConfig
cfg Double
lam Expr Double
target [DataFrame]
dfs
        | Double
lam Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 = (DataFrame -> LinearRegressor) -> [DataFrame] -> [LinearRegressor]
forall a b. (a -> b) -> [a] -> [b]
map (LinearConfig
-> Expr Double
-> FrameFor (Expr Double)
-> FitResult
     (FrameFor (Expr Double)) (ModelOf LinearConfig (Expr Double))
forall cfg input.
(Fit cfg input,
 CheckFrame (FrameReq cfg input) (FrameFor input)) =>
cfg
-> input
-> FrameFor input
-> FitResult (FrameFor input) (ModelOf cfg input)
fit LinearConfig
cfg Expr Double
target) [DataFrame]
dfs
        | [DataFrame] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [DataFrame]
dfs = []
        | Bool
otherwise = LinearConfig
-> Double -> Expr Double -> [DataFrame] -> [LinearRegressor]
shrinkLinear LinearConfig
cfg Double
lam Expr Double
target [DataFrame]
dfs

shrinkLinear ::
    LinearConfig -> Double -> Expr Double -> [DataFrame] -> [LinearRegressor]
shrinkLinear :: LinearConfig
-> Double -> Expr Double -> [DataFrame] -> [LinearRegressor]
shrinkLinear LinearConfig
cfg Double
lam Expr Double
target [DataFrame]
dfs = (Vector Double -> LinearRegressor)
-> [Vector Double] -> [LinearRegressor]
forall a b. (a -> b) -> [a] -> [b]
map Vector Double -> LinearRegressor
toReg [Vector Double]
dsegs
  where
    names :: [Text]
names = case [DataFrame]
dfs of
        (DataFrame
d0 : [DataFrame]
_) -> Expr Double -> DataFrame -> [Text]
forall a. Expr a -> DataFrame -> [Text]
featureNames Expr Double
target DataFrame
d0
        [] -> DataFrameException -> [Text]
forall a e. Exception e => e -> a
throw (Text -> DataFrameException
EmptyDataSetException Text
"shrinkLinear")
    d :: Int
d = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Text]
names
    mats :: [Matrix]
mats = [(Vector Text, Matrix) -> Matrix
forall a b. (a, b) -> b
snd ([Text] -> DataFrame -> (Vector Text, Matrix)
numericMatrix [Text]
names DataFrame
dframe) | DataFrame
dframe <- [DataFrame]
dfs]
    ys :: [Vector Double]
ys = [Expr Double -> DataFrame -> Vector Double
targetDoubles Expr Double
target DataFrame
dframe | DataFrame
dframe <- [DataFrame]
dfs]
    ns :: [Int]
ns = (Matrix -> Int) -> [Matrix] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Matrix -> Int
forall a. Vector a -> Int
V.length [Matrix]
mats
    pooledRows :: Matrix
pooledRows = [Matrix] -> Matrix
forall a. [Vector a] -> Vector a
V.concat [Matrix]
mats
    nP :: Int
nP = Matrix -> Int
forall a. Vector a -> Int
V.length Matrix
pooledRows
    means :: Vector Double
means =
        Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
d ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
            [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Matrix
pooledRows Matrix -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i) Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j | Int
i <- [Int
0 .. Int
nP Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nP
    sds :: Vector Double
sds =
        Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
d ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
            let m :: Double
m = Vector Double
means Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j
                var :: Double
var =
                    [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double -> Double
sq ((Matrix
pooledRows Matrix -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i) Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
m) | Int
i <- [Int
0 .. Int
nP Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                        Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nP
             in Double -> Double
forall a. Floating a => a -> a
sqrt Double
var
    stdValue :: Int -> Double -> Double
stdValue Int
j Double
x = let s :: Double
s = Vector Double
sds Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j in if Double
s Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Vector Double
means Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s
    augStd :: Vector Double -> Vector Double
augStd Vector Double
row = Double -> Vector Double -> Vector Double
forall a. Unbox a => a -> Vector a -> Vector a
VU.cons Double
1 ((Int -> Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap Int -> Double -> Double
stdValue Vector Double
row)
    zMats :: [Matrix]
zMats = [(Vector Double -> Vector Double) -> Matrix -> Matrix
forall a b. (a -> b) -> Vector a -> Vector b
V.map Vector Double -> Vector Double
augStd Matrix
m | Matrix
m <- [Matrix]
mats]
    olsAll :: [Maybe (Vector Double)]
olsAll = (Matrix -> Vector Double -> Maybe (Vector Double))
-> [Matrix] -> [Vector Double] -> [Maybe (Vector Double)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Matrix -> Vector Double -> Maybe (Vector Double)
fitOLS [Matrix]
zMats [Vector Double]
ys
    fitOLS :: Matrix -> Vector Double -> Maybe (Vector Double)
fitOLS Matrix
z Vector Double
y = ([Int] -> Maybe (Vector Double))
-> (Vector Double -> Maybe (Vector Double))
-> Either [Int] (Vector Double)
-> Maybe (Vector Double)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (Vector Double) -> [Int] -> Maybe (Vector Double)
forall a b. a -> b -> a
const Maybe (Vector Double)
forall a. Maybe a
Nothing) Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just (Matrix -> Vector Double -> Either [Int] (Vector Double)
qrLeastSquares Matrix
z Vector Double
y)
    good :: [(Int, Vector Double)]
good = [(Int
n, Vector Double
sol) | (Int
n, Just Vector Double
sol) <- [Int] -> [Maybe (Vector Double)] -> [(Int, Maybe (Vector Double))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
ns [Maybe (Vector Double)]
olsAll]
    totW :: Double
totW = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (((Int, Vector Double) -> Int) -> [(Int, Vector Double)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Vector Double) -> Int
forall a b. (a, b) -> a
fst [(Int, Vector Double)]
good)) :: Double
    dref :: Vector Double
dref
        | [(Int, Vector Double)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Int, Vector Double)]
good = Int -> Double -> Vector Double
forall a. Unbox a => Int -> a -> Vector a
VU.replicate (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Double
0
        | Bool
otherwise =
            Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
k ->
                [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Vector Double
sol Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
k) | (Int
n, Vector Double
sol) <- [(Int, Vector Double)]
good] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
totW
    dsegs :: [Vector Double]
dsegs = (Matrix -> Vector Double -> Vector Double)
-> [Matrix] -> [Vector Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Matrix -> Vector Double -> Vector Double
solveSeg [Matrix]
zMats [Vector Double]
ys
    solveSeg :: Matrix -> Vector Double -> Vector Double
solveSeg Matrix
z Vector Double
y =
        let r :: Vector Double
r = (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith (-) Vector Double
y (Matrix -> Vector Double -> Vector Double
matVec Matrix
z Vector Double
dref)
            a :: Matrix
a = Double -> Matrix -> Matrix
addDiagonal Double
lam (Matrix -> Matrix
gram Matrix
z)
            rhs :: Vector Double
rhs = Matrix -> Vector Double -> Vector Double
tMatVec Matrix
z Vector Double
r
         in case Matrix -> Vector Double -> Maybe (Vector Double)
choleskySolve Matrix
a Vector Double
rhs of
                Just Vector Double
eta -> (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
(a -> b -> c) -> Vector a -> Vector b -> Vector c
VU.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Vector Double
dref Vector Double
eta
                Maybe (Vector Double)
Nothing -> Vector Double
dref
    toReg :: Vector Double -> LinearRegressor
toReg Vector Double
dseg =
        let bStd :: Double
bStd = Vector Double
dseg Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
0
            wStd :: Vector Double
wStd = Int -> Vector Double -> Vector Double
forall a. Unbox a => Int -> Vector a -> Vector a
VU.drop Int
1 Vector Double
dseg
            rawCoef :: Vector Double
rawCoef = (Int -> Double -> Double) -> Vector Double -> Vector Double
forall a b.
(Unbox a, Unbox b) =>
(Int -> a -> b) -> Vector a -> Vector b
VU.imap (\Int
j Double
w -> let s :: Double
s = Vector Double
sds Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j in if Double
s Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
w Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) Vector Double
wStd
            adj :: Double
adj =
                [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                    [ let s :: Double
s = Vector Double
sds Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j
                       in if Double
s Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else (Vector Double
wStd Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Vector Double
means Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s
                    | Int
j <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
         in Vector Double
-> Double -> Vector Text -> Penalty -> LinearRegressor
LinearRegressor Vector Double
rawCoef (Double
bStd Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
adj) ([Text] -> Vector Text
forall a. [a] -> Vector a
V.fromList [Text]
names) (LinearConfig -> Penalty
lcPenalty LinearConfig
cfg)

sq :: Double -> Double
sq :: Double -> Double
sq Double
x = Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x

-- | Add @lam@ to the diagonal of a square matrix.
addDiagonal :: Double -> Matrix -> Matrix
addDiagonal :: Double -> Matrix -> Matrix
addDiagonal Double
lam = (Int -> Vector Double -> Vector Double) -> Matrix -> Matrix
forall a b. (Int -> a -> b) -> Vector a -> Vector b
V.imap (\Int
i Vector Double
row -> Vector Double
row Vector Double -> [(Int, Double)] -> Vector Double
forall a. Unbox a => Vector a -> [(Int, a)] -> Vector a
VU.// [(Int
i, (Vector Double
row Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i) Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
lam)])