diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ + + > 2.0 * x + 0.9999999999999989 -## A decision tree is a readable expression +## Type-safe linear regression -The tree compiles to nested `if/then/else` over your columns — no special -viewer, it is just an expression: +`fit` and `predict` work on both typed and untyped dataframes. You can +have the compiler enforce that you don't hand the fit function a frame +with nullable fields or a non-Double: ```haskell -import DataFrame.DecisionTree (defaultTreeConfig) -import DataFrame.DecisionTree.Model () +import qualified DataFrame.Typed as T +import Data.Maybe (fromJust) +salesT = T.unsafeFreeze @'[T.Column "x" Double, T.Column "y" Double] sales +typedModel = fit defaultLinearConfig (T.col @"y") salesT +scored = T.derive @"prediction" (predict typedModel) salesT + +putStr (unlines + [ "typed model: " ++ D.prettyPrint (T.unTExpr (predict typedModel)) + , "schema after: " ++ show (T.columnNames scored) ]) +``` + +> +> typed model: 2.0 * x + 0.9999999999999989 +> schema after: ["x","y","prediction"] + +## Decision trees + +The tree compiles to nested `if/then/else` over your columns: + +```haskell flowers = D.fromNamedColumns [ ("petal_length", D.fromList ([1.4, 1.3, 1.5, 1.4, 4.5, 4.7, 4.6, 4.4, 5.5, 5.8, 5.6, 5.7] :: [Double])) , ("petal_width", D.fromList ([0.2, 0.2, 0.1, 0.3, 1.5, 1.4, 1.6, 1.3, 2.0, 2.1, 1.9, 2.2] :: [Double])) @@ -70,10 +87,10 @@ > > if petal_length .<=. 2.95 -> then 0.0 -> else if petal_length .<=. 5.1 -> then 1.0 -> else 2.0 +> then 0.0 +> else if petal_length .<=. 5.1 +> then 1.0 +> else 2.0 ## Symbolic regression discovers a formula @@ -81,8 +98,6 @@ it as a dataframe `Expr` plus the accuracy/complexity Pareto front: ```haskell -import DataFrame.SymbolicRegression - curve = D.fromNamedColumns [ ("x", D.fromList xs) , ("y", D.fromList [x * x + x | x <- xs]) @@ -98,36 +113,9 @@ > > x + x * x (mse 0.0) -## When the formula is bigger than a glance - -Not every model is a one-liner. A linear model, a small tree, or a symbolic -expression you can *read*; a 40-tree gradient booster you cannot — its `predict` -is an exact sum of forty trees. Counting the characters in each printed formula -shows the gap: - -```haskell -import DataFrame.Boosting - -gbm = fit defaultGBConfig { gbNEstimators = 40, gbMaxDepth = 2 } (D.col @Double "y") sales -putStr (unlines - [ "linear prediction: " ++ show (length (D.prettyPrint (predict model))) ++ " characters" - , "GBM(40 trees): " ++ show (length (D.prettyPrint (predict gbm))) ++ " characters" ]) -``` - -> -> linear prediction: 28 characters -> GBM(40 trees): 7151 characters - -Even when it is too big to eyeball, the expression is still the whole story: -a self-contained, dependency-free artifact that scores a frame with `derive` — -no pickled blob, no runtime to ship. For the big ensembles the *interpretability* -comes from `gbFeatureImportances` and pretty-printing individual trees, not from -reading the summed formula. - ## Deploy: applying an expression to a frame -Because the model is an `Expr`, deploying it is just `derive` — you add the -prediction as a new column with the ordinary dataframe API: +Because the model is an `Expr` you can use `derive` to do inference. ```haskell D.columnNames (D.derive "prediction" (predict model) sales) @@ -139,16 +127,13 @@ ## A model and its preprocessing compose by substitution Preprocessing is an expression too, so a model trained in a transformed space and -the transform that produced it *compose* — and composition of expressions is +the transform that produced it compose. Composition of expressions is substitution of one into the other. `compileThrough` performs that composition, folding a fitted transform into a prediction so the result is a single formula -over the raw inputs. Here we standardize `x`, fit in the scaled space, then fold +over the raw inputs. Below we standardize `x`, fit in the scaled space, then fold the scaler back in to recover a raw-column model: ```haskell -import DataFrame.Transform -import DataFrame.Metrics - scaler = standardScaler ["x"] sales scaledSales = applyTransform (scalerTransform scaler) sales scaledModel = fit defaultLinearConfig (D.col @Double "y") scaledSales @@ -164,9 +149,7 @@ > folded to raw columns: 3.4156502553198655 * (x - 3.5) / 1.707825127659933 + 8.0 The folded expression is a function of the raw `x` alone, so it scores the -original frame with no preprocessing step at inference time — and by the -substitution lemma it computes the same result (up to floating point) as -transforming the frame and then predicting: +original frame with no preprocessing step at inference time. ```haskell evaluate rmse deployed (D.col @Double "y") sales @@ -175,12 +158,11 @@ > > 3.6259732146947156e-16 -## A realistic run: pick features, split, evaluate held-out, tune - -Real frames are noisy and carry columns you must not train on. Here is a noisy -linear signal with a spurious `id` column: +## Splitting the data, and evaluation ```haskell +import qualified DataFrame as D + realistic = D.fromNamedColumns [ ("id", D.fromList [fromIntegral ((i * 7919) `mod` 97) | i <- [1 .. 40 :: Int]]) , ("x", D.fromList xs) @@ -189,39 +171,20 @@ where xs = map fromIntegral [1 .. 40 :: Int] :: [Double] noise i = fromIntegral ((i * 2654435761 + 12345) `mod` 1000) / 100 - 5 -``` -> - -**Feature selection.** Supervised `fit` uses *every* non-target column as a -feature, so a naive fit drags `id` into the model. `selectFeatures` restricts to -the columns you mean (mirroring the explicit feature list the unsupervised -fitters take), which is the difference between a leaky model and a clean one: - -```haskell -import DataFrame.Model (selectFeatures) - -naive = fit defaultLinearConfig (D.col @Double "y") realistic -guarded = fit defaultLinearConfig (D.col @Double "y") - (selectFeatures ["x"] (D.col @Double "y") realistic) -putStr (unlines - [ "all columns: " ++ D.prettyPrint (predict naive) - , "selectFeatures [\"x\"]: " ++ D.prettyPrint (predict guarded) ]) +clean = D.select ["x", "y"] realistic ``` > -> all columns: -7.746701620642152e-3 * id + 1.9914915483217268 * x + 1.6474622984919354 -> selectFeatures ["x"]: 1.9918011257035637 * x + 1.2630769230769452 -**Hold-out evaluation.** `trainTestSplit` (seeded, deterministic) keeps the score -honest — evaluate on rows the model never saw, and the metrics are realistic, not -the `1e-15` of an in-sample toy: +**Hold-out evaluation.** `randomSplit` (seeded, deterministic) keeps the +score honest — evaluate on rows the model never saw, and the metrics are +realistic, not the `1e-15` of an in-sample toy: ```haskell -import DataFrame.ModelSelection +import System.Random (mkStdGen) -clean = selectFeatures ["x"] (D.col @Double "y") realistic -(train, test) = trainTestSplit 0.75 7 clean +(train, test) = D.randomSplit (mkStdGen 7) 0.75 clean heldModel = fit defaultLinearConfig (D.col @Double "y") train putStr (unlines [ "held-out R^2: " ++ show (evaluate r2 (predict heldModel) (D.col @Double "y") test) @@ -248,7 +211,7 @@ `gridSearch` tunes hyperparameters the same way, over a list of configs. -## Reports without hand-rolling metrics +## Reporting metrics Metrics are plain functions (`rmse`, `mse`, `r2`, `accuracy`, multiclass `precision`/`recall`/`f1`), and `classificationReport` bundles the common numbers @@ -256,8 +219,6 @@ macro/weighted averages): ```haskell -import DataFrame.Metrics.Report - clf = fit defaultLogisticConfig (D.col @Double "species") flowers putStr (show (classificationReportExpr (predict clf) (D.col @Double "species") flowers)) ``` @@ -279,8 +240,6 @@ expression over the raw columns for export: ```haskell -import DataFrame.PCA - features = ["petal_length", "petal_width"] scalerF = standardScaler features flowers pca = fit (PCAConfig (NComp 2) True) (map (D.col @Double) features) flowers @@ -292,21 +251,6 @@ > > ["petal_length","petal_width","species","pc1","pc2"] -## Unsupervised models are inspectable too - -k-means returns `cluster_centers_`-style centroids, and per-cluster distance / -assignment expressions: - -```haskell -import DataFrame.KMeans - -km = fit defaultKMeansConfig { kmK = 3, kmSeed = 1 } (map (D.col @Double) features) flowers -kmCenters km -``` - -> -> [[1.4,0.2],[5.65,2.05],[4.55,1.4500000000000002]] - ## Synthesize the feature you would have hand-engineered `DataFrame.Synthesis` is automated feature engineering: a bottom-up enumerative @@ -317,8 +261,6 @@ to exact — still a formula you can read: ```haskell -import DataFrame.Synthesis - interactions = D.fromNamedColumns [ ("a", D.fromList as) , ("b", D.fromList bs) @@ -333,7 +275,7 @@ withFeat = D.derive "synth" (predict feature) interactions fitModel = fit defaultLinearConfig (D.col @Double "y") - (selectFeatures ["synth"] (D.col @Double "y") withFeat) + (D.select ["synth", "y"] withFeat) putStr (unlines [ "discovered feature: " ++ D.prettyPrint (predict feature) @@ -349,125 +291,3 @@ `predict feature` is the single best expression; `sfFeatures feature` is the whole ranked, deduplicated bank, ready to `derive` as a batch of candidate columns. - -## What's in the box - -| Task | Models | -|------|--------| -| Regression | OLS, ridge, lasso, elastic net, regression trees, gradient boosting, symbolic regression | -| Classification | logistic regression, linear SVC, RFF kernel SVM, decision trees, gradient boosting, AdaBoost | -| Dimensionality reduction | PCA, Nyström kernel PCA | -| Clustering | k-means, Gaussian mixtures, DBSCAN | -| Feature engineering | `DataFrame.Synthesis` (enumerative feature synthesis), symbolic regression | -| Evaluation | `DataFrame.Metrics` (metrics + `evaluate`), `DataFrame.Metrics.Report` (reports) | -| Pipelines & tuning | `DataFrame.Transform` (composable transforms), `DataFrame.ModelSelection` (`trainTestSplit`, `crossValidate`, `gridSearch`) | - -Every model is a `Fit` instance, so there is one verb to train — `fit cfg input -df` — and every model with an honest out-of-sample prediction is a `Predict` -instance, so one verb to compile it — `predict model`. Auxiliary outputs -(`gbProbaExpr`, `logisticProbExprs`, `kmeansDistanceExprs`, `pcaTransform`, …) -keep descriptive names; transductive models like DBSCAN deliberately have no -`Predict` instance. Fits that use randomness take a `seed` in their config, so -results are reproducible across Linux, macOS, and Windows. Pure Haskell — the -only extra dependency beyond the dataframe packages is `random`. - -## Design notes: the categorical account - -The two verbs live in `DataFrame.Model`: - -``` -class Fit cfg input model | cfg input -> model where - fit :: cfg -> input -> DataFrame -> model - -class Predict model r | model -> r where - predict :: model -> Expr r -``` - -They are small on purpose, because the structure they hang on lives in the -expression language, not in the classes. The framing borrows from -[*Seven Sketches in Compositionality*](https://arxiv.org/abs/1803.05316) -(Fong & Spivak) and the Para/Lens account of learners -([Fong, Johnson & Spivak, *Lenses and Learners*](https://arxiv.org/abs/1903.03671); -[Cruttwell et al., *Categorical Foundations of Gradient-Based Learning*](https://arxiv.org/abs/2103.01931)). -What follows is deliberately careful about what is load-bearing and what is only -analogy. - -**The row-wise fragment is a category.** Restrict to the row-wise expression -constructors — `Col`, `Lit`, `Unary`, `Binary`, `If`. Take typed column contexts -as objects and, as an arrow `Γ → Δ`, a `Δ`-tuple of such expressions over `Γ`. -Composition is simultaneous substitution (`substituteColumns`, added to -`dataframe-core` for exactly this) and the identities are the column projections -(`Col`). This is the category of contexts (the Lawvere theory) of the column -signature. The restriction is load-bearing for *both* laws, not just composition: -`Agg` and `Over` are column-level/relational, not row-wise maps, and the raw-text -column reference inside `CastWith` is opaque to substitution — so identity-by-`Col` -fails on those constructors too. They are excluded by construction (transforms -reject `Agg`/`Over`), which is why composition and identity stay well defined. - -**`predict` gives every model a uniform codomain.** `Predict model r` interprets a -fitted model as an arrow in that category: `predict model :: Expr r` runs from the -model's feature context to the one-column context `{r}`, and the dependency -`model -> r` fixes the codomain object. This is *not* a functor or a denotation in -the technical sense — there is no category of models to be functorial over. The -real, useful property is uniformity: every model's prediction lands in the *same* -expression type (`Expr Double`/`Expr a`/`Expr Int`), so `derive`, the `Transform` -monoid, and `compileThrough` all apply with no per-model glue. That the compiled -`Expr` actually agrees with the fitted record's own parameters is a tested property -(`tests/Learn/Denotation.hs`), not a typeclass law — the class only knows the -symbolic half. - -**`fit` is the parametrized-morphism (Para) fragment.** `fit cfg input df` chooses a -parameter — the trained record — and `predict` is the forward map applied at it. -In the Para/Lens picture of learning a learner is a *parametrized lens* carrying a -forward map plus backward update/request maps; we inhabit only the forward (Para) -part and expose no backward maps, because this interface is batch training, not -online gradient exchange. That is a complete, self-contained sub-structure, not a -half-built one — but it does mean the Lens vocabulary is motivation here, not -something the code instantiates. (The functional dependency `cfg input -> model` -fixes the parameter *type*; `fit` is the value-level map that picks the point.) - -**`Transform` is a monoid of derived-column lists.** `Transform`'s `<>` keeps the -earlier step's outputs and rewrites the later step's column references through them -by substitution; `mempty` is the empty list. These are context-*extending* maps -(`applyTransform` adds columns), so this is an ordinary algebraic monoid — a monoid -*is* a one-object category (Seven Sketches ch. 3) — not the endomorphism monoid of a -fixed object. Associativity and identity hold for the row-wise fragment **provided -output names do not collide**: the implementation merges output maps with -`Data.Map.fromList`, which keeps the last binding on a clash, so reusing a column -name across steps is the one way to break the law. - -**Composition is the point.** `compileThrough t (predict m)` realizes the composite -`predict m ∘ t` (read right-to-left: first `t`, then `predict m`) by substituting -`t`'s definitions into `predict m`. By the substitution lemma it denotes the same -function as transforming the frame and then predicting — equal results up to -floating point, not syntactically identical expressions. That is exactly the -"compose by substitution" example above, and it is why a model trained in a -transformed space deploys as one formula over the raw columns. - -**What deliberately has no `predict` — two different reasons.** DBSCAN is -transductive: every clustering *fit* depends on the whole training set, but what -distinguishes the models is whether the *fitted* model induces an out-of-sample -rule. k-means (nearest centroid) and GMM (max posterior) do, so they have honest -`predict` arrows; DBSCAN's density-reachability assignment has no per-row rule, so -we give it no `Predict` instance rather than a fake `Maybe` or a throwing stub. -PCA and kernel PCA are the *opposite* case: they *are* arrows, but multi-output -feature maps with no privileged label column, so their canonical interface is a -`Transform` (`pcaTransform`/`pcaExprs`), not a one-column `predict`. - -**A note on classifiers.** A multiclass `predict` is a genuine arrow into the label -object, but it compiles arg-max to a nested-`If` cascade (`argMaxExpr`), quadratic -in the number of classes — so for a 5-class model `prettyPrint (predict m)` is an -If-tree, not a tidy formula. The "model is a readable formula" aesthetic is honest -for affine and tree models; for classifiers and clusterers the value is that the -arrow exists and composes, not that it is short. - -**An aside.** A linear or affine model's prediction is a signal-flow graph — a -weighted sum of inputs. `affineExpr` builds the arrow in the prop of affine *maps* -(the single-valued sub-prop of Seven Sketches ch. 5's signal-flow calculus of -affine relations), and dropping zero-weight terms is diagram simplification — -deleting a zero-gain wire. - -(The "instance is a functor `C → Set`" slogan from Spivak's functorial data model, -Seven Sketches ch. 3, is sometimes invoked for dataframes; a single flat table is -the degenerate case — a schema with no foreign-key morphisms — so it is an analogy -here, not a structure we use.) diff --git a/dataframe-learn.cabal b/dataframe-learn.cabal --- a/dataframe-learn.cabal +++ b/dataframe-learn.cabal @@ -1,6 +1,6 @@ -cabal-version: 2.4 +cabal-version: 3.4 name: dataframe-learn -version: 1.1.0.1 +version: 2.0.0.0 synopsis: Interpretable, expression-returning machine learning for the dataframe ecosystem. description: A small scikit-learn-style ML library where every model returns both an @@ -28,10 +28,13 @@ -Wunused-local-binds -Wunused-packages -library +-- Numeric kernels, linear solver, tree engine, and symbolic-regression +-- internals. Private: reachable only by this package's own components +-- (the public estimators below and the learn-internal test-suite). +library internal import: warnings + visibility: private exposed-modules: - DataFrame.DecisionTree DataFrame.DecisionTree.Types DataFrame.DecisionTree.CondVec DataFrame.DecisionTree.Cart @@ -50,18 +53,42 @@ DataFrame.LinearAlgebra.Eigen DataFrame.Random DataFrame.Featurize.Internal + DataFrame.SymbolicRegression.Expr + DataFrame.SymbolicRegression.Simplify + DataFrame.SymbolicRegression.Optimize + DataFrame.SymbolicRegression.GP + build-depends: base >= 4 && < 5, + containers >= 0.6.7 && < 0.10, + parallel >= 3.3 && < 4, + random >= 1.2 && < 2, + dataframe-core:internal >= 2.0 && < 2.1, + dataframe-operations >= 2.0 && < 2.1, + text >= 2.1 && < 3, + vector >= 0.13 && < 0.15, + vector-algorithms >= 0.9 && < 0.11 + hs-source-dirs: src-internal + default-language: Haskell2010 + +-- Curated estimator surface: fit/predict + configs + fitted records + the +-- DataFrame.Learn umbrella. Implementation sealed in the private sublib. +library + import: warnings + exposed-modules: + DataFrame.Learn DataFrame.Model DataFrame.LinearModel DataFrame.LinearModel.Regression DataFrame.LinearModel.Logistic DataFrame.SVM + DataFrame.SVM.RFF + DataFrame.DecisionTree DataFrame.DecisionTree.Regression DataFrame.DecisionTree.Model DataFrame.PCA DataFrame.PCA.Kernel - DataFrame.SVM.RFF DataFrame.KMeans DataFrame.Transform + DataFrame.Transform.Serialize DataFrame.Boosting DataFrame.Boosting.GBM DataFrame.Boosting.AdaBoost @@ -71,19 +98,59 @@ DataFrame.Metrics.Report DataFrame.ModelSelection DataFrame.SymbolicRegression - DataFrame.SymbolicRegression.Expr - DataFrame.SymbolicRegression.Simplify - DataFrame.SymbolicRegression.Optimize - DataFrame.SymbolicRegression.GP DataFrame.Synthesis + DataFrame.Segmented build-depends: base >= 4 && < 5, - containers >= 0.6.7 && < 0.9, - parallel ^>= 3.2, + aeson >= 0.11.0.0 && < 3, + containers >= 0.6.7 && < 0.10, + parallel >= 3.3 && < 4, random >= 1.2 && < 2, - dataframe-core ^>= 1.1, - dataframe-operations ^>= 1.1.1, + dataframe-core >= 2.0 && < 2.1, + dataframe-core:internal >= 2.0 && < 2.1, + dataframe-operations >= 2.0 && < 2.1, + dataframe-operations:internal >= 2.0 && < 2.1, + dataframe-expr-serializer >= 1.1 && < 1.2, + dataframe-learn:internal, text >= 2.1 && < 3, - vector ^>= 0.13, - vector-algorithms ^>= 0.9 + vector >= 0.13 && < 0.15 hs-source-dirs: src + default-language: Haskell2010 + +-- Tests that exercise the private `internal` sublib directly (solver, tree +-- engine, symbolic-regression internals). They live here rather than in the +-- meta `dataframe` test-suite because that suite cannot reach a private sublib. +test-suite learn-internal + import: warnings + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: DataFrameApi + Cart + DecisionTree + TreePruning + Worklist + LinearSolver + Properties.Simplify + Learn.Numerics + Learn.Symbolic + Learn.EdgeCases + Learn.NumericalRigor + -- Depends on the constituent packages directly, NOT the meta `dataframe` + -- (which depends back on dataframe-learn, a package-level cycle). + build-depends: base >= 4 && < 5, + aeson >= 0.11.0.0 && < 3, + bytestring >= 0.11 && < 0.14, + containers >= 0.6.7 && < 0.10, + dataframe-core >= 2.0 && < 2.1, + dataframe-core:internal >= 2.0 && < 2.1, + dataframe-csv >= 2.0 && < 2.1, + dataframe-learn, + dataframe-learn:internal, + dataframe-operations >= 2.0 && < 2.1, + dataframe-operations:internal >= 2.0 && < 2.1, + HUnit >= 1.6 && < 1.8, + QuickCheck >= 2 && < 3, + random >= 1 && < 2, + text >= 2.1 && < 3, + vector >= 0.13 && < 0.15 + hs-source-dirs: tests-internal default-language: Haskell2010 diff --git a/src-internal/DataFrame/DecisionTree/Cart.hs b/src-internal/DataFrame/DecisionTree/Cart.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Cart.hs @@ -0,0 +1,290 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | sklearn-faithful CART initializer used to seed TAO. One-hot encodes +categoricals and splits on exact (unsmoothed) Gini over midpoint thresholds +(@<=@ routes left), matching @DecisionTreeClassifier(criterion='gini')@. +-} +module DataFrame.DecisionTree.Cart ( + CartFeature (..), + CartNode (..), + sortIndicesByValue, + buildCartTree, + cartFeatures, + cartTargetLabels, +) where + +import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operations.Core (nRows) +import DataFrame.Operators + +import Data.Either (fromRight) +import Data.Function (on) +import Data.List (foldl') +import qualified Data.Map.Strict as M +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Merge as VA +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +{- | A one-hot feature column: per-row Double values plus the sklearn LEFT +predicate (@x <= threshold@) over the ORIGINAL DataFrame. +-} +data CartFeature = CartFeature + { cfValues :: !(VU.Vector Double) + , cfPred :: !(Double -> Expr Bool) + } + +-- | Pre-'Tree' CART node: a leaf class id, or a split on feature @j@. +data CartNode = CLeaf !Int | CSplit !Int !Double !CartNode !CartNode + +-- | Immutable per-fit context for the CART recursion. +data CartCtx = CartCtx + { ctxFeats :: !(V.Vector CartFeature) + , ctxNFeats :: !Int + , ctxCodes :: !(VU.Vector Int) + , ctxNClasses :: !Int + , ctxMaxDepth :: !Int + , ctxMinLeaf :: !Int + } + +{- | Indices @0..n-1@ stably sorted by their value (ascending), ties keeping +ascending index. In-place unboxed merge sort — no boxed-list allocation. +-} +sortIndicesByValue :: VU.Vector Double -> VU.Vector Int +sortIndicesByValue vs = + VU.create $ do + mv <- VU.thaw (VU.enumFromN 0 (VU.length vs)) + VA.sortBy (compare `on` (vs VU.!)) mv + pure mv + +buildCartTree :: + forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> Tree a +buildCartTree cfg target df = + cartToTree feats classes (buildCartNode ctx 0 (VU.enumFromN 0 nAll) featSorted) + where + nAll = nRows df + feats = V.fromList (cartFeatures target df) + featSorted = V.map (sortIndicesByValue . cfValues) feats + labels = cartLabels @a df target + classes = cartClasses labels + ctx = + CartCtx + feats + (V.length feats) + (classCodes classes labels) + (V.length classes) + (maxTreeDepth cfg) + (max 1 (minLeafSize cfg)) + +cartLabels :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a +cartLabels df target = case interpret @a df (Col target) of + Right (TColumn column) -> fromRight err (toVector @a column) + _ -> err + where + err = error "buildCartTree: cannot interpret target column" + +cartClasses :: (Ord a) => V.Vector a -> V.Vector a +cartClasses = V.fromList . Set.toList . Set.fromList . V.toList + +classCodes :: (Ord a) => V.Vector a -> V.Vector a -> VU.Vector Int +classCodes classes labels = VU.generate (V.length labels) (\i -> M.findWithDefault 0 (labels V.! i) ix) + where + ix = M.fromList (zip (V.toList classes) [0 ..]) + +cartToTree :: V.Vector CartFeature -> V.Vector a -> CartNode -> Tree a +cartToTree feats classes = go + where + go (CLeaf cid) = Leaf (classes V.! cid) + go (CSplit fj thr l r) = Branch (cfPred (feats V.! fj) thr) (go l) (go r) + +classCounts :: CartCtx -> VU.Vector Int -> VU.Vector Int +classCounts ctx idxs = + VU.accumulate + (+) + (VU.replicate (ctxNClasses ctx) 0) + (VU.map (\i -> (ctxCodes ctx VU.! i, 1)) idxs) + +isPure :: VU.Vector Int -> Bool +isPure counts = VU.length (VU.filter (> 0) counts) <= 1 + +buildCartNode :: + CartCtx -> Int -> VU.Vector Int -> V.Vector (VU.Vector Int) -> CartNode +buildCartNode ctx depth idxs sortedByFeat + | VU.length idxs < 2 || depth >= ctxMaxDepth ctx || isPure counts = leaf + | otherwise = + maybe + leaf + (splitNode ctx depth idxs sortedByFeat) + (bestSplit ctx sortedByFeat counts n) + where + n = VU.length idxs + counts = classCounts ctx idxs + leaf = CLeaf (VU.maxIndex counts) + +splitNode :: + CartCtx -> + Int -> + VU.Vector Int -> + V.Vector (VU.Vector Int) -> + (Int, Double) -> + CartNode +splitNode ctx depth idxs sortedByFeat (fj, thr) = + CSplit fj thr (rec leftIdx leftSorted) (rec rightIdx rightSorted) + where + vals = cfValues (ctxFeats ctx V.! fj) + leftIdx = VU.filter (\i -> vals VU.! i <= thr) idxs + rightIdx = VU.filter (\i -> vals VU.! i > thr) idxs + leftSorted = V.map (VU.filter (\i -> vals VU.! i <= thr)) sortedByFeat + rightSorted = V.map (VU.filter (\i -> vals VU.! i > thr)) sortedByFeat + rec = buildCartNode ctx (depth + 1) + +{- | Minimum weighted-child-Gini @(feature, threshold)@; the first feature wins +ties; 'Nothing' when no feature has a leaf-size-respecting threshold. +-} +bestSplit :: + CartCtx -> + V.Vector (VU.Vector Int) -> + VU.Vector Int -> + Int -> + Maybe (Int, Double) +bestSplit ctx sortedByFeat counts n = + fmap (\(_, j, t) -> (j, t)) (foldl' consider Nothing [0 .. ctxNFeats ctx - 1]) + where + total = VU.toList counts + consider acc fj = case sweepFeature ctx total (sortedByFeat V.! fj) (ctxFeats ctx V.! fj) n of + Just (g, thr) | maybe True (\(gB, _, _) -> g < gB) acc -> Just (g, fj, thr) + _ -> acc + +{- | Accumulator while sweeping a feature's sorted rows: best @(gini, thr)@ so +far, per-class left counts, rows moved left, and the previous value seen. +-} +data Sweep = Sweep + { swBest :: !(Maybe (Double, Double)) + , swLeft :: ![Int] + , swMoved :: !Int + , swPrev :: !Double + } + +sweepFeature :: + CartCtx -> + [Int] -> + VU.Vector Int -> + CartFeature -> + Int -> + Maybe (Double, Double) +sweepFeature ctx total si feat n = + swBest + ( foldl' + step + (Sweep Nothing (replicate (ctxNClasses ctx) 0) 0 (0 / 0)) + [0 .. VU.length si - 1] + ) + where + vals = cfValues feat + step s k = advance ctx total n (vals VU.! i) (ctxCodes ctx VU.! i) s + where + i = si VU.! k + +advance :: CartCtx -> [Int] -> Int -> Double -> Int -> Sweep -> Sweep +advance ctx total n v c s = + Sweep + (considerThreshold ctx total n v s) + (bumpClass c (swLeft s)) + (swMoved s + 1) + v + +considerThreshold :: + CartCtx -> [Int] -> Int -> Double -> Sweep -> Maybe (Double, Double) +considerThreshold ctx total n v s + | swMoved s >= ctxMinLeaf ctx + , n - swMoved s >= ctxMinLeaf ctx + , v > swPrev s + 1e-7 = + keepBetter + (swBest s) + (weightedGini total (swLeft s) (swMoved s) n) + ((swPrev s + v) / 2) + | otherwise = swBest s + +keepBetter :: + Maybe (Double, Double) -> Double -> Double -> Maybe (Double, Double) +keepBetter best g thr = case best of + Just (wb, _) | wb <= g -> best + _ -> Just (g, thr) + +weightedGini :: [Int] -> [Int] -> Int -> Int -> Double +weightedGini total leftAcc nl n = + ( fromIntegral nl * giniImpurity leftAcc nl + + fromIntegral nr * giniImpurity rightAcc nr + ) + / fromIntegral n + where + nr = n - nl + rightAcc = zipWith (-) total leftAcc + +-- | Gini impurity @1 - Σ (c/m)²@ of a class-count list of total @m@. +giniImpurity :: [Int] -> Int -> Double +giniImpurity _ 0 = 0 +giniImpurity cs m = 1 - sum [let p = fromIntegral c / fromIntegral m in p * p | c <- cs] + +bumpClass :: Int -> [Int] -> [Int] +bumpClass c = zipWith (\j x -> if j == c then x + 1 else x) [0 ..] + +-- | One-hot features in @pd.get_dummies(drop_first=False)@ column order. +cartFeatures :: T.Text -> DataFrame -> [CartFeature] +cartFeatures target df = concatMap (featuresOfColumn df) (filter (/= target) (columnNames df)) + +featuresOfColumn :: DataFrame -> T.Text -> [CartFeature] +featuresOfColumn df c = case unsafeGetColumn c df of + UnboxedColumn _ (v :: VU.Vector b) -> numericFeature @b c v + BoxedColumn _ (v :: V.Vector b) -> oneHotFeatures @b (nRows df) c v + pt@(PackedText _ _) -> case materializePacked pt of + BoxedColumn _ (v :: V.Vector b) -> oneHotFeatures @b (nRows df) c v + _ -> [] + +numericFeature :: + forall b. (Columnable b, VU.Unbox b) => T.Text -> VU.Vector b -> [CartFeature] +numericFeature c v = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [CartFeature v (\t -> F.col @Double c .<=. F.lit t)] + Nothing -> case sIntegral @b of + STrue -> + [ CartFeature (VU.map fromIntegral v) (\t -> F.toDouble (F.col @b c) .<=. F.lit t) + ] + SFalse -> [] + +oneHotFeatures :: + forall b. (Columnable b) => Int -> T.Text -> V.Vector b -> [CartFeature] +oneHotFeatures nAll c v = case testEquality (typeRep @b) (typeRep @T.Text) of + Just Refl -> [oneHot nAll c v cat | cat <- Set.toList (Set.fromList (V.toList v))] + Nothing -> [] + +oneHot :: Int -> T.Text -> V.Vector T.Text -> T.Text -> CartFeature +oneHot nAll c v cat = + CartFeature + (VU.generate nAll (\i -> if v V.! i == cat then 1 else 0)) + (const (F.col @T.Text c ./=. F.lit cat)) + +-- | Target column as string labels (matches pandas @y.astype(str)@). +cartTargetLabels :: T.Text -> DataFrame -> V.Vector T.Text +cartTargetLabels target df = case unsafeGetColumn target df of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @T.Text) of + Just Refl -> v + Nothing -> V.map (T.pack . show) v + UnboxedColumn _ (v :: VU.Vector b) -> V.map (T.pack . show) (V.convert v) + pt@(PackedText _ _) -> case materializePacked pt of + BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @T.Text) of + Just Refl -> v + Nothing -> V.map (T.pack . show) v + _ -> V.empty diff --git a/src-internal/DataFrame/DecisionTree/Categorical.hs b/src-internal/DataFrame/DecisionTree/Categorical.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Categorical.hs @@ -0,0 +1,380 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Categorical split candidates: Breiman prefixes for binary targets, +subset/singleton enumeration otherwise, and cross-column equality. Each +value-list becomes an OR-of-equalities condition. +-} +module DataFrame.DecisionTree.Categorical ( + TargetInfo (..), + mkTargetInfo, + distinctValuesUpTo, + validBoxedValues, + orEqs, + subsetSplits, + subsetLists, + singletonSplits, + singletonLists, + breimanPrefixSplits, + breimanPrefixLists, + catValueLists, + membershipVec, + crossColumnConds, + discreteConditions, + discreteCondVecs, +) where + +import DataFrame.DecisionTree.CondVec (CondVec (..), materializeCondVec) +import DataFrame.DecisionTree.Types ( + ColumnOrdering, + SynthConfig (..), + TreeConfig (..), + withOrdFrom, + ) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operators + +import Data.Either (fromRight) +import Data.Function (on) +import Data.List (inits, sort, sortBy, subsequences) +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe, mapMaybe) +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | Valid-slot view of a nullable boxed column (null slots hold crash-thunks). +validBoxedValues :: Bitmap -> V.Vector a -> V.Vector a +validBoxedValues bm = V.ifilter (\i _ -> bitmapTestBit bm i) + +{- | Target-column summary driving the categorical generator: binary vs +multi-class, the deterministic positive class, and the raw label vector. +-} +data TargetInfo target = TargetInfo + { tiIsBinary :: !Bool + , tiPositiveClass :: !(Maybe target) + , tiValues :: !(V.Vector target) + } + +{- | Compute 'TargetInfo' once per fit. The positive class for binary targets +is the lexicographically-first distinct value, for deterministic pools. +-} +mkTargetInfo :: + forall target. + (Columnable target, Ord target) => + T.Text -> DataFrame -> Maybe (TargetInfo target) +mkTargetInfo target df = case interpret @target df (Col target) of + Right (TColumn column) -> + either (const Nothing) (Just . targetInfoFromValues) (toVector @target column) + _ -> Nothing + +targetInfoFromValues :: (Ord target) => V.Vector target -> TargetInfo target +targetInfoFromValues vals = TargetInfo isBinary posClass vals + where + distinct = Set.toAscList (Set.fromList (V.toList vals)) + isBinary = length distinct == 2 + posClass = case distinct of + (p : _) | isBinary -> Just p + _ -> Nothing + +{- | Distinct values, capped: @Right vs@ (sorted) under the cap, else @Left@ +the count-so-far so the caller routes to the high-cardinality path. +-} +distinctValuesUpTo :: (Ord a) => Int -> V.Vector a -> Either Int [a] +distinctValuesUpTo cap values = go Set.empty 0 + where + n = V.length values + go !s !i + | i >= n = Right (Set.toAscList s) + | Set.size s > cap = Left (Set.size s) + | otherwise = go (Set.insert (V.unsafeIndex values i) s) (i + 1) + +{- | OR-of-equalities for a value-list, shared by the expression and +truth-vector discrete paths so they stay byte-identical. +-} +orEqs :: (a -> Expr Bool) -> [a] -> Expr Bool +orEqs eqLit = foldr1 (.||.) . map eqLit + +subsetSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] +subsetSplits eqLit = map (orEqs eqLit) . subsetLists + +-- | Proper non-empty, non-full subsets of the values. +subsetLists :: [a] -> [[a]] +subsetLists vs = drop 1 (init (subsequences vs)) + +singletonSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] +singletonSplits = map + +singletonLists :: [a] -> [[a]] +singletonLists = map (: []) + +breimanPrefixSplits :: + (Ord a, Ord target) => + target -> + V.Vector a -> + V.Vector target -> + [a] -> + (a -> Expr Bool) -> + [Expr Bool] +breimanPrefixSplits pc values targetVals distinctVals eqLit = + map (orEqs eqLit) (breimanPrefixLists pc values targetVals distinctVals) + +{- | Breiman's binary-target split set: sort levels by Laplace-smoothed +positive rate, then take every contiguous non-trivial prefix. +-} +breimanPrefixLists :: + (Ord a, Ord target) => target -> V.Vector a -> V.Vector target -> [a] -> [[a]] +breimanPrefixLists pc values targetVals distinctVals = + nonTrivialPrefixes (sortByRate (levelCounts pc values targetVals) distinctVals) + +levelCounts :: + (Ord a, Eq target) => + target -> V.Vector a -> V.Vector target -> M.Map a (Int, Int) +levelCounts pc values targetVals = V.ifoldl' add M.empty values + where + add acc i v = M.insertWith plus v (indicator (V.unsafeIndex targetVals i == pc), 1) acc + plus (p1, n1) (p2, n2) = (p1 + p2, n1 + n2) + indicator b = if b then 1 else 0 + +laplaceRate :: (Ord a) => M.Map a (Int, Int) -> a -> Double +laplaceRate counts v = case M.lookup v counts of + Nothing -> 0.5 + Just (pos, n) -> (fromIntegral pos + 1) / (fromIntegral n + 2) + +sortByRate :: (Ord a) => M.Map a (Int, Int) -> [a] -> [a] +sortByRate counts = sortBy (compare `on` (\v -> (laplaceRate counts v, v))) + +nonTrivialPrefixes :: [a] -> [[a]] +nonTrivialPrefixes = drop 1 . init . inits + +{- | Value-lists a categorical column contributes; shared by the expression and +truth-vector paths so both enumerate identical candidates in the same order. +-} +catValueLists :: + (Ord a, Ord target) => + Bool -> Maybe target -> V.Vector target -> Int -> V.Vector a -> [[a]] +catValueLists isBinary posClass targetVals subsetCap values + | V.null values = [] + | isBinary, Just pc <- posClass = binaryLists pc targetVals values + | otherwise = multiclassLists subsetCap values + +binaryLists :: + (Ord a, Ord target) => target -> V.Vector target -> V.Vector a -> [[a]] +binaryLists pc targetVals values + | length distinct < 2 = [] + | otherwise = breimanPrefixLists pc values targetVals distinct + where + distinct = fromRight (ascDistinct values) (distinctValuesUpTo 64 values) + +multiclassLists :: (Ord a) => Int -> V.Vector a -> [[a]] +multiclassLists subsetCap values = case distinctValuesUpTo subsetCap values of + Right vs | length vs >= 2 -> subsetLists vs + Right _ -> [] + Left _ -> singletonLists (ascDistinct values) + +ascDistinct :: (Ord a) => V.Vector a -> [a] +ascDistinct = Set.toAscList . Set.fromList . V.toList + +{- | Truth vector of @col ∈ values@ read directly from the column; equal to +interpreting @orEqs (== v) values@ because the values are distinct. +-} +membershipVec :: (Ord a) => V.Vector a -> [a] -> VU.Vector Bool +membershipVec colVals vs = + let !s = Set.fromList vs + in VU.generate (V.length colVals) (\i -> Set.member (colVals `V.unsafeIndex` i) s) + +{- | Per-fit categorical generation context bundling the target summary and +the column-ordering registry. +-} +data CatCtx target = CatCtx + { ccBinary :: !Bool + , ccPos :: !(Maybe target) + , ccTargets :: !(V.Vector target) + , ccSubsetCap :: !Int + , ccOrds :: !ColumnOrdering + } + +catCtx :: TargetInfo target -> TreeConfig -> CatCtx target +catCtx ti cfg = + CatCtx + (tiIsBinary ti) + (tiPositiveClass ti) + (tiValues ti) + (maxCategoricalSubsetCardinality (synthConfig cfg)) + (columnOrdering cfg) + +catValueListsFor :: (Ord a, Ord target) => CatCtx target -> V.Vector a -> [[a]] +catValueListsFor ctx = catValueLists (ccBinary ctx) (ccPos ctx) (ccTargets ctx) (ccSubsetCap ctx) + +-- | True for numeric columns (handled by the numeric pool, not here). +isNumericKind :: forall a. (Columnable a) => Bool +isNumericKind = case sFloating @a of + STrue -> True + SFalse -> case sIntegral @a of + STrue -> True + SFalse -> False + +{- | All equality-based candidate splits from non-numeric columns: per-column +categorical conditions plus cross-column equality/order conditions. +-} +discreteConditions :: + forall target. + (Columnable target, Ord target) => + TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool] +discreteConditions targetInfo cfg df = + concatMap (columnConds (catCtx targetInfo cfg) df) (columnNames df) + ++ crossColumnConds cfg df + +columnConds :: + (Columnable target, Ord target) => + CatCtx target -> DataFrame -> T.Text -> [Expr Bool] +columnConds ctx df colName = case unsafeGetColumn colName df of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColConds ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> nullableColConds ctx colName bm column + UnboxedColumn _ (_ :: VU.Vector a) -> [] + pt@(PackedText _ _) -> case materializePacked pt of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColConds ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> nullableColConds ctx colName bm column + _ -> [] + +nonNullColConds :: + forall a target. + (Columnable a, Ord target) => + CatCtx target -> T.Text -> V.Vector a -> [Expr Bool] +nonNullColConds ctx colName column = + fromMaybe + [] + ( withOrdFrom @a + (ccOrds ctx) + (map (orEqs (eqExprFor @a colName)) (catValueListsFor ctx column)) + ) + +nullableColConds :: + forall a target. + (Columnable a, Ord target) => + CatCtx target -> T.Text -> Bitmap -> V.Vector a -> [Expr Bool] +nullableColConds ctx colName bm column + | isNumericKind @a || V.null valid = [] + | otherwise = + fromMaybe + [] + ( withOrdFrom @a + (ccOrds ctx) + (map (orEqs (eqJustFor @a colName)) (catValueListsFor ctx valid)) + ) + where + valid = validBoxedValues bm column + +eqExprFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool +eqExprFor colName v = Col @a colName .==. Lit v + +eqJustFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool +eqJustFor colName v = Col @(Maybe a) colName .==. Lit (Just v) + +-- | Cross-column equality/order conditions over pairs of same-typed columns. +crossColumnConds :: TreeConfig -> DataFrame -> [Expr Bool] +crossColumnConds cfg df = concatMap (pairConds (columnOrdering cfg) df) (allowedPairs cfg df) + +allowedPairs :: TreeConfig -> DataFrame -> [(T.Text, T.Text)] +allowedPairs cfg df = + [ (l, r) + | l <- columnNames df + , r <- columnNames df + , l /= r + , not (isDisallowedPair cfg l r) + ] + +isDisallowedPair :: TreeConfig -> T.Text -> T.Text -> Bool +isDisallowedPair cfg l r = + any + (\(l', r') -> sort [l', r'] == sort [l, r]) + (disallowedCombinations (synthConfig cfg)) + +pairConds :: ColumnOrdering -> DataFrame -> (T.Text, T.Text) -> [Expr Bool] +pairConds ords df (l, r) = case ( materializePacked (unsafeGetColumn l df) + , materializePacked (unsafeGetColumn r df) + ) of + (BoxedColumn Nothing (_ :: V.Vector a), BoxedColumn Nothing (_ :: V.Vector b)) -> strictPairConds @a @b l r + (BoxedColumn (Just _) (_ :: V.Vector a), BoxedColumn (Just _) (_ :: V.Vector b)) -> nullablePairConds @a @b ords l r + _ -> [] + +strictPairConds :: + forall a b. (Columnable a, Columnable b) => T.Text -> T.Text -> [Expr Bool] +strictPairConds l r = case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> [Col @a l .==. Col @a r] + Nothing -> [] + +nullablePairConds :: + forall a b. + (Columnable a, Columnable b) => + ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] +nullablePairConds ords l r = case testEquality (typeRep @a) (typeRep @b) of + Nothing -> [] + Just Refl -> nullableEqOrLe @a ords l r + +nullableEqOrLe :: + forall a. (Columnable a) => ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] +nullableEqOrLe ords l r + | isTextType @a = eqOnly + | otherwise = + maybe + eqOnly + (++ eqOnly) + (withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r]) + where + eqOnly = [Col @(Maybe a) l .==. Col @(Maybe a) r] + +isTextType :: forall a. (Columnable a) => Bool +isTextType = case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> True + Nothing -> False + +{- | 'discreteConditions' materialized with shared per-column reads: the +non-nullable categorical path builds truth vectors directly from one read +per column; nullable and cross-column fall back to interpret. +-} +discreteCondVecs :: + forall target. + (Columnable target, Ord target) => + TargetInfo target -> TreeConfig -> DataFrame -> [CondVec] +discreteCondVecs targetInfo cfg df = + concatMap (columnCondVecs (catCtx targetInfo cfg) df) (columnNames df) + ++ mapMaybe (materializeCondVec df) (crossColumnConds cfg df) + +columnCondVecs :: + (Columnable target, Ord target) => + CatCtx target -> DataFrame -> T.Text -> [CondVec] +columnCondVecs ctx df colName = case unsafeGetColumn colName df of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColCondVecs ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> mapMaybe (materializeCondVec df) (nullableColConds ctx colName bm column) + UnboxedColumn _ (_ :: VU.Vector a) -> [] + pt@(PackedText _ _) -> case materializePacked pt of + BoxedColumn Nothing (column :: V.Vector a) -> nonNullColCondVecs ctx colName column + BoxedColumn (Just bm) (column :: V.Vector a) -> mapMaybe (materializeCondVec df) (nullableColConds ctx colName bm column) + _ -> [] + +nonNullColCondVecs :: + forall a target. + (Columnable a, Ord target) => CatCtx target -> T.Text -> V.Vector a -> [CondVec] +nonNullColCondVecs ctx colName column = + fromMaybe + [] + ( withOrdFrom @a + (ccOrds ctx) + (map (membershipCondVec colName column) (catValueListsFor ctx column)) + ) + +membershipCondVec :: + forall a. (Columnable a, Ord a) => T.Text -> V.Vector a -> [a] -> CondVec +membershipCondVec colName column vs = CondVec (orEqs (eqExprFor @a colName) vs) (membershipVec column vs) diff --git a/src-internal/DataFrame/DecisionTree/CondVec.hs b/src-internal/DataFrame/DecisionTree/CondVec.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/CondVec.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Cached condition truth vectors and the per-fit cache keyed by structural +form. A condition's truth over a fixed DataFrame is invariant for a whole +fit, so it is materialized once and reused. +-} +module DataFrame.DecisionTree.CondVec ( + CondVec (..), + materializeCondVec, + CondCache, + condCacheKey, + condCacheFromVecs, + addTreeCondsToCache, + lookupCondVec, + partitionByVec, + countErrorsByVec, + consolidateThreshold, + combineAndVec, + combineOrVec, +) where + +import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..), Tree (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression ( + BinaryOp (binaryName), + Expr (..), + eqExpr, + normalize, + ) +import DataFrame.Internal.Interpreter (interpret) + +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | A boolean condition paired with its truth vector over the full DataFrame. +data CondVec = CondVec + { cvExpr :: !(Expr Bool) + , cvVec :: !(VU.Vector Bool) + } + +{- | Interpret a condition once over the DataFrame; 'Nothing' on a +type/interpret failure so the candidate is silently dropped. +-} +materializeCondVec :: DataFrame -> Expr Bool -> Maybe CondVec +materializeCondVec df cond = case interpret @Bool df cond of + Left _ -> Nothing + Right (TColumn column) -> CondVec cond <$> eitherToMaybe (toVector @Bool @VU.Vector column) + +eitherToMaybe :: Either e a -> Maybe a +eitherToMaybe = either (const Nothing) Just + +{- | Full-DataFrame truth vectors keyed by structural form, read-only once +built. Seeded for free from the candidate pool plus the initial tree so the +predict/loss passes index a vector instead of re-interpreting per node. +-} +type CondCache = M.Map T.Text (VU.Vector Bool) + +{- | Structural key matching the candidate-dedup key, so a tree branch whose +condition came from the pool hits the cache (equal keys ⟹ equal vector). +-} +condCacheKey :: Expr Bool -> T.Text +condCacheKey = T.pack . show . normalize + +-- | Seed a cache from already-materialized candidate 'CondVec's (no interpret). +condCacheFromVecs :: [CondVec] -> CondCache +condCacheFromVecs cvs = M.fromList [(condCacheKey (cvExpr cv), cvVec cv) | cv <- cvs] + +{- | Add a tree's branch-condition vectors to a cache (one interpret per +distinct, not-yet-cached condition). +-} +addTreeCondsToCache :: DataFrame -> Tree a -> CondCache -> CondCache +addTreeCondsToCache df = go + where + go (Leaf _) c = c + go (Branch cond l r) c = go r (go l (insertCond df cond c)) + +insertCond :: DataFrame -> Expr Bool -> CondCache -> CondCache +insertCond df cond c + | M.member k c = c + | otherwise = + maybe c (\cv -> M.insert k (cvVec cv) c) (materializeCondVec df cond) + where + k = condCacheKey cond + +{- | A condition's truth vector: a cache hit, else interpret over the +DataFrame. 'Nothing' mirrors the interpret-failure fallback (route left). +-} +lookupCondVec :: CondCache -> DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) +lookupCondVec cache df cond = case M.lookup (condCacheKey cond) cache of + hit@(Just _) -> hit + Nothing -> cvVec <$> materializeCondVec df cond + +-- | Partition row indices by a truth vector: @True@ → left, @False@ → right. +partitionByVec :: VU.Vector Bool -> V.Vector Int -> (V.Vector Int, V.Vector Int) +partitionByVec boolVals = V.partition (boolVals VU.!) + +-- | Count care points the truth vector routes to the wrong child. +countErrorsByVec :: VU.Vector Bool -> [CarePoint] -> Int +countErrorsByVec boolVals = length . filter misrouted + where + misrouted cp = (boolVals VU.! cpIndex cp) /= (cpCorrectDir cp == GoLeft) + +{- | A same-column same-direction Double threshold comparison, with a rebuild +function to swap in a new threshold. +-} +data ThreshCmp = ThreshCmp + { tcCol :: !T.Text + , tcName :: !T.Text + , tcThr :: !Double + , tcRebuild :: Double -> Expr Bool + } + +asDoubleThreshold :: Expr Bool -> Maybe ThreshCmp +asDoubleThreshold (Binary op (Col c :: Expr cc) (Lit (t :: tt))) = + case ( testEquality (typeRep @cc) (typeRep @Double) + , testEquality (typeRep @tt) (typeRep @Double) + ) of + (Just Refl, Just Refl) -> Just (ThreshCmp c (binaryName op) t (Binary op (Col c) . Lit)) + _ -> Nothing +asDoubleThreshold _ = Nothing + +directionalNames :: [T.Text] +directionalNames = ["lt", "leq", "gt", "geq"] + +{- | Tighter (AND) or looser (OR) of two same-direction thresholds: @<@/@<=@ +are left-half-spaces (AND = min), @>@/@>=@ are right-half-spaces (AND = max). +-} +chooseThreshold :: Bool -> T.Text -> Double -> Double -> Double +chooseThreshold isAnd name t1 t2 + | leftDir = if isAnd then min t1 t2 else max t1 t2 + | otherwise = if isAnd then max t1 t2 else min t1 t2 + where + leftDir = name == "lt" || name == "leq" + +{- | Collapse two same-column same-direction strict-Double comparisons into one +comparison (the @True@ argument selects AND, @False@ OR); 'Nothing' otherwise. +-} +consolidateThreshold :: Bool -> Expr Bool -> Expr Bool -> Maybe (Expr Bool) +consolidateThreshold isAnd ea eb = do + a <- asDoubleThreshold ea + b <- asDoubleThreshold eb + if tcCol a == tcCol b && tcName a == tcName b && tcName a `elem` directionalNames + then Just (tcRebuild a (chooseThreshold isAnd (tcName a) (tcThr a) (tcThr b))) + else Nothing + +{- | AND-combine two cached conditions: idempotence and threshold consolidation +first, else the generic @F.and@; the vector is always the elementwise AND. +-} +combineAndVec :: CondVec -> CondVec -> CondVec +combineAndVec a b + | eqExpr (cvExpr a) (cvExpr b) = a + | otherwise = CondVec expr (VU.zipWith (&&) (cvVec a) (cvVec b)) + where + expr = + fromMaybe + (F.and (cvExpr a) (cvExpr b)) + (consolidateThreshold True (cvExpr a) (cvExpr b)) + +{- | OR-combine two cached conditions (see 'combineAndVec'; AND/OR direction +differs in 'consolidateThreshold'). +-} +combineOrVec :: CondVec -> CondVec -> CondVec +combineOrVec a b + | eqExpr (cvExpr a) (cvExpr b) = a + | otherwise = CondVec expr (VU.zipWith (||) (cvVec a) (cvVec b)) + where + expr = + fromMaybe + (F.or (cvExpr a) (cvExpr b)) + (consolidateThreshold False (cvExpr a) (cvExpr b)) diff --git a/src-internal/DataFrame/DecisionTree/Fit.hs b/src-internal/DataFrame/DecisionTree/Fit.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Fit.hs @@ -0,0 +1,213 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Top-level fitting: assemble the candidate pool, seed from CART, run TAO, +and convert the result to an expression. Also the probability-tree variant +('fitProbTree') that annotates leaves with class distributions. +-} +module DataFrame.DecisionTree.Fit ( + treeToExpr, + fitDecisionTree, + buildTree, + pruneTree, + partitionDataFrame, + calculateGini, + majorityValue, + getCounts, + percentile, + ProbTree, + probsFromIndices, + buildProbTree, + fitProbTree, + probExprs, +) where + +import DataFrame.DecisionTree.Cart (buildCartTree) +import DataFrame.DecisionTree.Categorical ( + TargetInfo (..), + discreteCondVecs, + discreteConditions, + mkTargetInfo, + ) +import DataFrame.DecisionTree.CondVec (CondVec) +import DataFrame.DecisionTree.Numeric (numericCondVecs, numericConditions) +import DataFrame.DecisionTree.Pool (dedupCVByExpr, nubByExpr) +import DataFrame.DecisionTree.Predict (partitionIndices) +import DataFrame.DecisionTree.Prune (pruneDead, pruneExpr) +import DataFrame.DecisionTree.Tao (taoOptimize, taoOptimizeCV) +import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Operations.Core (nRows) +import DataFrame.Operations.Subset (exclude, filterWhere) + +import Control.Exception (throw) +import Data.Function (on) +import Data.List (foldl', maximumBy, nub, sort) +import qualified Data.Map.Strict as M +import Data.Maybe (fromMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V + +-- | Convert a fitted tree to a nested-conditional expression. +treeToExpr :: (Columnable a) => Tree a -> Expr a +treeToExpr (Leaf v) = Lit v +treeToExpr (Branch cond left right) = F.ifThenElse cond (treeToExpr left) (treeToExpr right) + +-- | Fit a TAO decision tree (CART-seeded) and return it as an expression. +fitDecisionTree :: + forall a. (Columnable a, Ord a) => TreeConfig -> Expr a -> DataFrame -> Expr a +fitDecisionTree cfg (Col target) df = + pruneExpr + (treeToExpr (taoOptimizeCV @a cfg target condVecs df indices initialTree)) + where + condVecs = candidatePool @a cfg target df + initialTree = buildCartTree @a cfg target df + indices = V.enumFromN 0 (nRows df) +fitDecisionTree _ expr _ = error ("Cannot create tree for compound expression: " ++ show expr) + +-- | The deduplicated numeric + discrete candidate pool for a target column. +candidatePool :: + forall a. + (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> [CondVec] +candidatePool cfg target df = dedupCVByExpr (numericCVs ++ discreteCVs) + where + dfNoTarget = exclude [target] df + numericCVs = numericCondVecs cfg dfNoTarget df + discreteCVs = discreteCondVecs (targetInfoOrEmpty @a target df) cfg dfNoTarget + +targetInfoOrEmpty :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> TargetInfo a +targetInfoOrEmpty target df = fromMaybe (TargetInfo False Nothing V.empty) (mkTargetInfo @a target df) + +-- | Fit a tree at a given depth from a raw condition list (CART + TAO + prune). +buildTree :: + forall a. + (Columnable a, Ord a) => + TreeConfig -> Int -> T.Text -> [Expr Bool] -> DataFrame -> Expr a +buildTree cfg depth target conds df = + pruneExpr (treeToExpr (taoOptimize @a cfg target conds df indices tree)) + where + tree = buildCartTree @a cfg{maxTreeDepth = depth} target df + indices = V.enumFromN 0 (nRows df) + +pruneTree :: forall a. (Columnable a) => Expr a -> Expr a +pruneTree = pruneExpr + +partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) +partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) + +-- | Laplace-smoothed Gini impurity of the target distribution. +calculateGini :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double +calculateGini target df + | n == 0 = 0 + | otherwise = 1 - sum (map (^ (2 :: Int)) probs) + where + n = fromIntegral (nRows df) + counts = getCounts @a target df + numClasses = fromIntegral (M.size counts) + probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts) + +majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a +majorityValue target df + | M.null counts = error "Empty DataFrame in leaf" + | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) + where + counts = getCounts @a target df + +getCounts :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int +getCounts target df = case interpret @a df (Col target) of + Left e -> throw e + Right (TColumn column) -> case toVector @a column of + Left e -> throw e + Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals) + +-- | The @p@-th percentile of an expression's values (@0@ on failure/empty). +percentile :: Int -> Expr Double -> DataFrame -> Double +percentile p expr df = case interpret @Double df expr of + Right (TColumn column) -> either (const 0) (percentileOfVec p) (toVector @Double column) + _ -> 0 + +percentileOfVec :: Int -> V.Vector Double -> Double +percentileOfVec p vals + | n == 0 = 0 + | otherwise = sorted V.! min (n - 1) (max 0 ((p * n) `div` 100)) + where + sorted = V.fromList (sort (V.toList vals)) + n = V.length sorted + +-- | A tree whose leaves hold class-probability distributions. +type ProbTree a = Tree (M.Map a Double) + +-- | Normalised class probabilities over a subset of training rows. +probsFromIndices :: + forall a. + (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> M.Map a Double +probsFromIndices target df indices = case interpret @a df (Col target) of + Right (TColumn column) -> either (const M.empty) (normaliseCounts indices) (toVector @a column) + _ -> M.empty + +normaliseCounts :: (Ord a) => V.Vector Int -> V.Vector a -> M.Map a Double +normaliseCounts indices vals = M.map (\c -> fromIntegral c / total) counts + where + counts = + V.foldl' + (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) + M.empty + indices + total = fromIntegral (V.length indices) :: Double + +{- | Re-label a fitted tree's leaves with class distributions, routing the +training data through the (unchanged) split conditions. +-} +buildProbTree :: + forall a. + (Columnable a, Ord a) => + Tree a -> T.Text -> DataFrame -> V.Vector Int -> ProbTree a +buildProbTree (Leaf _) target df indices = Leaf (probsFromIndices @a target df indices) +buildProbTree (Branch cond left right) target df indices = + Branch + cond + (buildProbTree @a left target df l) + (buildProbTree @a right target df r) + where + (l, r) = partitionIndices cond df indices + +-- | Fit a TAO tree and return one probability expression per class. +fitProbTree :: + forall a. + (Columnable a, Ord a) => + TreeConfig -> Expr a -> DataFrame -> M.Map a (Expr Double) +fitProbTree cfg (Col target) df = probExprs (buildProbTree @a pruned target df indices) + where + conds = + nubByExpr + ( numericConditions cfg dfNoTarget + ++ discreteConditions (targetInfoOrEmpty @a target df) cfg dfNoTarget + ) + dfNoTarget = exclude [target] df + indices = V.enumFromN 0 (nRows df) + pruned = + pruneDead + (taoOptimize @a cfg target conds df indices (buildCartTree @a cfg target df)) +fitProbTree _ expr _ = error ("Cannot create prob tree for compound expression: " ++ show expr) + +-- | Convert a 'ProbTree' into one @Expr Double@ per class. +probExprs :: + forall a. (Columnable a, Ord a) => ProbTree a -> M.Map a (Expr Double) +probExprs tree = M.fromList [(c, classExpr c tree) | c <- nub (allClasses tree)] + +allClasses :: ProbTree a -> [a] +allClasses (Leaf m) = M.keys m +allClasses (Branch _ l r) = allClasses l ++ allClasses r + +classExpr :: (Ord a) => a -> ProbTree a -> Expr Double +classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m) +classExpr c (Branch cond l r) = F.ifThenElse cond (classExpr c l) (classExpr c r) diff --git a/src-internal/DataFrame/DecisionTree/Linear.hs b/src-internal/DataFrame/DecisionTree/Linear.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Linear.hs @@ -0,0 +1,161 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | Oblique split candidates: fit an L1-regularised logistic hyperplane to the +care points (class-balanced) and convert it to a boolean condition, rejecting +all-zero and degenerate (single-side) hyperplanes. +-} +module DataFrame.DecisionTree.Linear ( + bestLinearCandidate, + fitLinearCandidate, + careRowsFromFeatures, + careLabels, + featName, + imputeMean, + materializeFeatureForCare, +) where + +import DataFrame.DecisionTree.Numeric (NumExpr (..), numericCols) +import DataFrame.DecisionTree.Types ( + CarePoint (..), + Direction (..), + TreeConfig (..), + ) +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr, getColumns) +import DataFrame.Internal.Interpreter (interpret) +import qualified DataFrame.LinearSolver as LS + +import Data.Maybe (catMaybes, fromMaybe, mapMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +{- | Best oblique candidate, or 'Nothing' when the linear path is disabled or +there are too few care points to fit on. +-} +bestLinearCandidate :: + TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) +bestLinearCandidate cfg df carePoints + | not (useLinearSolver cfg) = Nothing + | length carePoints < minCarePointsForLinear cfg = Nothing + | otherwise = fitLinearCandidate cfg df carePoints + +{- | Fit an L1 logistic regression to the care points and convert the resulting +hyperplane to a condition, or 'Nothing' when no numeric features exist or the +fitted model is all-zero or degenerate. +-} +fitLinearCandidate :: + TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) +fitLinearCandidate cfg df carePoints = case materializedFeatures df carePoints of + [] -> Nothing + mats -> linearFromFeatures cfg carePoints mats + +materializedFeatures :: DataFrame -> [CarePoint] -> [(T.Text, VU.Vector Double)] +materializedFeatures df carePoints = mapMaybe (materializeFeatureForCare df carePoints) (numericCols df) + +linearFromFeatures :: + TreeConfig -> [CarePoint] -> [(T.Text, VU.Vector Double)] -> Maybe (Expr Bool) +linearFromFeatures cfg carePoints mats + | VU.all (== 0) weights = Nothing + | degenerateHyperplane rows weights (LS.lmIntercept model) = Nothing + | otherwise = Just (LS.modelToExpr model) + where + rows = careRowsFromFeatures (length carePoints) mats + labels = careLabels carePoints + model = + LS.fitL1Logistic + (solverConfigFor cfg labels) + rows + labels + (V.fromList (map fst mats)) + weights = LS.lmWeights model + +solverConfigFor :: TreeConfig -> VU.Vector Double -> LS.SolverConfig +solverConfigFor cfg labels = (linearSolverConfig cfg){LS.scSampleWeights = classBalancedWeights labels} + +{- | Class-balanced sklearn-form weights @w_i = N / (2 · N_class)@ (mean 1), or +'Nothing' in the degenerate one-class case (uniform weighting). +-} +classBalancedWeights :: VU.Vector Double -> Maybe (VU.Vector Double) +classBalancedWeights labels + | nPos > 0 && nNeg > 0 = Just (VU.generate nCare weightAt) + | otherwise = Nothing + where + nCare = VU.length labels + nPos = VU.length (VU.filter (> 0) labels) + nNeg = nCare - nPos + weightAt i + | VU.unsafeIndex labels i > 0 = fromIntegral nCare / (2 * fromIntegral nPos) + | otherwise = fromIntegral nCare / (2 * fromIntegral nNeg) + +{- | A hyperplane is degenerate when every care row scores on the same side of +zero (equivalent to an invalid split, caught upstream). +-} +degenerateHyperplane :: + V.Vector (VU.Vector Double) -> VU.Vector Double -> Double -> Bool +degenerateHyperplane rows weights bias = + nCare > 0 && (VU.minimum scores > 0 || VU.maximum scores < 0) + where + nCare = V.length rows + scores = + VU.generate + nCare + (\i -> VU.sum (VU.zipWith (*) weights (V.unsafeIndex rows i)) + bias) + +{- | Per-care-point feature rows from materialized columns (each of length +@nCare@, so indexing is in range). +-} +careRowsFromFeatures :: + Int -> [(T.Text, VU.Vector Double)] -> V.Vector (VU.Vector Double) +careRowsFromFeatures nCare mats = + V.generate nCare (\i -> VU.generate nFeat (\j -> snd (matsVec V.! j) VU.! i)) + where + matsVec = V.fromList mats + nFeat = V.length matsVec + +-- | Solver labels: @+1@ when 'GoLeft' is correct, @-1@ otherwise. +careLabels :: [CarePoint] -> VU.Vector Double +careLabels carePoints = + VU.fromList [if cpCorrectDir cp == GoLeft then 1.0 else -1.0 | cp <- carePoints] + +-- | First column referenced by an expression, or a placeholder when none. +featName :: Expr b -> T.Text +featName expr = case getColumns expr of + (c : _) -> c + [] -> "" + +{- | Replace missing values with the mean of present ones; 'Nothing' when +nothing is present so the caller can drop the feature. +-} +imputeMean :: [Maybe Double] -> Maybe (VU.Vector Double) +imputeMean careRaw = case catMaybes careRaw of + [] -> Nothing + present -> Just (VU.fromList [fromMaybe (mean present) mv | mv <- careRaw]) + where + mean xs = sum xs / fromIntegral (length xs) + +interpretDoubleVals :: DataFrame -> Expr Double -> Maybe (V.Vector Double) +interpretDoubleVals df expr = case interpret @Double df expr of + Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) + _ -> Nothing + +interpretMaybeDoubleVals :: + DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) +interpretMaybeDoubleVals df expr = case interpret @(Maybe Double) df expr of + Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) + _ -> Nothing + +{- | Materialize a 'NumExpr' over the care rows; 'Nothing' on interpret failure +or (nullable) when no care point has a present value, else mean-imputed. +-} +materializeFeatureForCare :: + DataFrame -> [CarePoint] -> NumExpr -> Maybe (T.Text, VU.Vector Double) +materializeFeatureForCare df carePoints (NDouble expr) = do + vals <- interpretDoubleVals df expr + Just (featName expr, VU.fromList [vals V.! cpIndex cp | cp <- carePoints]) +materializeFeatureForCare df carePoints (NMaybeDouble expr) = do + vals <- interpretMaybeDoubleVals df expr + imputed <- imputeMean [vals V.! cpIndex cp | cp <- carePoints] + Just (featName expr, imputed) diff --git a/src-internal/DataFrame/DecisionTree/Numeric.hs b/src-internal/DataFrame/DecisionTree/Numeric.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Numeric.hs @@ -0,0 +1,254 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Numeric split candidates: per-column Double expressions, arithmetic +expansion, and threshold conditions. 'numericCondVecs' materializes the +pool with one interpret per distinct expression. +-} +module DataFrame.DecisionTree.Numeric ( + NumExpr (..), + numExprCols, + numExprEq, + combineNumExprs, + numericConditions, + generateNumericConds, + percentilesOf, + numericCondVecs, + numericExprsWithTerms, + numericCols, +) where + +import DataFrame.DecisionTree.CondVec (CondVec (..)) +import DataFrame.DecisionTree.Types (SynthConfig (..), TreeConfig (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) +import DataFrame.Internal.Expression (Expr (..), eqExpr, getColumns, normalize) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Types +import DataFrame.Operators + +import Data.List (sort) +import Data.Maybe (fromMaybe) +import qualified Data.Set as Set +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (typeRep) + +-- | A numeric feature expression, non-nullable or nullable. +data NumExpr + = NDouble !(Expr Double) + | NMaybeDouble !(Expr (Maybe Double)) + +numExprCols :: NumExpr -> [T.Text] +numExprCols (NDouble e) = getColumns e +numExprCols (NMaybeDouble e) = getColumns e + +numExprEq :: NumExpr -> NumExpr -> Bool +numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2 +numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2 +numExprEq _ _ = False + +-- | Safe division: @0@ (or @Nothing@) where the divisor is zero. +safeDivD :: Expr Double -> Expr Double -> Expr Double +safeDivD a b = F.ifThenElse (b ./= F.lit (0 :: Double)) (a ./ b) (F.lit (0 :: Double)) + +safeDivMaybe :: Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double) +safeDivMaybe nonZero q = F.ifThenElse nonZero q (F.lit (Nothing :: Maybe Double)) + +-- | Arithmetic combinations (@+@, @-@, @*@, safe @/@) of two numeric exprs. +combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] +combineNumExprs (NDouble e1) (NDouble e2) = + map NDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivD e1 e2] +combineNumExprs (NDouble e1) (NMaybeDouble e2) = + map + NMaybeDouble + [ e1 .+ e2 + , e1 .- e2 + , e1 .* e2 + , safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2) + ] +combineNumExprs (NMaybeDouble e1) (NDouble e2) = + map + NMaybeDouble + [ e1 .+ e2 + , e1 .- e2 + , e1 .* e2 + , safeDivMaybe (e2 ./= F.lit (0 :: Double)) (e1 ./ e2) + ] +combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) = + map + NMaybeDouble + [ e1 .+ e2 + , e1 .- e2 + , e1 .* e2 + , safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2) + ] + +numericConditions :: TreeConfig -> DataFrame -> [Expr Bool] +numericConditions = generateNumericConds + +generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool] +generateNumericConds cfg df = do + expr <- numericExprsWithTerms (synthConfig cfg) df + threshold <- numericThresholds cfg df expr + condsFromExpr expr threshold + +numericThresholds :: TreeConfig -> DataFrame -> NumExpr -> [Double] +numericThresholds cfg df (NDouble e) = thresholdsForExpr cfg df e +numericThresholds cfg df (NMaybeDouble e) = thresholdsForExpr cfg df (F.fromMaybe 0 e) + +thresholdsForExpr :: TreeConfig -> DataFrame -> Expr Double -> [Double] +thresholdsForExpr cfg df e = + maybe [] (percentilesOf (percentiles cfg) . V.toList) (interpretDoubleCol df e) + +condsFromExpr :: NumExpr -> Double -> [Expr Bool] +condsFromExpr (NDouble e) t = [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] +condsFromExpr (NMaybeDouble e) t = + map + (F.fromMaybe False) + [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] + +{- | Percentile thresholds for a value list: sort once, index each percentile. +Shared by 'generateNumericConds' and 'numericCondVecs' for identical results. +-} +percentilesOf :: [Int] -> [Double] -> [Double] +percentilesOf ps valsList + | n == 0 = [] + | otherwise = map (\p -> sortedV V.! min (n - 1) (max 0 (p * n `div` 100))) ps + where + !sortedV = V.fromList (sort valsList) + !n = V.length sortedV + +interpretDoubleCol :: DataFrame -> Expr Double -> Maybe (V.Vector Double) +interpretDoubleCol df e = case interpret @Double df e of + Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) + _ -> Nothing + +interpretMaybeDoubleCol :: + DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) +interpretMaybeDoubleCol df e = case interpret @(Maybe Double) df e of + Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) + _ -> Nothing + +{- | Materialize the numeric pool with one interpret per distinct expression, +deriving each threshold/operator truth vector by direct comparison. +Byte-identical to materializing 'numericConditions' one at a time. +-} +numericCondVecs :: TreeConfig -> DataFrame -> DataFrame -> [CondVec] +numericCondVecs cfg dfGen df = concatMap forExpr (numericExprsWithTerms (synthConfig cfg) dfGen) + where + forExpr (NDouble e) = maybe [] (condsForDouble cfg e) (interpretDoubleCol df e) + forExpr (NMaybeDouble e) = maybe [] (condsForMaybe cfg e) (interpretMaybeDoubleCol df e) + +condsForDouble :: TreeConfig -> Expr Double -> V.Vector Double -> [CondVec] +condsForDouble cfg e vals = concatMap (doubleCondsAt e vals (V.length vals)) ts + where + ts = percentilesOf (percentiles cfg) (V.toList vals) + +doubleCondsAt :: Expr Double -> V.Vector Double -> Int -> Double -> [CondVec] +doubleCondsAt e vals n t = + [ CondVec (e .<= F.lit t) (gen (<= t)) + , CondVec (e .>= F.lit t) (gen (>= t)) + , CondVec (e .< F.lit t) (gen (< t)) + , CondVec (e .> F.lit t) (gen (> t)) + ] + where + gen p = VU.generate n (\i -> p (vals V.! i)) + +condsForMaybe :: + TreeConfig -> Expr (Maybe Double) -> V.Vector (Maybe Double) -> [CondVec] +condsForMaybe cfg e mvals = concatMap (maybeCondsAt e mvals (V.length mvals)) ts + where + ts = percentilesOf (percentiles cfg) (map (fromMaybe 0) (V.toList mvals)) + +maybeCondsAt :: + Expr (Maybe Double) -> V.Vector (Maybe Double) -> Int -> Double -> [CondVec] +maybeCondsAt e mvals n t = + [ CondVec (F.fromMaybe False (e .<= F.lit t)) (gen (<= t)) + , CondVec (F.fromMaybe False (e .>= F.lit t)) (gen (>= t)) + , CondVec (F.fromMaybe False (e .< F.lit t)) (gen (< t)) + , CondVec (F.fromMaybe False (e .> F.lit t)) (gen (> t)) + ] + where + gen p = VU.generate n (\i -> maybe False p (mvals V.! i)) + +{- | Arithmetic candidate expansion, generated already-deduped: each round +combines @frontier × base@ and admits only normalized-novel candidates. +Produces @base@ plus @maxExprDepth-1@ combination rounds. +-} +numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr] +numericExprsWithTerms cfg df + | not (enableArithOps cfg) = base + | otherwise = + base ++ expandRounds cfg base (max 0 (maxExprDepth cfg - 1)) base seen0 + where + base = numericCols df + seen0 = Set.fromList (map keyNum base) + +keyNum :: NumExpr -> String +keyNum (NDouble e) = show (normalize e) +keyNum (NMaybeDouble e) = show (normalize e) + +isDisallowed :: SynthConfig -> NumExpr -> NumExpr -> Bool +isDisallowed cfg e1 e2 = + any (\(l, r) -> l `elem` cols && r `elem` cols) (disallowedCombinations cfg) + where + cols = numExprCols e1 <> numExprCols e2 + +roundProducts :: SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr] +roundProducts cfg frontier base = + [ c + | e1 <- frontier + , e2 <- base + , not (numExprEq e1 e2) + , not (isDisallowed cfg e1 e2) + , c <- combineNumExprs e1 e2 + ] + +expandRounds :: + SynthConfig -> [NumExpr] -> Int -> [NumExpr] -> Set.Set String -> [NumExpr] +expandRounds _ _ 0 _ _ = [] +expandRounds cfg base d frontier seen + | null admitted = [] + | otherwise = admitted ++ expandRounds cfg base (d - 1) admitted seen' + where + (admitted, seen') = admitNovel seen (roundProducts cfg frontier base) + +admitNovel :: Set.Set String -> [NumExpr] -> ([NumExpr], Set.Set String) +admitNovel seen0 = go seen0 [] + where + go seen acc [] = (reverse acc, seen) + go seen acc (c : cs) + | keyNum c `Set.member` seen = go seen acc cs + | otherwise = go (Set.insert (keyNum c) seen) (c : acc) cs + +numericCols :: DataFrame -> [NumExpr] +numericCols df = concatMap (numExprsOfColumn df) (columnNames df) + +numExprsOfColumn :: DataFrame -> T.Text -> [NumExpr] +numExprsOfColumn df colName = case unsafeGetColumn colName df of + UnboxedColumn Nothing (_ :: VU.Vector b) -> strictNumeric @b colName + BoxedColumn (Just _) (_ :: V.Vector b) -> nullableNumeric @b colName + UnboxedColumn (Just _) (_ :: VU.Vector b) -> nullableNumeric @b colName + _ -> [] + +strictNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] +strictNumeric c = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [NDouble (Col c)] + Nothing -> case sIntegral @b of + STrue -> [NDouble (F.toDouble (Col @b c))] + SFalse -> [] + +nullableNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] +nullableNumeric c = case testEquality (typeRep @b) (typeRep @Double) of + Just Refl -> [NMaybeDouble (Col @(Maybe b) c)] + Nothing -> case sIntegral @b of + STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) c))] + SFalse -> [] diff --git a/src-internal/DataFrame/DecisionTree/Pool.hs b/src-internal/DataFrame/DecisionTree/Pool.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Pool.hs @@ -0,0 +1,224 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +{- | Candidate-pool scoring and boolean expansion: penalized scoring, diverse +top-K selection, AND/OR saturation, and structural/truth-vector dedup. The +per-node scoring scans run in parallel chunks. +-} +module DataFrame.DecisionTree.Pool ( + evalWithPenaltyVec, + primaryColExpr, + primaryColCV, + takeDiverse, + candidateParChunk, + bestDiscreteCandidate, + boolExprsVec, + DedupMode (..), + saturateCandidates, + roundProducts, + admitKeys, + admitVecs, + dedupCVByExpr, + nubByExpr, +) where + +import DataFrame.DecisionTree.CondVec ( + CondVec (..), + combineAndVec, + combineOrVec, + countErrorsByVec, + ) +import DataFrame.DecisionTree.Types ( + CarePoint, + SynthConfig (..), + TreeConfig (..), + ) +import DataFrame.Internal.Expression ( + Expr, + compareExpr, + eSize, + eqExpr, + getColumns, + normalize, + ) + +import Control.Parallel.Strategies (parListChunk, rdeepseq, using) +import Data.Function (on) +import Data.List (minimumBy, sortBy) +import qualified Data.Map.Strict as M +import qualified Data.Set as Set +import qualified Data.Text as T +import qualified Data.Vector.Unboxed as VU + +{- | Penalized score of a candidate: care-point errors plus a complexity +penalty, tie-broken by expression size. +-} +evalWithPenaltyVec :: TreeConfig -> [CarePoint] -> CondVec -> (Int, Int) +evalWithPenaltyVec cfg carePoints cv = (countErrorsByVec (cvVec cv) carePoints + penalty, sz) + where + sz = eSize (cvExpr cv) + penalty = floor (complexityPenalty (synthConfig cfg) * fromIntegral sz) + +{- | First referenced column of a condition (a sentinel for literal-only ones), +used by 'takeDiverse' to enforce per-column diversity. +-} +primaryColExpr :: Expr Bool -> T.Text +primaryColExpr e = case getColumns e of + [] -> "" + (c : _) -> c + +primaryColCV :: CondVec -> T.Text +primaryColCV = primaryColExpr . cvExpr + +{- | Keep the first @k@ of an already-sorted list, admitting at most @quota@ per +primary column (@Nothing@ disables the per-column cap). +-} +takeDiverse :: Int -> Maybe Int -> (a -> T.Text) -> [a] -> [a] +takeDiverse k Nothing _ = take k +takeDiverse k (Just quota) primary = go M.empty 0 + where + go !_ !_ [] = [] + go !seen !n (x : xs) + | n >= k = [] + | M.findWithDefault 0 col seen >= quota = go seen n xs + | otherwise = x : go (M.insertWith (+) col 1 seen) (n + 1) xs + where + !col = primary x + +{- | Chunk size for the parallel per-node candidate scans; tuned by an -N +sweep, not correctness-affecting. +-} +candidateParChunk :: Int +candidateParChunk = 64 + +{- | Decorate candidates with their penalty in parallel chunks, forcing only +the @(Int, Int)@ key so the order (hence later sorts/minima) is preserved. +-} +decorate :: (CondVec -> (Int, Int)) -> [CondVec] -> [((Int, Int), CondVec)] +decorate penaltyCV xs = zip (map penaltyCV xs `using` parListChunk candidateParChunk rdeepseq) xs + +-- | The diverse top-@expressionPairs@ valid candidates by penalty. +sortedTopK :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> [CondVec] +sortedTopK cfg penaltyCV validCondVecs = + map + snd + ( takeDiverse + (expressionPairs cfg) + (perColumnQuota (synthConfig cfg)) + (primaryColCV . snd) + sorted + ) + where + sorted = sortBy (compare `on` fst) (decorate penaltyCV validCondVecs) + +-- | Lowest-penalty candidate after boolean saturation of the diverse top-K. +bestDiscreteCandidate :: + TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> Maybe CondVec +bestDiscreteCandidate _ _ [] = Nothing +bestDiscreteCandidate cfg penaltyCV validCondVecs = + case saturateCandidates + Structural + (boolExpansion (synthConfig cfg)) + (sortedTopK cfg penaltyCV validCondVecs) of + [] -> Nothing + xs -> Just (snd (minimumBy (compare `on` fst) (decorate penaltyCV xs))) + +{- | AND/OR expansion of cached conditions to depth @maxDepth@ (each +combination is a single vector op, not an interpret). +-} +boolExprsVec :: [CondVec] -> [CondVec] -> Int -> Int -> [CondVec] +boolExprsVec baseExprs prevExprs depth maxDepth + | depth == 0 = + baseExprs ++ boolExprsVec baseExprs prevExprs (depth + 1) maxDepth + | depth >= maxDepth = [] + | otherwise = combined ++ boolExprsVec baseExprs combined (depth + 1) maxDepth + where + combined = roundProducts prevExprs baseExprs + +data DedupMode = Structural | TruthVector + deriving (Eq, Show) + +{- | Saturate the pool with AND/OR combinations, deduplicating structurally +(byte-identical, first occurrence kept) or by truth vector (opt-in). +-} +saturateCandidates :: DedupMode -> Int -> [CondVec] -> [CondVec] +saturateCandidates Structural maxDepth base = base' ++ go 1 base' seen0 + where + (base', seen0) = admitKeys Set.empty base + go !depth frontier seen + | depth >= maxDepth || null frontier = [] + | otherwise = + let (admitted, seen') = admitKeys seen (roundProducts frontier base) + in admitted ++ go (depth + 1) admitted seen' +saturateCandidates TruthVector maxDepth base = M.elems (go 1 frontier0 reps0) + where + (reps0, frontier0) = admitVecs M.empty base + go !depth frontier reps + | depth >= maxDepth || null frontier = reps + | otherwise = + let (reps', admitted) = admitVecs reps (roundProducts frontier base) + in go (depth + 1) admitted reps' + +{- | One combination round: @frontier × base@ via AND then OR, skipping +self-pairs (mirrors 'boolExprsVec' for byte-identical structural output). +-} +roundProducts :: [CondVec] -> [CondVec] -> [CondVec] +roundProducts frontier base = + [ c + | e1 <- frontier + , e2 <- base + , not (eqExpr (cvExpr e1) (cvExpr e2)) + , c <- [combineAndVec e1 e2, combineOrVec e1 e2] + ] + +-- | Admit candidates with a not-yet-seen normalized form, preserving order. +admitKeys :: Set.Set String -> [CondVec] -> ([CondVec], Set.Set String) +admitKeys = go [] + where + go acc seen [] = (reverse acc, seen) + go acc !seen (c : cs) + | structuralKey c `Set.member` seen = go acc seen cs + | otherwise = go (c : acc) (Set.insert (structuralKey c) seen) cs + +structuralKey :: CondVec -> String +structuralKey = show . normalize . cvExpr + +{- | Admit candidates by distinct truth vector, keeping the smallest-expression +representative per vector. +-} +admitVecs :: + M.Map (VU.Vector Bool) CondVec -> + [CondVec] -> + (M.Map (VU.Vector Bool) CondVec, [CondVec]) +admitVecs = go [] + where + go acc reps [] = (reps, reverse acc) + go acc !reps (c : cs) = case M.lookup (cvVec c) reps of + Nothing -> go (c : acc) (M.insert (cvVec c) c reps) cs + Just r -> go acc (M.insert (cvVec c) (smaller r c) reps) cs + +smaller :: CondVec -> CondVec -> CondVec +smaller a b = case compare (eSize (cvExpr a)) (eSize (cvExpr b)) of + LT -> a + GT -> b + EQ -> if compareExpr (cvExpr a) (cvExpr b) /= GT then a else b + +-- | Deduplicate 'CondVec's by normalized 'cvExpr', keeping the first. +dedupCVByExpr :: [CondVec] -> [CondVec] +dedupCVByExpr = go Set.empty + where + go _ [] = [] + go seen (cv : cvs) + | structuralKey cv `Set.member` seen = go seen cvs + | otherwise = cv : go (Set.insert (structuralKey cv) seen) cvs + +-- | Deduplicate expressions by normalized form, keeping the first. +nubByExpr :: [Expr Bool] -> [Expr Bool] +nubByExpr = go Set.empty + where + go _ [] = [] + go seen (e : es) + | k `Set.member` seen = go seen es + | otherwise = e : go (Set.insert k seen) es + where + k = show (normalize e) diff --git a/src-internal/DataFrame/DecisionTree/Predict.hs b/src-internal/DataFrame/DecisionTree/Predict.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Predict.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Prediction, care-point identification, node validity, and tree loss. The +batched, cache-aware variants resolve each branch condition's truth vector +once per call instead of once per row. +-} +module DataFrame.DecisionTree.Predict ( + predictWithTree, + predictManyWithTree, + predictManyWithTreeCached, + identifyCarePoints, + identifyCarePointsCached, + countCarePointErrors, + partitionIndices, + partitionIndicesCached, + majorityValueFromIndices, + computeTreeLoss, + computeTreeLossCached, + isValidAtNode, +) where + +import DataFrame.DecisionTree.CondVec ( + CondCache, + countErrorsByVec, + lookupCondVec, + ) +import DataFrame.DecisionTree.Types ( + CarePoint (..), + Direction (..), + Tree (..), + TreeConfig (..), + ) +import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Internal.Interpreter (interpret) + +import Control.Exception (throw) +import Control.Monad.ST (ST) +import Data.Function (on) +import Data.List (maximumBy) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU + +{- | A condition's truth vector over the DataFrame, or 'Nothing' on a +type/interpret failure (callers default such rows to the left child). +-} +branchBool :: DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) +branchBool df cond = case interpret @Bool df cond of + Right (TColumn column) -> either (const Nothing) Just (toVector @Bool @VU.Vector column) + _ -> Nothing + +-- | The target column as a label vector, or 'Nothing' on failure. +interpretLabelCol :: + forall a. (Columnable a) => DataFrame -> T.Text -> Maybe (V.Vector a) +interpretLabelCol df target = case interpret @a df (Col target) of + Right (TColumn column) -> either (const Nothing) Just (toVector @a column) + _ -> Nothing + +-- | Predict the label for a single row by walking a fixed tree (@True@ → left). +predictWithTree :: + forall a. (Columnable a) => T.Text -> DataFrame -> Int -> Tree a -> a +predictWithTree _ _ _ (Leaf v) = v +predictWithTree target df idx (Branch cond left right) = + predictWithTree @a target df idx (childFor cond left right idx df) + +childFor :: Expr Bool -> Tree a -> Tree a -> Int -> DataFrame -> Tree a +childFor cond left right idx df = case branchBool df cond of + Nothing -> left + Just boolVals -> if boolVals VU.! idx then left else right + +predictManyWithTree :: + forall a. (Columnable a) => Tree a -> DataFrame -> V.Vector Int -> V.Vector a +predictManyWithTree = predictManyWithTreeCached @a M.empty + +{- | 'predictManyWithTree' resolving each branch condition through a 'CondCache'. +Each condition is read at most once per call rather than once per row. +-} +predictManyWithTreeCached :: + forall a. + (Columnable a) => CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a +predictManyWithTreeCached cache tree df indices = V.create $ do + mv <- VM.new (V.length indices) + fill mv (V.zip (V.enumFromN 0 (V.length indices)) indices) tree + pure mv + where + fill :: VM.MVector s a -> V.Vector (Int, Int) -> Tree a -> ST s () + fill mv prs (Leaf v) = V.mapM_ (\(p, _) -> VM.write mv p v) prs + fill mv prs (Branch cond left right) = case lookupCondVec cache df cond of + Nothing -> fill mv prs left + Just boolVals -> fillSplit mv (V.partition (\(_, i) -> boolVals VU.! i) prs) left right + + fillSplit :: + VM.MVector s a -> + (V.Vector (Int, Int), V.Vector (Int, Int)) -> + Tree a -> + Tree a -> + ST s () + fillSplit mv (leftPrs, rightPrs) left right = fill mv leftPrs left >> fill mv rightPrs right + +identifyCarePoints :: + forall a. + (Columnable a) => + T.Text -> DataFrame -> V.Vector Int -> Tree a -> Tree a -> [CarePoint] +identifyCarePoints = identifyCarePointsCached @a M.empty + +{- | Rows the parent must route to a specific child for the (fixed) subtrees to +classify correctly; a 'CondCache' avoids re-interpreting subtree conditions. +-} +identifyCarePointsCached :: + forall a. + (Columnable a) => + CondCache -> + T.Text -> + DataFrame -> + V.Vector Int -> + Tree a -> + Tree a -> + [CarePoint] +identifyCarePointsCached cache target df indices leftTree rightTree = + maybe [] carePoints (interpretLabelCol @a df target) + where + leftPreds = predictManyWithTreeCached cache leftTree df indices + rightPreds = predictManyWithTreeCached cache rightTree df indices + carePoints targetVals = V.toList (V.imapMaybe (checkPoint targetVals leftPreds rightPreds) indices) + +checkPoint :: + (Eq a) => + V.Vector a -> V.Vector a -> V.Vector a -> Int -> Int -> Maybe CarePoint +checkPoint targetVals leftPreds rightPreds k idx = + case (leftPreds V.! k == trueLabel, rightPreds V.! k == trueLabel) of + (True, False) -> Just (CarePoint idx GoLeft) + (False, True) -> Just (CarePoint idx GoRight) + _ -> Nothing + where + trueLabel = targetVals V.! idx + +-- | Care points a free condition misroutes (uncached; for the linear path). +countCarePointErrors :: Expr Bool -> DataFrame -> [CarePoint] -> Int +countCarePointErrors cond df carePoints = + maybe (length carePoints) (`countErrorsByVec` carePoints) (branchBool df cond) + +partitionIndices :: + Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int) +partitionIndices = partitionIndicesCached M.empty + +{- | 'partitionIndices' resolving the condition through a 'CondCache'; a miss +routes every index left (matching the uncached fallback). +-} +partitionIndicesCached :: + CondCache -> + Expr Bool -> + DataFrame -> + V.Vector Int -> + (V.Vector Int, V.Vector Int) +partitionIndicesCached cache cond df indices = case lookupCondVec cache df cond of + Nothing -> (indices, V.empty) + Just boolVals -> V.partition (boolVals VU.!) indices + +-- | A split is valid at a node when both children keep at least 'minLeafSize'. +isValidAtNode :: TreeConfig -> DataFrame -> V.Vector Int -> Expr Bool -> Bool +isValidAtNode cfg df indices c = + V.length t >= minLeafSize cfg && V.length f >= minLeafSize cfg + where + (t, f) = partitionIndices c df indices + +majorityValueFromIndices :: + forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> a +majorityValueFromIndices target df indices = majorityOf (countLabels (labelColOrThrow @a df target) indices) + +labelColOrThrow :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a +labelColOrThrow df target = case interpret @a df (Col target) of + Left e -> throw e + Right (TColumn column) -> either throw id (toVector @a column) + +countLabels :: (Ord a) => V.Vector a -> V.Vector Int -> M.Map a Int +countLabels vals = V.foldl' (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty + +majorityOf :: M.Map a Int -> a +majorityOf counts + | M.null counts = error "Empty indices in majorityValueFromIndices" + | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) + +computeTreeLoss :: + forall a. + (Columnable a) => T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double +computeTreeLoss = computeTreeLossCached @a M.empty + +-- | 0/1 loss of a tree over @indices@, with a 'CondCache' for the predictions. +computeTreeLossCached :: + forall a. + (Columnable a) => + CondCache -> T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double +computeTreeLossCached cache target df indices tree + | V.null indices = 0 + | otherwise = + maybe 1.0 (treeLoss cache tree df indices) (interpretLabelCol @a df target) + +treeLoss :: + (Columnable a) => + CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a -> Double +treeLoss cache tree df indices targetVals = + fromIntegral (countMismatches targetVals indices preds) + / fromIntegral (V.length indices) + where + preds = predictManyWithTreeCached cache tree df indices + +countMismatches :: (Eq a) => V.Vector a -> V.Vector Int -> V.Vector a -> Int +countMismatches targetVals indices preds = + V.length + (V.ifilter (\k _ -> targetVals V.! (indices V.! k) /= preds V.! k) preds) diff --git a/src-internal/DataFrame/DecisionTree/Prune.hs b/src-internal/DataFrame/DecisionTree/Prune.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Prune.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Post-convergence simplification of a fitted tree and its expression form: +drop branches forced by path-condition entailment, collapse identical +siblings, and fold redundant nested conditionals. +-} +module DataFrame.DecisionTree.Prune ( + pruneDead, + treeEq, + pruneExpr, +) where + +import DataFrame.DecisionTree.Types (Tree (..)) +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr (..), eqExpr) +import DataFrame.Internal.Simplify (PredFact, entails, factFalse, factTrue) + +{- | Drop branches whose test is forced by the path conditions reaching them, +and collapse @Branch c t t@ to @t@. Sound for the decidable threshold subset; +other tests are left untouched. +-} +pruneDead :: forall a. (Columnable a) => Tree a -> Tree a +pruneDead = go [] + where + go :: [PredFact] -> Tree a -> Tree a + go _ (Leaf v) = Leaf v + go facts (Branch cond left right) = case entails facts cond of + Just True -> go facts left + Just False -> go facts right + Nothing -> + reconcile + cond + (go (addFact (factTrue cond) facts) left) + (go (addFact (factFalse cond) facts) right) + +reconcile :: (Columnable a) => Expr Bool -> Tree a -> Tree a -> Tree a +reconcile cond left right + | treeEq left right = left + | otherwise = Branch cond left right + +addFact :: Maybe PredFact -> [PredFact] -> [PredFact] +addFact (Just f) fs = f : fs +addFact Nothing fs = fs + +treeEq :: (Columnable a) => Tree a -> Tree a -> Bool +treeEq (Leaf x) (Leaf y) = x == y +treeEq (Branch c1 l1 r1) (Branch c2 l2 r2) = eqExpr c1 c2 && treeEq l1 l2 && treeEq r1 r2 +treeEq _ _ = False + +{- | Recursively fold @If@ expressions whose branches coincide or nest the same +condition; leave other expressions structurally unchanged. +-} +pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a +pruneExpr (If cond t0 f0) = collapseIf cond (pruneExpr t0) (pruneExpr f0) +pruneExpr (Unary op e) = Unary op (pruneExpr e) +pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) +pruneExpr e = e + +collapseIf :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a +collapseIf cond t f + | eqExpr t f = t + | If ci ti _ <- t, eqExpr cond ci = If cond ti f + | If ci _ fi <- f, eqExpr cond ci = If cond t fi + | otherwise = If cond t f diff --git a/src-internal/DataFrame/DecisionTree/Tao.hs b/src-internal/DataFrame/DecisionTree/Tao.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Tao.hs @@ -0,0 +1,266 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Tree Alternating Optimization: hold the tree fixed and re-optimize one node +at a time, bottom-up, minimizing care-point misroutes. Sibling subtrees at a +depth level are independent and optimized in parallel. +-} +module DataFrame.DecisionTree.Tao ( + taoOptimize, + taoOptimizeCV, + taoIteration, + taoIterationCV, + optimizeNode, + findBestSplitTAO, +) where + +import DataFrame.DecisionTree.CondVec +import DataFrame.DecisionTree.Linear (bestLinearCandidate) +import DataFrame.DecisionTree.Pool ( + bestDiscreteCandidate, + candidateParChunk, + evalWithPenaltyVec, + ) +import DataFrame.DecisionTree.Predict +import DataFrame.DecisionTree.Prune (pruneDead) +import DataFrame.DecisionTree.Types +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.DataFrame (DataFrame) +import DataFrame.Internal.Expression (Expr) + +import Control.Parallel (par, pseq) +import Control.Parallel.Strategies (parListChunk, rdeepseq, using) +import Data.Function (on) +import Data.List (foldl', minimumBy) +import Data.Maybe (catMaybes, mapMaybe) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +{- | The constant per-fit context threaded through the node-optimization +recursion (the cache is rebuilt each iteration). +-} +data TaoEnv = TaoEnv + { teCache :: !CondCache + , teCfg :: !TreeConfig + , teTarget :: !T.Text + , teConds :: ![CondVec] + , teDf :: !DataFrame + } + +-- | Public TAO entry point over raw conditions; materializes each once. +taoOptimize :: + forall a. + (Columnable a, Ord a) => + TreeConfig -> + T.Text -> + [Expr Bool] -> + DataFrame -> + V.Vector Int -> + Tree a -> + Tree a +taoOptimize cfg target conds df = + taoOptimizeCV @a cfg target (mapMaybe (materializeCondVec df) conds) df + +{- | TAO outer loop over pre-evaluated candidates: iterate until the iteration +budget or convergence tolerance is reached, then prune dead branches. +-} +taoOptimizeCV :: + forall a. + (Columnable a, Ord a) => + TreeConfig -> + T.Text -> + [CondVec] -> + DataFrame -> + V.Vector Int -> + Tree a -> + Tree a +taoOptimizeCV cfg target condVecs df rootIndices initialTree = + go 0 initialTree (lossWith baseCache initialTree) + where + baseCache = condCacheFromVecs condVecs + lossWith cache = computeTreeLossCached @a cache target df rootIndices + go iter tree prevLoss + | iter >= taoIterations cfg = pruneDead tree + | prevLoss - newLoss < taoConvergenceTol cfg = pruneDead tree' + | otherwise = go (iter + 1) tree' newLoss + where + cache = addTreeCondsToCache df tree baseCache + tree' = taoIterationCV @a cache cfg target condVecs df rootIndices tree + newLoss = lossWith cache tree' + +-- | Public single-iteration entry point. +taoIteration :: + forall a. + (Columnable a, Ord a) => + TreeConfig -> + T.Text -> + [Expr Bool] -> + DataFrame -> + V.Vector Int -> + Tree a -> + Tree a +taoIteration cfg target conds df rootIndices tree = + let condVecs = mapMaybe (materializeCondVec df) conds + cache = addTreeCondsToCache df tree (condCacheFromVecs condVecs) + in taoIterationCV @a cache cfg target condVecs df rootIndices tree + +-- | One bottom-to-top sweep: re-optimize every node level by level. +taoIterationCV :: + forall a. + (Columnable a, Ord a) => + CondCache -> + TreeConfig -> + T.Text -> + [CondVec] -> + DataFrame -> + V.Vector Int -> + Tree a -> + Tree a +taoIterationCV cache cfg target condVecs df rootIndices tree = + foldl' + (optimizeDepthLevel env rootIndices) + tree + [treeDepth tree, treeDepth tree - 1 .. 0] + where + env = TaoEnv cache cfg target condVecs df + +optimizeDepthLevel :: + forall a. + (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Int -> Tree a +optimizeDepthLevel env rootIndices tree = optimizeAtDepth @a env rootIndices tree 0 + +optimizeAtDepth :: + forall a. + (Columnable a, Ord a) => + TaoEnv -> V.Vector Int -> Tree a -> Int -> Int -> Tree a +optimizeAtDepth env indices tree currentDepth targetDepth + | currentDepth == targetDepth = optimizeNode @a env indices tree + | otherwise = case tree of + Leaf v -> Leaf v + Branch cond left right -> optimizeChildren @a env indices cond left right currentDepth targetDepth + +{- | Optimize the two subtrees over their disjoint index sets, scoring the left +in parallel with the right (the cache is read-only, so this is pure). +-} +optimizeChildren :: + forall a. + (Columnable a, Ord a) => + TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Int -> Int -> Tree a +optimizeChildren env indices cond left right currentDepth targetDepth = + forceTreeWork left' `par` (forceTreeWork right' `pseq` Branch cond left' right') + where + (indicesL, indicesR) = partitionIndicesCached (teCache env) cond (teDf env) indices + left' = optimizeAtDepth @a env indicesL left (currentDepth + 1) targetDepth + right' = optimizeAtDepth @a env indicesR right (currentDepth + 1) targetDepth + +{- | Force a subtree's optimization work to WHNF so the parallel scheduler has +something substantial to evaluate; pure and value-preserving. +-} +forceTreeWork :: Tree a -> () +forceTreeWork (Leaf v) = v `seq` () +forceTreeWork (Branch c l r) = c `seq` forceTreeWork l `seq` forceTreeWork r + +{- | Re-optimize one node: pick its best split, or collapse to a leaf when the +node is empty or the chosen split underflows 'minLeafSize'. +-} +optimizeNode :: + forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Tree a +optimizeNode env indices tree + | V.null indices = tree + | otherwise = case tree of + Leaf _ -> leaf + Branch oldCond left right -> rebuiltBranch env indices oldCond left right leaf + where + leaf = Leaf (majorityValueFromIndices @a (teTarget env) (teDf env) indices) + +rebuiltBranch :: + forall a. + (Columnable a, Ord a) => + TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Tree a -> Tree a +rebuiltBranch env indices oldCond left right leaf + | underflows = leaf + | otherwise = Branch newCond left right + where + newCond = findBestSplitTAO @a env indices left right oldCond + (l, r) = partitionIndicesCached (teCache env) newCond (teDf env) indices + underflows = V.length l < minLeafSize (teCfg env) || V.length r < minLeafSize (teCfg env) + +{- | The lowest-penalty replacement condition for a node, falling back to the +current condition when no valid candidate beats it. +-} +findBestSplitTAO :: + forall a. + (Columnable a) => + TaoEnv -> V.Vector Int -> Tree a -> Tree a -> Expr Bool -> Expr Bool +findBestSplitTAO env indices leftTree rightTree currentCond + | V.null indices || null carePoints = currentCond + | pureReplacementLinear cfg + , Just c <- linearCandidate + , isValidAtNode cfg (teDf env) indices c = + c + | otherwise = bestOfPool penaltyCV currentCond pool + where + cfg = teCfg env + carePoints = + identifyCarePointsCached @a + (teCache env) + (teTarget env) + (teDf env) + indices + leftTree + rightTree + penaltyCV = evalWithPenaltyVec cfg carePoints + linearCandidate = bestLinearCandidate cfg (teDf env) carePoints + valid = filterValidCandidates cfg indices (teConds env) + pool = + candidatePool + env + indices + currentCond + (bestDiscreteCandidate cfg penaltyCV valid) + linearCandidate + +bestOfPool :: (CondVec -> (Int, Int)) -> Expr Bool -> [CondVec] -> Expr Bool +bestOfPool _ currentCond [] = currentCond +bestOfPool penaltyCV _ pool = cvExpr (minimumBy (compare `on` penaltyCV) pool) + +{- | Validity-filtered candidates the node could split on: both children must +keep at least 'minLeafSize'. Scored in parallel chunks, order preserved. +-} +filterValidCandidates :: TreeConfig -> V.Vector Int -> [CondVec] -> [CondVec] +filterValidCandidates cfg indices condVecs = map snd (filter fst (zip validity condVecs)) + where + validity = + map (validAtNode cfg indices) condVecs + `using` parListChunk candidateParChunk rdeepseq + +validAtNode :: TreeConfig -> V.Vector Int -> CondVec -> Bool +validAtNode cfg indices cv = nTrue >= minLeaf && (V.length indices - nTrue) >= minLeaf + where + minLeaf = minLeafSize cfg + nTrue = + V.foldl' + (\ !acc i -> if cvVec cv VU.! i then acc + 1 else acc) + (0 :: Int) + indices + +{- | The candidate pool to minimize over: the current condition, the best +discrete candidate, and the linear candidate, each kept only if valid. +-} +candidatePool :: + TaoEnv -> + V.Vector Int -> + Expr Bool -> + Maybe CondVec -> + Maybe (Expr Bool) -> + [CondVec] +candidatePool env indices currentCond discreteCV linearCandidate = + filter + (isValidAtNode (teCfg env) (teDf env) indices . cvExpr) + (catMaybes [currentCV, discreteCV, linearCV]) + where + currentCV = CondVec currentCond <$> lookupCondVec (teCache env) (teDf env) currentCond + linearCV = linearCandidate >>= materializeCondVec (teDf env) diff --git a/src-internal/DataFrame/DecisionTree/Types.hs b/src-internal/DataFrame/DecisionTree/Types.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/DecisionTree/Types.hs @@ -0,0 +1,190 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Shared types, configuration and ordering machinery for the decision-tree +learner. Imported by every other @DataFrame.DecisionTree.*@ module. +-} +module DataFrame.DecisionTree.Types ( + Tree (..), + treeDepth, + TreeConfig (..), + SynthConfig (..), + defaultTreeConfig, + defaultSynthConfig, + ColumnOrdering (..), + orderable, + defaultColumnOrdering, + withOrdFrom, + CarePoint (..), + Direction (..), + ttrace, +) where + +import DataFrame.Internal.Column (Columnable) +import DataFrame.Internal.Expression (Expr (..)) +import qualified DataFrame.LinearSolver as LS + +import Data.Int (Int16, Int32, Int64, Int8) +import qualified Data.Map.Strict as M +import Data.Proxy (Proxy (..)) +import qualified Data.Text as T +import Data.Type.Equality (testEquality, (:~:) (..)) +import Data.Word (Word16, Word32, Word64, Word8) +import qualified Debug.Trace as Trace +import System.Environment (lookupEnv) +import System.IO.Unsafe (unsafePerformIO) +import Type.Reflection (SomeTypeRep (..), typeRep) + +{- | A fitted tree: a leaf value, or an internal node testing a boolean +expression with @True@ routing left. +-} +data Tree a + = Leaf !a + | Branch !(Expr Bool) !(Tree a) !(Tree a) + deriving (Show) + +treeDepth :: Tree a -> Int +treeDepth (Leaf _) = 0 +treeDepth (Branch _ l r) = 1 + max (treeDepth l) (treeDepth r) + +{- | A row the parent node must route to a specific child for the subtrees to +classify it correctly (the TAO objective is the count of misroutes). +-} +data CarePoint = CarePoint + { cpIndex :: !Int + , cpCorrectDir :: !Direction + } + deriving (Eq, Show) + +data Direction = GoLeft | GoRight + deriving (Eq, Show) + +data TreeConfig = TreeConfig + { maxTreeDepth :: Int + , minSamplesSplit :: Int + , minLeafSize :: Int + , percentiles :: [Int] + , expressionPairs :: Int + , synthConfig :: SynthConfig + , taoIterations :: Int + , taoConvergenceTol :: Double + , columnOrdering :: ColumnOrdering + , useLinearSolver :: Bool + , linearSolverConfig :: LS.SolverConfig + , minCarePointsForLinear :: Int + , pureReplacementLinear :: Bool + } + +data SynthConfig = SynthConfig + { maxExprDepth :: Int + , boolExpansion :: Int + , disallowedCombinations :: [(T.Text, T.Text)] + , complexityPenalty :: Double + , enableStringOps :: Bool + , enableCrossCols :: Bool + , enableArithOps :: Bool + , maxCategoricalSubsetCardinality :: Int + , perColumnQuota :: Maybe Int + } + deriving (Eq, Show) + +defaultSynthConfig :: SynthConfig +defaultSynthConfig = + SynthConfig + { maxExprDepth = 2 + , boolExpansion = 2 + , disallowedCombinations = [] + , complexityPenalty = 0.05 + , enableStringOps = True + , enableCrossCols = True + , enableArithOps = True + , maxCategoricalSubsetCardinality = 4 + , perColumnQuota = Just 3 + } + +defaultTreeConfig :: TreeConfig +defaultTreeConfig = + TreeConfig + { maxTreeDepth = 4 + , minSamplesSplit = 5 + , minLeafSize = 1 + , percentiles = [0, 10 .. 100] + , expressionPairs = 10 + , synthConfig = defaultSynthConfig + , taoIterations = 10 + , taoConvergenceTol = 1e-6 + , columnOrdering = defaultColumnOrdering + , useLinearSolver = True + , linearSolverConfig = LS.defaultSolverConfig + , minCarePointsForLinear = 10 + , pureReplacementLinear = False + } + +{- | Which column types support ordering for splits. Register a type with +'orderable' and combine with @<>@. +-} +newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict) + +instance Semigroup ColumnOrdering where + ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b) + +instance Monoid ColumnOrdering where + mempty = ColumnOrdering M.empty + +-- | Register a type as orderable for decision-tree splits. +orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering +orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a))) + +-- | All standard numeric, text, and primitive types. +defaultColumnOrdering :: ColumnOrdering +defaultColumnOrdering = mconcat (numericOrderings ++ otherOrderings) + +numericOrderings :: [ColumnOrdering] +numericOrderings = + [ orderable @Int + , orderable @Int8 + , orderable @Int16 + , orderable @Int32 + , orderable @Int64 + , orderable @Word + , orderable @Word8 + , orderable @Word16 + , orderable @Word32 + , orderable @Word64 + , orderable @Integer + , orderable @Double + , orderable @Float + ] + +otherOrderings :: [ColumnOrdering] +otherOrderings = + [orderable @Bool, orderable @Char, orderable @T.Text, orderable @String] + +-- | Existential @Ord@ dictionary keyed by type representation. +data OrdDict where + OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict + +{- | Run @k@ with the @Ord a@ instance recovered from the ordering registry, +or 'Nothing' when @a@ is not registered. +-} +withOrdFrom :: + forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r +withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of + Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of + Just Refl -> Just k + Nothing -> Nothing + Nothing -> Nothing + +{-# NOINLINE taoTraceEnabled #-} +taoTraceEnabled :: Bool +taoTraceEnabled = unsafePerformIO (fmap (== Just "1") (lookupEnv "TAO_TRACE")) + +-- | Emit a trace line when @TAO_TRACE=1@; a no-op otherwise. +ttrace :: String -> a -> a +ttrace msg x + | taoTraceEnabled = Trace.trace ("[TAO] " ++ msg) x + | otherwise = x diff --git a/src-internal/DataFrame/Featurize/Internal.hs b/src-internal/DataFrame/Featurize/Internal.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Featurize/Internal.hs @@ -0,0 +1,206 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Shared internal helpers used across the model fitters: turning a 'DataFrame' +plus a target/feature 'Expr' into the numeric matrices the algorithms consume, +plus common expression builders (affine score, arg-max/arg-min over scores). +-} +module DataFrame.Featurize.Internal ( + -- * Supervised extraction + featureNames, + numericMatrix, + targetDoubles, + targetValues, + + -- * Unsupervised extraction + Features (..), + extractFeatures, + columnExprName, + materializeColumn, + + -- * Expression builders + affineExpr, + argMaxExpr, + argMinExpr, +) where + +import Control.Exception (throw) +import Data.Maybe (isJust) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Type.Reflection (TypeRep, typeRep) + +import Data.Either (fromRight) +import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (Columnable, columnBitmap, columnTypeString) +import qualified DataFrame.Internal.Column as D +import DataFrame.Internal.DataFrame (DataFrame, columnNames, getColumn) +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.LinearAlgebra (Matrix, transposeM) +import DataFrame.Operations.Core ( + columnAsDoubleVector, + columnAsUnboxedVector, + columnAsVector, + ) +import DataFrame.Operators ((.&&.), (.*.), (.+.), (.<=.), (.>=.)) + +-- | Every column name except the supervised target's. +featureNames :: Expr a -> DataFrame -> [T.Text] +featureNames (Col target) df = filter (/= target) (columnNames df) +featureNames _ df = columnNames df + +{- | The named columns as a row-major @n×d@ matrix, paired with the names. Every +column must already be stored as non-null 'Double'; a nullable ('Maybe Double') +or non-'Double' column is a fit-time error rather than a silent coercion. +-} +numericMatrix :: [T.Text] -> DataFrame -> (V.Vector T.Text, Matrix) +numericMatrix names df = (V.fromList names, transposeM colMajor) + where + colMajor = V.fromList (map (`requireDoubleColumn` df) names) + +{- | Read a feature\/target column strictly as a non-null 'Double' vector. A +nullable ('Maybe Double') column, or a column of any other type, is a fit-time +error naming the column and pointing at the fix. +-} +requireDoubleColumn :: T.Text -> DataFrame -> VU.Vector Double +requireDoubleColumn name df = case getColumn name df of + Nothing -> + throw + ( TypeMismatchException + ( MkTypeErrorContext + (Right (typeRep @Double)) + (Left "missing" :: Either String (TypeRep Double)) + (Just (T.unpack name)) + Nothing + ) + ) + Just c -> + if isJust (columnBitmap c) + then throw (nullableInputError name (columnTypeString c)) + else + if D.hasElemType @Double c + then fromRight undefined (columnAsUnboxedVector (F.col @Double name) df) + else + throw + ( asDoubleInputError + name + ( TypeMismatchException + ( MkTypeErrorContext + (Right (typeRep @Double)) + (Left (columnTypeString c) :: Either String (TypeRep Double)) + (Just (T.unpack name)) + Nothing + ) + ) + ) + +-- | Actionable error for a non-'Double' input column (wraps the type mismatch). +asDoubleInputError :: T.Text -> DataFrameException -> DataFrameException +asDoubleInputError name (TypeMismatchException ctx) = + TypeMismatchException + ctx + { errorColumnName = Just (T.unpack name) + , callingFunctionName = + Just + "model fit (input columns must be Double; convert with toDouble or build the column as Double)" + } +asDoubleInputError _ e = e + +-- | Actionable error for a nullable ('Maybe Double') input column. +nullableInputError :: T.Text -> String -> DataFrameException +nullableInputError name found = + TypeMismatchException + ( MkTypeErrorContext + { userType = Right (typeRep @Double) + , expectedType = Left found + , errorColumnName = Just (T.unpack name) + , callingFunctionName = + Just + "model fit (input columns must be non-null Double; drop or impute missing values before fitting)" + } :: + TypeErrorContext Double () + ) + +{- | The target column as a non-null 'Double' vector. A nullable or non-'Double' +target is a fit-time error, not a coercion. +-} +targetDoubles :: Expr Double -> DataFrame -> VU.Vector Double +targetDoubles (Col name) df = requireDoubleColumn name df +targetDoubles expr df = case columnAsUnboxedVector expr df of + Right v -> v + Left e -> throw (asDoubleInputError (T.pack "") e) + +-- | The target column as a vector of its own type (for classifiers). +targetValues :: (Columnable a) => Expr a -> DataFrame -> V.Vector a +targetValues expr df = case columnAsVector expr df of + Right v -> v + Left e -> throw e + +{- | The extracted feature columns of an unsupervised fit, in the shapes the +algorithms need: names, column-major vectors, the row-major matrix, and the +@(n, d)@ dimensions. +-} +data Features = Features + { ftNames :: ![T.Text] + , ftCols :: ![VU.Vector Double] + , ftRows :: !Matrix + , ftN :: !Int + , ftD :: !Int + } + +-- | Extract the given feature columns once, in every shape the fitters use. +extractFeatures :: [Expr Double] -> DataFrame -> Features +extractFeatures features df = Features names cols rows n d + where + names = map columnExprName features + cols = map (materializeColumn df) features + n = case cols of + (x : _) -> VU.length x + _ -> 0 + d = length cols + rows = V.generate n (\i -> VU.generate d (\j -> (cols !! j) VU.! i)) + +-- | The column name behind a @Col@ feature expression. +columnExprName :: Expr Double -> T.Text +columnExprName (Col n) = n +columnExprName e = error ("expected a column expression, got " ++ show e) + +-- | Interpret a @Col@ (or numeric) expression to a @Double@ vector. +materializeColumn :: DataFrame -> Expr Double -> VU.Vector Double +materializeColumn df e = case columnAsDoubleVector e df of + Right v -> v + Left err -> throw err + +{- | An affine score @b + Σ wⱼ·colⱼ@ over named columns, dropping zero weights +(the shared core of linear/logistic/SVM margins). +-} +affineExpr :: Double -> [(Double, T.Text)] -> Expr Double +affineExpr b terms = + foldr + (.+.) + (F.lit b) + [F.lit w .*. (Col n :: Expr Double) | (w, n) <- terms, w /= 0] + +{- | The class whose score is greatest, as a nested-@If@ expression; ties go to +the earlier class. +-} +argMaxExpr :: (Columnable a) => [(a, Expr Double)] -> Expr a +argMaxExpr = argExtreme (.>=.) + +-- | The class whose score is smallest (e.g. nearest cluster by distance). +argMinExpr :: (Columnable a) => [(a, Expr Double)] -> Expr a +argMinExpr = argExtreme (.<=.) + +argExtreme :: + (Columnable a) => + (Expr Double -> Expr Double -> Expr Bool) -> [(a, Expr Double)] -> Expr a +argExtreme _ [] = error "argExtreme: no classes" +argExtreme _ [(c, _)] = Lit c +argExtreme cmp ((c, sc) : rest) = + If + (foldr ((\o acc -> cmp sc o .&&. acc) . snd) (F.lit True) rest) + (Lit c) + (argExtreme cmp rest) diff --git a/src-internal/DataFrame/LinearAlgebra.hs b/src-internal/DataFrame/LinearAlgebra.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/LinearAlgebra.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Dependency-free dense linear algebra over row-major matrices, shared by the +models in @dataframe-learn@. Solvers live in "DataFrame.LinearAlgebra.Solve" +and eigenproblems in "DataFrame.LinearAlgebra.Eigen". +-} +module DataFrame.LinearAlgebra ( + Matrix, + dot, + axpy, + scaleV, + matVec, + tMatVec, + gram, + transposeM, + identityM, + logSumExp, + sqDist, + nearestCenter, + epsNeighbors, +) where + +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +{- | Row-major dense matrix: an outer boxed vector of equal-length rows. An +@n×d@ matrix has @n@ rows of length @d@. +-} +type Matrix = V.Vector (VU.Vector Double) + +-- | Inner product of two equal-length vectors. +dot :: VU.Vector Double -> VU.Vector Double -> Double +dot a b = VU.sum (VU.zipWith (*) a b) +{-# INLINE dot #-} + +-- | @axpy a x y = a*x + y@. +axpy :: Double -> VU.Vector Double -> VU.Vector Double -> VU.Vector Double +axpy a = VU.zipWith (\xi yi -> a * xi + yi) +{-# INLINE axpy #-} + +-- | Scalar-vector product. +scaleV :: Double -> VU.Vector Double -> VU.Vector Double +scaleV a = VU.map (* a) +{-# INLINE scaleV #-} + +-- | @matVec A v@ for @A@ of shape @n×d@ and @v@ of length @d@; result length @n@. +matVec :: Matrix -> VU.Vector Double -> VU.Vector Double +matVec a v = VU.convert (V.map (`dot` v) a) + +-- | @tMatVec A v = Aᵀ v@ for @A@ of shape @n×d@, @v@ of length @n@; result length @d@. +tMatVec :: Matrix -> VU.Vector Double -> VU.Vector Double +tMatVec a v + | V.null a = VU.empty + | otherwise = V.foldl' step (VU.replicate d 0) (V.zipWith (,) vBoxed a) + where + d = VU.length (V.head a) + vBoxed = V.generate (V.length a) (v VU.!) + step !acc (vi, row) = axpy vi row acc + +-- | @gram A = Aᵀ A@, the @d×d@ symmetric matrix of column inner products. +gram :: Matrix -> Matrix +gram a + | V.null a = V.empty + | otherwise = + V.generate d $ \i -> + VU.generate d $ \j -> + V.foldl' (\ !acc row -> acc + (row VU.! i) * (row VU.! j)) 0 a + where + d = VU.length (V.head a) + +-- | Transpose an @n×d@ matrix to @d×n@. +transposeM :: Matrix -> Matrix +transposeM a + | V.null a = V.empty + | otherwise = V.generate d $ \j -> VU.generate n $ \i -> (a V.! i) VU.! j + where + n = V.length a + d = VU.length (V.head a) + +-- | @d×d@ identity matrix. +identityM :: Int -> Matrix +identityM d = V.generate d $ \i -> VU.generate d $ \j -> if i == j then 1 else 0 + +-- | Numerically stable @log Σ exp xᵢ@. +logSumExp :: VU.Vector Double -> Double +logSumExp xs + | VU.null xs = negate (1 / 0) + | otherwise = m + log (VU.sum (VU.map (\x -> exp (x - m)) xs)) + where + m = VU.maximum xs + +-- | Squared Euclidean distance. +sqDist :: VU.Vector Double -> VU.Vector Double -> Double +sqDist a b = VU.sum (VU.zipWith (\x y -> let z = x - y in z * z) a b) +{-# INLINE sqDist #-} + +-- | Index of and squared distance to the nearest centre. +nearestCenter :: + V.Vector (VU.Vector Double) -> VU.Vector Double -> (Int, Double) +nearestCenter centers p = + V.ifoldl' + ( \(!bi, !bd) i c -> + let dd = sqDist c p in if dd < bd then (i, dd) else (bi, bd) + ) + (-1, 1 / 0) + centers + +{- | Indices @j@ (excluding @i@) within squared radius @eps²@ of row @i@, by +brute force; @O(n d)@ per query. +-} +epsNeighbors :: Double -> Matrix -> Int -> VU.Vector Int +epsNeighbors eps rows i = + VU.fromList + [ j + | j <- [0 .. n - 1] + , j /= i + , sqDist (rows V.! i) (rows V.! j) <= eps2 + ] + where + n = V.length rows + eps2 = eps * eps diff --git a/src-internal/DataFrame/LinearAlgebra/Eigen.hs b/src-internal/DataFrame/LinearAlgebra/Eigen.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/LinearAlgebra/Eigen.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Symmetric eigenproblems in pure Haskell: cyclic Jacobi for full +decomposition (PCA covariance, @m×m@ kernels) and power iteration for the +dominant eigenpair (FISTA step sizes). Deterministic, sign-canonicalised output. +-} +module DataFrame.LinearAlgebra.Eigen ( + jacobiEigenSym, + powerIterTop, +) where + +import Control.Monad (forM_, when) +import Control.Monad.ST (runST) +import Data.List (sortBy) +import Data.Ord (Down (..), comparing) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import DataFrame.LinearAlgebra (Matrix, dot, matVec, scaleV) + +{- | Cyclic Jacobi eigendecomposition of a symmetric matrix. Eigenvalues are +returned in descending order paired with eigenvectors as rows, each +sign-canonicalised (largest-magnitude component positive) for unique output. +-} +jacobiEigenSym :: Matrix -> (VU.Vector Double, Matrix) +jacobiEigenSym a0 + | V.null a0 = (VU.empty, V.empty) + | otherwise = runST $ do + a <- VUM.new (d * d) + forM_ [0 .. d - 1] $ \i -> + forM_ [0 .. d - 1] $ \j -> + VUM.write a (i * d + j) ((a0 V.! i) VU.! j) + v <- VUM.replicate (d * d) 0 + forM_ [0 .. d - 1] $ \i -> VUM.write v (i * d + i) 1 + sweep a v 0 + afrozen <- VU.freeze a + vmat <- VU.freeze v + let diag = VU.generate d (\i -> afrozen VU.! (i * d + i)) + vecs = + V.generate d $ \col -> + VU.generate d $ \row -> vmat VU.! (row * d + col) + paired = + sortBy + (comparing (Down . fst)) + (zip (VU.toList diag) (V.toList vecs)) + pure + ( VU.fromList (map fst paired) + , V.fromList (map (canonicalSign . snd) paired) + ) + where + d = V.length a0 + maxSweeps = 100 + tol = 1e-12 + sweep a v s + | s >= maxSweeps = pure () + | otherwise = do + off <- offNorm a + when (off >= tol) $ do + forM_ [0 .. d - 2] $ \p -> + forM_ [p + 1 .. d - 1] $ \q -> rotate a v p q + sweep a v (s + 1) + offNorm a = go 0 0 + where + go i !acc + | i >= d = pure acc + | otherwise = do + r <- goRow i (i + 1) acc + go (i + 1) r + goRow i j !acc + | j >= d = pure acc + | otherwise = do + x <- VUM.read a (i * d + j) + goRow i (j + 1) (acc + x * x) + rotate a v p q = do + apq <- VUM.read a (p * d + q) + when (abs apq > 1e-300) $ do + app <- VUM.read a (p * d + p) + aqq <- VUM.read a (q * d + q) + let theta = (aqq - app) / (2 * apq) + s' = if theta == 0 then 1 else signum theta + t = s' / (abs theta + sqrt (theta * theta + 1)) + c = 1 / sqrt (t * t + 1) + sn = t * c + forM_ [0 .. d - 1] $ \i -> do + aip <- VUM.read a (i * d + p) + aiq <- VUM.read a (i * d + q) + VUM.write a (i * d + p) (c * aip - sn * aiq) + VUM.write a (i * d + q) (sn * aip + c * aiq) + forM_ [0 .. d - 1] $ \j -> do + apj <- VUM.read a (p * d + j) + aqj <- VUM.read a (q * d + j) + VUM.write a (p * d + j) (c * apj - sn * aqj) + VUM.write a (q * d + j) (sn * apj + c * aqj) + forM_ [0 .. d - 1] $ \i -> do + vip <- VUM.read v (i * d + p) + viq <- VUM.read v (i * d + q) + VUM.write v (i * d + p) (c * vip - sn * viq) + VUM.write v (i * d + q) (sn * vip + c * viq) + +canonicalSign :: VU.Vector Double -> VU.Vector Double +canonicalSign vec = + let idx = VU.maxIndex (VU.map abs vec) + in if vec VU.! idx < 0 then VU.map negate vec else vec + +{- | Dominant eigenvalue and eigenvector of a symmetric PSD matrix via power +iteration with a deterministic all-ones start. +-} +powerIterTop :: Int -> Matrix -> (Double, VU.Vector Double) +powerIterTop iters a + | V.null a = (0, VU.empty) + | otherwise = go iters (normalize (VU.replicate d 1)) + where + d = V.length a + normalize v = + let nrm = sqrt (dot v v) in if nrm == 0 then v else scaleV (1 / nrm) v + go 0 v = (dot v (matVec a v), v) + go k v = + let av = matVec a v + nrm = sqrt (dot av av) + in if nrm < 1e-300 then (0, v) else go (k - 1) (scaleV (1 / nrm) av) diff --git a/src-internal/DataFrame/LinearAlgebra/Solve.hs b/src-internal/DataFrame/LinearAlgebra/Solve.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/LinearAlgebra/Solve.hs @@ -0,0 +1,213 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Householder QR (for ordinary least squares) and Cholesky factorisation (for +ridge normal equations and Gaussian log-densities). Pure, deterministic, no +LAPACK; sound at the @d@ ≤ low-hundreds scales this library targets. +-} +module DataFrame.LinearAlgebra.Solve ( + qrLeastSquares, + cholesky, + choleskySolve, + logDetFromChol, + forwardSubst, + backSubst, +) where + +import Control.Monad (forM_) +import Control.Monad.ST (ST, runST) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import DataFrame.LinearAlgebra (Matrix) + +{- | Solve @min ‖A x − b‖₂@ for an @n×d@ matrix @A@ (@n ≥ d@) via Householder QR. +@Left cols@ reports rank deficiency (near-zero @R@ diagonal) with the offending +column indices; @Right x@ is the least-squares solution. +-} +qrLeastSquares :: Matrix -> VU.Vector Double -> Either [Int] (VU.Vector Double) +qrLeastSquares a b + | V.null a = Right VU.empty + | n < d = Left [0 .. d - 1] + | otherwise = runST $ do + mat <- VUM.new (n * d) + forM_ [0 .. n - 1] $ \i -> + forM_ [0 .. d - 1] $ \j -> + VUM.write mat (j * n + i) ((a V.! i) VU.! j) + rhs <- VUM.new n + forM_ [0 .. n - 1] $ \i -> VUM.write rhs i (b VU.! i) + deficient <- householder mat rhs n d + if not (null deficient) + then pure (Left deficient) + else do + x <- VUM.new d + backSubstQR mat rhs n d x + Right <$> VU.freeze x + where + n = V.length a + d = VU.length (V.head a) + +householder :: + VUM.STVector s Double -> VUM.STVector s Double -> Int -> Int -> ST s [Int] +householder mat rhs n d = go 0 [] + where + tol = 1e-10 + go k acc + | k >= d = pure (reverse acc) + | otherwise = do + normSq <- sumSq k + let alphaMag = sqrt normSq + akk <- VUM.read mat (k * n + k) + let alpha = if akk > 0 then negate alphaMag else alphaMag + if alphaMag < tol + then go (k + 1) (k : acc) + else do + VUM.write mat (k * n + k) (akk - alpha) + vNormSq <- sumSq k + if vNormSq < tol * tol + then go (k + 1) acc + else do + forM_ [k + 1 .. d - 1] $ \j -> reflectColumn k j vNormSq + reflectRhs k vNormSq + VUM.write mat (k * n + k) alpha + go (k + 1) acc + sumSq k = foldRows k + where + foldRows i + | i >= n = pure 0 + | otherwise = do + x <- VUM.read mat (k * n + i) + rest <- foldRows (i + 1) + pure (x * x + rest) + reflectColumn k j vNormSq = do + dotv <- dotV k j k + let beta = 2 * dotv / vNormSq + forM_ [k .. n - 1] $ \i -> do + vi <- VUM.read mat (k * n + i) + aij <- VUM.read mat (j * n + i) + VUM.write mat (j * n + i) (aij - beta * vi) + reflectRhs k vNormSq = do + dotv <- dotRhs k k + let beta = 2 * dotv / vNormSq + forM_ [k .. n - 1] $ \i -> do + vi <- VUM.read mat (k * n + i) + bi <- VUM.read rhs i + VUM.write rhs i (bi - beta * vi) + dotV k j i + | i >= n = pure 0 + | otherwise = do + vi <- VUM.read mat (k * n + i) + aij <- VUM.read mat (j * n + i) + rest <- dotV k j (i + 1) + pure (vi * aij + rest) + dotRhs k i + | i >= n = pure 0 + | otherwise = do + vi <- VUM.read mat (k * n + i) + bi <- VUM.read rhs i + rest <- dotRhs k (i + 1) + pure (vi * bi + rest) + +backSubstQR :: + VUM.STVector s Double -> + VUM.STVector s Double -> + Int -> + Int -> + VUM.STVector s Double -> + ST s () +backSubstQR mat rhs n d x = forM_ [d - 1, d - 2 .. 0] $ \i -> do + bi <- VUM.read rhs i + s <- sumAbove i (i + 1) 0 + rii <- VUM.read mat (i * n + i) + VUM.write x i ((bi - s) / rii) + where + sumAbove i j !acc + | j >= d = pure acc + | otherwise = do + rij <- VUM.read mat (j * n + i) + xj <- VUM.read x j + sumAbove i (j + 1) (acc + rij * xj) + +{- | Cholesky factor @L@ (lower-triangular, @A = L Lᵀ@) of a symmetric +positive-definite matrix, or 'Nothing' if a non-positive pivot is hit. +-} +cholesky :: Matrix -> Maybe Matrix +cholesky a + | V.null a = Just V.empty + | otherwise = runST $ do + l <- VUM.replicate (d * d) 0 + ok <- buildL l + if ok then Just <$> freezeLower l else pure Nothing + where + d = V.length a + buildL l = go 0 + where + go j + | j >= d = pure True + | otherwise = do + s <- sumLk l j j (j - 1) 0 + let ajj = (a V.! j) VU.! j + diag = ajj - s + if diag <= 0 + then pure False + else do + let ljj = sqrt diag + VUM.write l (j * d + j) ljj + forM_ [j + 1 .. d - 1] $ \i -> do + sij <- sumLk l i j (j - 1) 0 + let aij = (a V.! i) VU.! j + VUM.write l (i * d + j) ((aij - sij) / ljj) + go (j + 1) + sumLk l i j k !acc + | k < 0 = pure acc + | otherwise = do + lik <- VUM.read l (i * d + k) + ljk <- VUM.read l (j * d + k) + sumLk l i j (k - 1) (acc + lik * ljk) + freezeLower l = do + frozen <- VU.freeze l + pure $ V.generate d $ \i -> VU.slice (i * d) d frozen + +-- | Solve @L y = b@ for lower-triangular @L@. +forwardSubst :: Matrix -> VU.Vector Double -> VU.Vector Double +forwardSubst l b = runST $ do + y <- VUM.new d + forM_ [0 .. d - 1] $ \i -> do + let row = l V.! i + s <- sumKnown y row i 0 0 + VUM.write y i ((b VU.! i - s) / (row VU.! i)) + VU.freeze y + where + d = V.length l + sumKnown y row i j !acc + | j >= i = pure acc + | otherwise = do + yj <- VUM.read y j + sumKnown y row i (j + 1) (acc + (row VU.! j) * yj) + +-- | Solve @Lᵀ x = y@ for lower-triangular @L@. +backSubst :: Matrix -> VU.Vector Double -> VU.Vector Double +backSubst l y = runST $ do + x <- VUM.new d + forM_ [d - 1, d - 2 .. 0] $ \i -> do + s <- sumKnown x i (i + 1) 0 + VUM.write x i ((y VU.! i - s) / ((l V.! i) VU.! i)) + VU.freeze x + where + d = V.length l + sumKnown x i j !acc + | j >= d = pure acc + | otherwise = do + xj <- VUM.read x j + sumKnown x i (j + 1) (acc + ((l V.! j) VU.! i) * xj) + +{- | Solve the SPD system @A x = b@ via Cholesky; 'Nothing' when @A@ is not +positive-definite. +-} +choleskySolve :: Matrix -> VU.Vector Double -> Maybe (VU.Vector Double) +choleskySolve a b = do + l <- cholesky a + pure (backSubst l (forwardSubst l b)) + +-- | @log det A = 2 Σ log Lᵢᵢ@ from a Cholesky factor @L@. +logDetFromChol :: Matrix -> Double +logDetFromChol l = 2 * V.sum (V.imap (\i row -> log (row VU.! i)) l) diff --git a/src-internal/DataFrame/LinearSolver.hs b/src-internal/DataFrame/LinearSolver.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/LinearSolver.hs @@ -0,0 +1,450 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Proximal-gradient (FISTA) solver for L1/L2-regularized generalized linear +models. 'fitL1Logistic' is the binary logistic split solver; 'fitProx' +generalizes it to any 'SmoothLoss'. Features are standardized internally. +-} +module DataFrame.LinearSolver ( + -- * Model + LinearModel (..), + + -- * Configuration + SolverConfig (..), + defaultSolverConfig, + + -- * Solvers + fitL1Logistic, + fitProx, + + -- * Expr conversion + modelToExpr, + + -- * Internals (exposed for testing) + standardize, + columnStats, + softThreshold, + sigmoid, + dotProduct, +) where + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.LinearAlgebra (Matrix, gram, scaleV) +import DataFrame.LinearAlgebra.Eigen (powerIterTop) +import DataFrame.LinearSolver.Loss ( + SmoothLoss (..), + logisticLoss, + sigmoid, + ) +import DataFrame.Operators ((.*.), (.+.), (.>.)) + +import Control.Monad.ST (ST, runST) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +{- | A fitted linear classifier: predicts the positive class when +@sum (weights .* features) + intercept > 0@. Weights of exactly @0@ mark +features dropped by the L1 penalty (filtered out by 'modelToExpr'). +-} +data LinearModel = LinearModel + { lmWeights :: !(VU.Vector Double) + , lmIntercept :: !Double + , lmFeatureNames :: !(V.Vector T.Text) + } + deriving (Eq, Show) + +-- | Hyper-parameters for the FISTA solver. +data SolverConfig = SolverConfig + { scL1Lambda :: !Double + -- ^ Strength of the L1 penalty on weights (intercept is not regularized). + , scL2Lambda :: !Double + {- ^ Strength of the L2 penalty @(λ₂/2)·|w|²@ (Elastic Net; Zou & Hastie + 2005). Combined with @scL1Lambda@ this is the elastic-net objective; + @0@ reduces the solver to pure L1. + -} + , scMaxIter :: !Int + -- ^ Maximum number of FISTA iterations. + , scTol :: !Double + -- ^ Convergence tolerance on the weight delta (L-inf norm). + , scSampleWeights :: !(Maybe (VU.Vector Double)) + {- ^ Optional per-row sample weights, length @n@ (@Nothing@ is uniform). + Weights should have mean 1 (i.e. @Σ w_i = N@) so the Lipschitz bound stays + valid; see 'fitLinearCandidate' for the class-balanced construction. + -} + } + deriving (Eq, Show) + +defaultSolverConfig :: SolverConfig +defaultSolverConfig = + SolverConfig + { scL1Lambda = 0.005 + , scL2Lambda = 0.005 + , scMaxIter = 200 + , scTol = 1.0e-4 + , scSampleWeights = Nothing + } + +{- | Fit L1-regularized binary logistic regression by FISTA. Rows are feature +vectors of equal length; labels are in @{\-1,+1}@. Features are standardized +internally and weights de-standardized, so the model applies to raw values. +-} +fitL1Logistic :: + SolverConfig -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + V.Vector T.Text -> + LinearModel +{-# INLINEABLE fitL1Logistic #-} +fitL1Logistic = runFista logisticLoss (specNormLipschitz logisticLoss) + +{- | Fit any 'SmoothLoss' with the elastic-net proximal-gradient engine. The +Lipschitz constant uses the spectral norm of the standardized Gram matrix +(power iteration), tight for squared and squared-hinge losses. +-} +fitProx :: + SmoothLoss -> + SolverConfig -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + V.Vector T.Text -> + LinearModel +fitProx loss = runFista loss (specNormLipschitz loss) + +{- | FISTA smooth-part Lipschitz bound: the loss curvature bound times the +spectral norm of the standardized Gram (+1 for the intercept), via power +iteration. Tight for every smooth loss. +-} +specNormLipschitz :: SmoothLoss -> Matrix -> Int -> Double +specNormLipschitz loss xKept _ = + let n = V.length xKept + gramN = V.map (scaleV (1 / fromIntegral n)) (gram xKept) + (specNorm, _) = powerIterTop 50 gramN + in slCurvBound loss * (specNorm + 1) + +{- | Shared FISTA scaffolding: standardize, drop near-constant columns, run the +inner loop, de-standardize. @lipschitzOf@ returns the smooth-part Lipschitz +bound from the kept-feature matrix; the L2 contribution @λ₂@ is added here. +-} +runFista :: + SmoothLoss -> + (Matrix -> Int -> Double) -> + SolverConfig -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + V.Vector T.Text -> + LinearModel +runFista loss lipschitzOf cfg rows labels featureNames + | n == 0 || d == 0 = zeroModel + | otherwise = + let (!means, !stds, !variances) = columnStats rows + !keep = keptIndices variances + in if VU.null keep + then zeroModel + else + let !meansKept = gatherBy keep means + !stdsKept = gatherBy keep stds + !xKept = V.map (standardizeRowKept keep means stds) rows + !lipschitz = + lipschitzOf xKept (VU.length keep) + scL2Lambda cfg + (!wStdKept, !bStd) = + fistaLoop + loss + (scL1Lambda cfg) + (scL2Lambda cfg) + lipschitz + (scMaxIter cfg) + (scTol cfg) + (scSampleWeights cfg) + xKept + labels + (VU.replicate (VU.length keep) 0) + 0 + !wRawKept = VU.zipWith (/) wStdKept stdsKept + !bRaw = bStd - VU.sum (VU.zipWith (*) wRawKept meansKept) + in LinearModel (expandWeights d keep wRawKept) bRaw featureNames + where + !n = V.length rows + !d = V.length featureNames + zeroModel = LinearModel (VU.replicate d 0) 0 featureNames + +{- | Indices of columns whose variance clears the near-constant threshold. +Columns below it are dropped before fitting; their weight ends up @0@. +-} +keptIndices :: VU.Vector Double -> VU.Vector Int +keptIndices variances = + VU.fromList + [ j + | j <- [0 .. VU.length variances - 1] + , VU.unsafeIndex variances j >= 1.0e-12 + ] + +{- | Gather the entries of @v@ at @idxs@, preserving order. unsafeIndex is +safe: every index in @idxs@ is in range by construction. +-} +gatherBy :: VU.Vector Int -> VU.Vector Double -> VU.Vector Double +gatherBy idxs v = VU.map (VU.unsafeIndex v) idxs + +{- | Standardize one row to the kept columns only (subtract column mean, divide +by column std). unsafeIndex is safe: rows share the column layout. +-} +standardizeRowKept :: + VU.Vector Int -> + VU.Vector Double -> + VU.Vector Double -> + VU.Vector Double -> + VU.Vector Double +standardizeRowKept keep means stds row = VU.map standardizeAt keep + where + standardizeAt j = + (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j + +{- | Scatter kept-column weights back into a full-width vector, with @0@ for +the dropped (near-constant) columns. +-} +expandWeights :: Int -> VU.Vector Int -> VU.Vector Double -> VU.Vector Double +expandWeights d keep wKept = VU.create $ do + mv <- VUM.replicate d 0 + VU.iforM_ keep $ \k j -> VUM.unsafeWrite mv j (VU.unsafeIndex wKept k) + pure mv + +{- | Convert a fitted model to an 'Expr Bool' over its feature columns, +dropping zero-weight features. With no non-zero weights it returns the +constant @Lit (intercept > 0)@. +-} +modelToExpr :: LinearModel -> Expr Bool +modelToExpr m = + case nonZero of + [] -> F.lit (b > 0) + (w0, n0) : rest -> score rest (term w0 n0) .>. F.lit (0 :: Double) + where + b = lmIntercept m + nonZero = + [ (w, n) + | (w, n) <- zip (VU.toList (lmWeights m)) (V.toList (lmFeatureNames m)) + , w /= 0 + ] + term w n = F.lit w .*. (Col n :: Expr Double) + score rest first = foldl (\acc (w, n) -> acc .+. term w n) first rest .+. F.lit b + +{- | Per-column @(means, stds, variances)@ of a feature matrix. Cheaper than +'standardize' when only the statistics are needed. unsafeIndex within is +safe: all rows share width @d@. +-} +columnStats :: + V.Vector (VU.Vector Double) -> + (VU.Vector Double, VU.Vector Double, VU.Vector Double) +columnStats x + | V.null x = (VU.empty, VU.empty, VU.empty) + | otherwise = + let !d = VU.length (V.unsafeHead x) + !invN = 1 / fromIntegral (V.length x) + !means = columnMeans d invN x + !variances = columnVariances d invN means x + !stds = VU.map (\v -> if v < 1e-12 then 1 else sqrt v) variances + in (means, stds, variances) + +-- | Mean of each of the @d@ columns; @invN@ is @1 / nRows@. +columnMeans :: Int -> Double -> V.Vector (VU.Vector Double) -> VU.Vector Double +columnMeans d invN x = runST $ do + acc <- VUM.replicate d 0 + V.forM_ x $ \row -> + VU.iforM_ row $ \j v -> VUM.unsafeModify acc (+ v) j + scaleInPlace invN acc + VU.unsafeFreeze acc + +-- | Variance of each of the @d@ columns about the supplied @means@. +columnVariances :: + Int -> + Double -> + VU.Vector Double -> + V.Vector (VU.Vector Double) -> + VU.Vector Double +columnVariances d invN means x = runST $ do + acc <- VUM.replicate d 0 + V.forM_ x $ \row -> + VU.iforM_ row $ \j v -> + let !c = v - VU.unsafeIndex means j in VUM.unsafeModify acc (+ c * c) j + scaleInPlace invN acc + VU.unsafeFreeze acc + +-- | Multiply every element of a mutable vector by @factor@ in place. +scaleInPlace :: Double -> VUM.MVector s Double -> ST s () +scaleInPlace factor mv = go 0 + where + go !j + | j >= VUM.length mv = pure () + | otherwise = VUM.unsafeModify mv (* factor) j >> go (j + 1) + +{- | Standardize each column to zero mean and unit variance, also returning +@(means, stds, variances)@. Near-constant columns get std @1@; callers use +the raw variances to detect and drop them (see 'fitL1Logistic'). +-} +standardize :: + V.Vector (VU.Vector Double) -> + ( V.Vector (VU.Vector Double) + , VU.Vector Double + , VU.Vector Double + , VU.Vector Double + ) +standardize x + | V.null x = (x, VU.empty, VU.empty, VU.empty) + | otherwise = + let (!means, !stds, !variances) = columnStats x + !d = VU.length (V.unsafeHead x) + standardizeRow row = + VU.generate d $ \j -> + (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j + in (V.map standardizeRow x, means, stds, variances) + +{- | Proximal operator for the L1 norm: shrink @v@ toward zero by @lambda@, +clamping at zero. +-} +softThreshold :: Double -> Double -> Double +softThreshold lambda v + | v > lambda = v - lambda + | v < -lambda = v + lambda + | otherwise = 0 + +{- | Dot product of two unboxed vectors. Caller must ensure equal length; +lengths are not checked. +-} +dotProduct :: VU.Vector Double -> VU.Vector Double -> Double +dotProduct u v = VU.sum (VU.zipWith (*) u v) + +{- | Gradient of the average loss at @(w, b)@, returning @(gradW, gradB)@. +When @sampleWeights@ is @Just ws@ each row is scaled by @ws[i]@; with mean-1 +weights the @1/N@ normalisation is preserved exactly. +-} +lossGradient :: + SmoothLoss -> + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +lossGradient loss sampleWeights features labels w b = (gradW, gradB) + where + !invN = 1 / fromIntegral (V.length features) + !coeffs = rowCoeffs loss sampleWeights features labels w b invN + !gradW = accumulateGradW (VU.length w) features coeffs + !gradB = VU.sum coeffs + +{- | Per-row loss coefficient @c_i = ℓ'(y_i, z_i) / N@ at margin +@z_i = w·x_i + b@, optionally scaled by @ws[i]@. +-} +rowCoeffs :: + SmoothLoss -> + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + Double -> + VU.Vector Double +rowCoeffs loss sampleWeights features labels w b invN = + VU.generate (V.length features) $ \i -> + let !yi = VU.unsafeIndex labels i + !row = V.unsafeIndex features i + !z = dotProduct w row + b + !base = slGradZ loss yi z * invN + in case sampleWeights of + Nothing -> base + Just ws -> base * VU.unsafeIndex ws i + +{- | Accumulate the weight gradient in one pass over every (row, feature) +pair, scattering into a length-@d@ mutable vector. +-} +accumulateGradW :: + Int -> V.Vector (VU.Vector Double) -> VU.Vector Double -> VU.Vector Double +accumulateGradW d features coeffs = runST $ do + mv <- VUM.replicate d 0 + V.iforM_ features $ \i row -> + let !c = VU.unsafeIndex coeffs i + in VU.iforM_ row $ \j v -> VUM.unsafeModify mv (+ c * v) j + VU.unsafeFreeze mv + +{- | Inner FISTA loop over standardized features, returning the final @(w, b)@ +(the caller de-standardizes). @lambda1@/@lambda2@ are the L1/L2 strengths and +@lp@ the smooth-part Lipschitz constant driving the elastic-net prox step. +-} +fistaLoop :: + SmoothLoss -> + Double -> + Double -> + Double -> + Int -> + Double -> + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +fistaLoop loss lambda1 lambda2 lp maxIter tol sampleWeights features labels w0 b0 = + go 0 w0 b0 w0 b0 1.0 + where + !shrink = lambda1 / lp + !ridgeDenom = 1 + lambda2 / lp + !stepInv = 1 / lp + proxStep = fistaProxStep loss sampleWeights features labels shrink ridgeDenom stepInv + go !iter !xWPrev !xBPrev !yW !yB !t + | iter >= maxIter = (xWPrev, xBPrev) + | iter > 0 && delta < tol = (xW, xB) + | otherwise = go (iter + 1) xW xB yWNew yBNew tNew + where + (!xW, !xB) = proxStep yW yB + !delta = if VU.null xW then 0 else deltaInf xWPrev xW + (!yWNew, !yBNew, !tNew) = fistaMomentum t xWPrev xBPrev xW xB + +{- | One fused FISTA prox step: gradient step plus the Elastic-Net proximal +operator @softThreshold(z, λ₁/lp) / (1 + λ₂/lp)@ (soft-threshold then ridge +shrinkage). The intercept is unregularised. +-} +fistaProxStep :: + SmoothLoss -> + Maybe (VU.Vector Double) -> + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + Double -> + Double -> + Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double) +fistaProxStep loss sampleWeights features labels shrink ridgeDenom stepInv yW yB = + let (gW, gB) = lossGradient loss sampleWeights features labels yW yB + !wNew = + VU.zipWith + (\yi gi -> softThreshold shrink (yi - gi * stepInv) / ridgeDenom) + yW + gW + !bNew = yB - gB * stepInv + in (wNew, bNew) + +{- | Nesterov momentum extrapolation: new look-ahead point @(yW, yB)@ and the +updated step size @t@. +-} +fistaMomentum :: + Double -> + VU.Vector Double -> + Double -> + VU.Vector Double -> + Double -> + (VU.Vector Double, Double, Double) +fistaMomentum t xWPrev xBPrev xW xB = + let !tNew = (1 + sqrt (1 + 4 * t * t)) / 2 + !mom = (t - 1) / tNew + !yW = VU.zipWith (\new old -> new + mom * (new - old)) xW xWPrev + !yB = xB + mom * (xB - xBPrev) + in (yW, yB, tNew) + +{- | L-inf norm of the weight delta. unsafeIndex is safe: both vectors share +the same length by construction. +-} +{-# INLINE deltaInf #-} +deltaInf :: VU.Vector Double -> VU.Vector Double -> Double +deltaInf xWPrev = VU.ifoldl' (\acc i x -> max acc (abs (x - VU.unsafeIndex xWPrev i))) 0 diff --git a/src-internal/DataFrame/LinearSolver/Loss.hs b/src-internal/DataFrame/LinearSolver/Loss.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/LinearSolver/Loss.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE OverloadedStrings #-} + +{- | Smooth losses for the proximal-gradient engine. Each carries its +derivative @∂ℓ/∂z@ at @z = w·x + b@ and a global bound on the curvature +@∂²ℓ/∂z²@ (used for the FISTA step size). +-} +module DataFrame.LinearSolver.Loss ( + SmoothLoss (..), + sigmoid, + logisticLoss, + squaredLoss, + sqHingeLoss, +) where + +import qualified Data.Text as T + +{- | A convex, @C¹@ per-sample loss @ℓ(y, z)@. 'slGradZ' is @∂ℓ/∂z@; +'slCurvBound' bounds @∂²ℓ/∂z²@ over all @(y, z)@. +-} +data SmoothLoss = SmoothLoss + { slName :: !T.Text + , slGradZ :: Double -> Double -> Double + , slCurvBound :: !Double + } + +-- | Numerically stable logistic sigmoid. +sigmoid :: Double -> Double +sigmoid z + | z >= 0 = 1 / (1 + exp (-z)) + | otherwise = let ez = exp z in ez / (1 + ez) + +-- | Binary logistic loss for labels in @{\-1,+1}@: @ℓ = log(1 + exp(-y z))@. +logisticLoss :: SmoothLoss +logisticLoss = + SmoothLoss "logistic" (\y z -> negate (y * sigmoid (negate (y * z)))) 0.25 + +-- | Squared error for regression: @ℓ = ½ (z - y)²@. +squaredLoss :: SmoothLoss +squaredLoss = SmoothLoss "squared" (flip (-)) 1.0 + +-- | Squared hinge for classification (LinearSVC default), labels @{\-1,+1}@. +sqHingeLoss :: SmoothLoss +sqHingeLoss = + SmoothLoss + "squared_hinge" + (\y z -> let m = 1 - y * z in if m > 0 then negate (2 * y * m) else 0) + 2.0 diff --git a/src-internal/DataFrame/Random.hs b/src-internal/DataFrame/Random.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/Random.hs @@ -0,0 +1,119 @@ +{-# LANGUAGE CPP #-} + +{- | Deterministic, platform-independent random sampling for the stochastic +fitters. Built on @random@'s SplitMix 'StdGen'; the distributions here are our +own so a seeded fit is bit-reproducible across platforms. +-} +module DataFrame.Random ( + Gen, + mkGen, + splitGen, + nextWord64, + nextDouble, + nextIntR, + gaussianPair, + gaussianVector, + shuffleInts, + sampleIndices, +) where + +import Control.Monad (forM_) +import Control.Monad.ST (runST) +import Data.Bits (shiftR) +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word64) +import System.Random (StdGen, genWord64, mkStdGen) +import qualified System.Random as R + +-- | The pure splittable generator. A fit is a function of @(seed, data)@. +type Gen = StdGen + +-- | Seed a generator from an 'Int'. +mkGen :: Int -> Gen +mkGen = mkStdGen + +-- | Split into two independent generators. +splitGen :: Gen -> (Gen, Gen) +#if MIN_VERSION_random(1,3,0) +splitGen = R.splitGen +#else +splitGen = R.split +#endif + +-- | Raw 64-bit draw. +nextWord64 :: Gen -> (Word64, Gen) +nextWord64 = genWord64 + +-- | Uniform 'Double' in @[0, 1)@ from the top 53 bits (exact mantissa). +nextDouble :: Gen -> (Double, Gen) +nextDouble g = + let (w, g') = genWord64 g + d = fromIntegral (w `shiftR` 11) * (1 / 9007199254740992) + in (d, g') + +{- | Uniform 'Int' in the inclusive range @[lo, hi]@ by rejection sampling +(unbiased). Returns @lo@ when @hi <= lo@. +-} +nextIntR :: (Int, Int) -> Gen -> (Int, Gen) +nextIntR (lo, hi) g + | hi <= lo = (lo, g) + | otherwise = loop g + where + range = fromIntegral (hi - lo + 1) :: Word64 + threshold = negate range `mod` range + loop gg = + let (w, gg') = genWord64 gg + in if w >= threshold + then (lo + fromIntegral (w `mod` range), gg') + else loop gg' + +{- | A pair of independent standard normals via Box-Muller, consuming exactly two +uniforms so stream offsets stay data-independent. +-} +gaussianPair :: Gen -> ((Double, Double), Gen) +gaussianPair g = + let (u1, g1) = nextDouble g + (u2, g2) = nextDouble g1 + u1' = if u1 <= 0 then 2.220446049250313e-16 else u1 + r = sqrt (-(2 * log u1')) + a = 2 * pi * u2 + in ((r * cos a, r * sin a), g2) + +-- | A length-@n@ vector of standard normals. +gaussianVector :: Int -> Gen -> (VU.Vector Double, Gen) +gaussianVector n g0 = go n g0 [] + where + go k g acc + | k <= 0 = (VU.fromList (take n (reverse acc)), g) + | otherwise = + let ((z0, z1), g') = gaussianPair g + in go (k - 2) g' (z1 : z0 : acc) + +{- | A uniformly random permutation of @[0 .. n-1]@ (Fisher-Yates), threading the +generator purely. +-} +shuffleInts :: Int -> Gen -> (VU.Vector Int, Gen) +shuffleInts n g0 + | n <= 1 = (VU.enumFromN 0 (max 0 n), g0) + | otherwise = + let (swaps, g1) = genSwaps (n - 1) g0 [] + v = runST $ do + m <- VU.thaw (VU.enumFromN 0 n) + forM_ swaps $ uncurry (VUM.swap m) + VU.freeze m + in (v, g1) + where + genSwaps i g acc + | i < 1 = (reverse acc, g) + | otherwise = + let (j, g') = nextIntR (0, i) g + in genSwaps (i - 1) g' ((i, j) : acc) + +{- | @sampleIndices k n@ draws @k@ distinct indices from @[0 .. n-1]@ (the first +@k@ of a full shuffle); returns all @n@ when @k >= n@. +-} +sampleIndices :: Int -> Int -> Gen -> (VU.Vector Int, Gen) +sampleIndices k n g = + let (perm, g') = shuffleInts n g + in (VU.take (min k n) perm, g') diff --git a/src-internal/DataFrame/SymbolicRegression/Expr.hs b/src-internal/DataFrame/SymbolicRegression/Expr.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/SymbolicRegression/Expr.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE FlexibleContexts #-} + +{- | The symbolic-regression expression tree: a small first-order ADT with +vectorized evaluation and a total translation to a dataframe 'Expr Double'. +Division, log, and sqrt are protected so evaluation never produces @NaN@. +-} +module DataFrame.SymbolicRegression.Expr ( + SRExpr (..), + BinOp (..), + UnOp (..), + evalSR, + toDataFrameExpr, + srSize, + constants, + setConstants, + allBinOps, + allUnOps, +) where + +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression (Expr (..)) +import DataFrame.Operators ((.*.), (.+.), (.-.), (./.)) + +data BinOp = SAdd | SSub | SMul | SDiv + deriving (Eq, Ord, Show, Enum, Bounded) + +data UnOp = SNeg | SSin | SCos | SExp | SLog | SSqrt + deriving (Eq, Ord, Show, Enum, Bounded) + +-- | A symbolic-regression expression over feature variables and constants. +data SRExpr + = SVar !Int + | SConst !Double + | SUn !UnOp SRExpr + | SBin !BinOp SRExpr SRExpr + deriving (Eq, Ord, Show) + +allBinOps :: [BinOp] +allBinOps = [minBound .. maxBound] + +allUnOps :: [UnOp] +allUnOps = [minBound .. maxBound] + +{- | Evaluate over a feature matrix given column-major (@feats ! j@ is feature +@j@ across all rows). Protected operators keep results finite. +-} +evalSR :: V.Vector (VU.Vector Double) -> Int -> SRExpr -> VU.Vector Double +evalSR feats n = go + where + go (SVar j) + | j < V.length feats = feats V.! j + | otherwise = VU.replicate n 0 + go (SConst c) = VU.replicate n c + go (SUn op e) = VU.map (unFn op) (go e) + go (SBin op a b) = VU.zipWith (binFn op) (go a) (go b) + +binFn :: BinOp -> Double -> Double -> Double +binFn SAdd a b = a + b +binFn SSub a b = a - b +binFn SMul a b = a * b +binFn SDiv a b = if abs b < 1e-9 then 1 else a / b + +unFn :: UnOp -> Double -> Double +unFn SNeg = negate +unFn SSin = sin +unFn SCos = cos +unFn SExp = exp . min 50 +unFn SLog = \x -> log (abs x + 1e-9) +unFn SSqrt = sqrt . abs + +-- | Translate to a dataframe expression over the named feature columns. +toDataFrameExpr :: V.Vector T.Text -> SRExpr -> Expr Double +toDataFrameExpr names = go + where + go (SVar j) + | j < V.length names = Col (names V.! j) + | otherwise = F.lit 0 + go (SConst c) = F.lit c + go (SUn op e) = unExpr op (go e) + go (SBin op a b) = binExpr op (go a) (go b) + unExpr SNeg = negate + unExpr SSin = sin + unExpr SCos = cos + unExpr SExp = exp + unExpr SLog = log + unExpr SSqrt = sqrt + binExpr SAdd = (.+.) + binExpr SSub = (.-.) + binExpr SMul = (.*.) + binExpr SDiv = (./.) + +srSize :: SRExpr -> Int +srSize (SVar _) = 1 +srSize (SConst _) = 1 +srSize (SUn _ e) = 1 + srSize e +srSize (SBin _ a b) = 1 + srSize a + srSize b + +-- | The constant values in left-to-right traversal order. +constants :: SRExpr -> [Double] +constants (SConst c) = [c] +constants (SVar _) = [] +constants (SUn _ e) = constants e +constants (SBin _ a b) = constants a ++ constants b + +-- | Replace the constants in traversal order; extra values are ignored. +setConstants :: [Double] -> SRExpr -> SRExpr +setConstants vals e = fst (go vals e) + where + go vs (SConst _) = case vs of + (v : rest) -> (SConst v, rest) + [] -> (SConst 0, []) + go vs (SVar j) = (SVar j, vs) + go vs (SUn op a) = let (a', vs') = go vs a in (SUn op a', vs') + go vs (SBin op a b) = + let (a', vs') = go vs a + (b', vs'') = go vs' b + in (SBin op a' b', vs'') diff --git a/src-internal/DataFrame/SymbolicRegression/GP.hs b/src-internal/DataFrame/SymbolicRegression/GP.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/SymbolicRegression/GP.hs @@ -0,0 +1,206 @@ +{- | A compact generational genetic-programming search over 'SRExpr': ramped +initialization, tournament selection, subtree crossover/mutation, elitism, and a +complexity-keyed Pareto archive. Deterministic given the seed. +-} +module DataFrame.SymbolicRegression.GP ( + GPParams (..), + runGP, +) where + +import Data.List (foldl', minimumBy, sortBy) +import qualified Data.Map.Strict as M +import Data.Ord (comparing) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +import DataFrame.Random (Gen, nextDouble, nextIntR) +import DataFrame.SymbolicRegression.Expr +import DataFrame.SymbolicRegression.Optimize ( + meanSquaredError, + optimizeConstants, + ) +import DataFrame.SymbolicRegression.Simplify (simplify) + +-- | GP hyper-parameters resolved from the public config. +data GPParams = GPParams + { gpFeats :: !(V.Vector (VU.Vector Double)) + , gpN :: !Int + , gpTarget :: !(VU.Vector Double) + , gpNVars :: !Int + , gpUnOps :: ![UnOp] + , gpPopSize :: !Int + , gpGenerations :: !Int + , gpMaxSize :: !Int + , gpTournament :: !Int + , gpCrossoverP :: !Double + , gpMutationP :: !Double + , gpOptimizeP :: !Double + , gpParsimony :: !Double + } + +type Scored = (SRExpr, Double) + +-- | Run the search; returns @(best, pareto front, generations run)@. +runGP :: GPParams -> Gen -> (SRExpr, [(Int, Double, SRExpr)], Int) +runGP p g0 = + let (pop0, g1) = initPop p g0 + scored0 = map (scoreOf p) pop0 + arch0 = foldl' (archiveInsert p) M.empty scored0 + (_, finalArch, gN, _) = + iterate' 0 scored0 arch0 g1 + best = bestOfArchive finalArch + front = + [ (sz, mse, e) + | (sz, (mse, e)) <- M.toList finalArch + ] + in (snd3 best, sortBy (comparing fst3) front, gN) + where + iterate' gen pop arch g + | gen >= gpGenerations p = (pop, arch, gen, g) + | otherwise = + let (pop', g') = nextGen p pop g + arch' = foldl' (archiveInsert p) arch pop' + in iterate' (gen + 1) pop' arch' g' + fst3 (a, _, _) = a + snd3 (_, b, _) = b + bestOfArchive arch = + case M.toList arch of + [] -> (0 :: Int, SConst 0, 1 / 0) + xs -> + let (sz, (mse, e)) = minimumBy (comparing (fst . snd)) xs + in (sz, e, mse) + +scoreOf :: GPParams -> SRExpr -> Scored +scoreOf p e = (e, meanSquaredError (gpFeats p) (gpN p) (gpTarget p) e) + +fitness :: GPParams -> Scored -> Double +fitness p (e, mse) = mse + gpParsimony p * fromIntegral (srSize e) + +archiveInsert :: + GPParams -> M.Map Int (Double, SRExpr) -> Scored -> M.Map Int (Double, SRExpr) +archiveInsert _ arch (e, mse) + | isNaN mse || isInfinite mse = arch + | otherwise = + let key = srSize (simplify e) + in M.insertWith better key (mse, e) arch + where + better newv@(m1, _) oldv@(m2, _) = if m1 < m2 then newv else oldv + +initPop :: GPParams -> Gen -> ([SRExpr], Gen) +initPop p = go (gpPopSize p) [] + where + go 0 acc g = (acc, g) + go k acc g = + let (depth, g1) = nextIntR (1, 4) g + (e, g2) = randomExpr p depth g1 + in go (k - 1) (e : acc) g2 + +randomExpr :: GPParams -> Int -> Gen -> (SRExpr, Gen) +randomExpr p depth g + | depth <= 1 = randomLeaf p g + | otherwise = + let (r, g1) = nextDouble g + in if r < 0.3 + then randomLeaf p g1 + else + let (isUn, g2) = nextDouble g1 + in if isUn < 0.3 && not (null (gpUnOps p)) + then + let (oi, g3) = nextIntR (0, length (gpUnOps p) - 1) g2 + (e, g4) = randomExpr p (depth - 1) g3 + in (SUn (gpUnOps p !! oi) e, g4) + else + let (oi, g3) = nextIntR (0, length allBinOps - 1) g2 + (a, g4) = randomExpr p (depth - 1) g3 + (b, g5) = randomExpr p (depth - 1) g4 + in (SBin (allBinOps !! oi) a b, g5) + +randomLeaf :: GPParams -> Gen -> (SRExpr, Gen) +randomLeaf p g = + let (r, g1) = nextDouble g + in if r < 0.6 && gpNVars p > 0 + then let (j, g2) = nextIntR (0, gpNVars p - 1) g1 in (SVar j, g2) + else let (c, g2) = nextDouble g1 in (SConst (c * 4 - 2), g2) + +nextGen :: GPParams -> [Scored] -> Gen -> ([Scored], Gen) +nextGen p pop g0 = + let elite = minimumBy (comparing (fitness p)) pop + (rest, g1) = go (gpPopSize p - 1) [] g0 + in (elite : rest, g1) + where + go 0 acc g = (acc, g) + go k acc g = + let (child, g') = breed p pop g + scored = optimizeMaybe p child g' + in go (k - 1) (fst scored : acc) (snd scored) + +optimizeMaybe :: GPParams -> SRExpr -> Gen -> (Scored, Gen) +optimizeMaybe p e g = + let (r, g1) = nextDouble g + e' = + if r < gpOptimizeP p + then optimizeConstants (gpFeats p) (gpN p) (gpTarget p) 15 e + else e + in (scoreOf p e', g1) + +breed :: GPParams -> [Scored] -> Gen -> (SRExpr, Gen) +breed p pop g0 = + let (pa, g1) = tournament p pop g0 + (doX, g2) = nextDouble g1 + (child, g3) = + if doX < gpCrossoverP p + then + let (pb, g2') = tournament p pop g2 + (c, g3') = crossover pa pb g2' + in (c, g3') + else (pa, g2) + (doM, g4) = nextDouble g3 + (child', g5) = + if doM < gpMutationP p then mutate p child g4 else (child, g4) + capped = if srSize child' > gpMaxSize p then pa else child' + in (simplify capped, g5) + +tournament :: GPParams -> [Scored] -> Gen -> (SRExpr, Gen) +tournament p pop g0 = + let (picks, g1) = pickN (gpTournament p) g0 + chosen = map (pop !!) picks + in (fst (minimumBy (comparing (fitness p)) chosen), g1) + where + n = length pop + pickN 0 g = ([], g) + pickN k g = + let (i, g') = nextIntR (0, n - 1) g + (is, g'') = pickN (k - 1) g' + in (i : is, g'') + +crossover :: SRExpr -> SRExpr -> Gen -> (SRExpr, Gen) +crossover a b g0 = + let (ia, g1) = nextIntR (0, srSize a - 1) g0 + (ib, g2) = nextIntR (0, srSize b - 1) g1 + sub = subtreeAt ib b + in (replaceAt ia a sub, g2) + +mutate :: GPParams -> SRExpr -> Gen -> (SRExpr, Gen) +mutate p e g0 = + let (i, g1) = nextIntR (0, srSize e - 1) g0 + (depth, g2) = nextIntR (1, 3) g1 + (newSub, g3) = randomExpr p depth g2 + in (replaceAt i e newSub, g3) + +subtreeAt :: Int -> SRExpr -> SRExpr +subtreeAt 0 e = e +subtreeAt i (SUn _ e) = subtreeAt (i - 1) e +subtreeAt i (SBin _ a b) = + let sa = srSize a + in if i <= sa then subtreeAt (i - 1) a else subtreeAt (i - 1 - sa) b +subtreeAt _ e = e + +replaceAt :: Int -> SRExpr -> SRExpr -> SRExpr +replaceAt 0 _ new = new +replaceAt i (SUn op e) new = SUn op (replaceAt (i - 1) e new) +replaceAt i (SBin op a b) new = + let sa = srSize a + in if i <= sa + then SBin op (replaceAt (i - 1) a new) b + else SBin op a (replaceAt (i - 1 - sa) b new) +replaceAt _ e _ = e diff --git a/src-internal/DataFrame/SymbolicRegression/Optimize.hs b/src-internal/DataFrame/SymbolicRegression/Optimize.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/SymbolicRegression/Optimize.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE BangPatterns #-} + +{- | Constant optimization for symbolic-regression candidates: finite-difference +gradient descent with backtracking line search on the embedded constants. Pure +and dependency-free; effective at the one-to-few constants a tree carries. +-} +module DataFrame.SymbolicRegression.Optimize ( + optimizeConstants, + meanSquaredError, +) where + +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +import DataFrame.SymbolicRegression.Expr ( + SRExpr, + constants, + evalSR, + setConstants, + ) + +-- | Mean squared error of an expression's predictions against the target. +meanSquaredError :: + V.Vector (VU.Vector Double) -> Int -> VU.Vector Double -> SRExpr -> Double +meanSquaredError feats n target e = + let pred = evalSR feats n e + diff = VU.zipWith (-) pred target + in VU.sum (VU.map (\x -> x * x) diff) / fromIntegral (max 1 n) + +-- | Refine an expression's constants to reduce MSE (no-op when constant-free). +optimizeConstants :: + V.Vector (VU.Vector Double) -> + Int -> + VU.Vector Double -> + Int -> + SRExpr -> + SRExpr +optimizeConstants feats n target iters expr + | null theta0 = expr + | otherwise = setConstants (descend iters theta0) expr + where + theta0 = constants expr + eps = 1e-6 + mseAt theta = meanSquaredError feats n target (setConstants theta expr) + descend 0 theta = theta + descend k theta = + let f0 = mseAt theta + g = numGrad theta + gn = sqrt (sum (map (\x -> x * x) g)) + in if gn < 1e-10 + then theta + else + let theta' = lineSearch theta g f0 + in if theta' == theta then theta else descend (k - 1) theta' + numGrad theta = + [ (mseAt (bump i eps theta) - mseAt (bump i (negate eps) theta)) / (2 * eps) + | i <- [0 .. length theta - 1] + ] + bump i delta theta = + [if j == i then t + delta else t | (j, t) <- zip [0 ..] theta] + lineSearch theta g f0 = go (1.0 :: Double) + where + go !step + | step < 1e-8 = theta + | otherwise = + let theta' = zipWith (\t gi -> t - step * gi) theta g + in if mseAt theta' < f0 then theta' else go (step / 2) diff --git a/src-internal/DataFrame/SymbolicRegression/Simplify.hs b/src-internal/DataFrame/SymbolicRegression/Simplify.hs new file mode 100644 --- /dev/null +++ b/src-internal/DataFrame/SymbolicRegression/Simplify.hs @@ -0,0 +1,63 @@ +{- | A fuel-bounded, deterministic algebraic simplifier — the dependency-light +stand-in for equality saturation. Used as a canonical dedup key for the Pareto +archive and to tidy reported expressions; total and evaluation-preserving. +-} +module DataFrame.SymbolicRegression.Simplify ( + simplify, +) where + +import DataFrame.SymbolicRegression.Expr (BinOp (..), SRExpr (..), UnOp (..)) + +-- | Simplify to a fixed point (bounded by a fuel counter). +simplify :: SRExpr -> SRExpr +simplify = go (10 :: Int) + where + go 0 e = e + go fuel e = + let e' = step e + in if e' == e then e else go (fuel - 1) e' + +step :: SRExpr -> SRExpr +step (SUn op e) = simplifyUn op (step e) +step (SBin op a b) = simplifyBin op (step a) (step b) +step e = e + +simplifyUn :: UnOp -> SRExpr -> SRExpr +simplifyUn SNeg (SUn SNeg e) = e +simplifyUn op (SConst c) = SConst (foldUn op c) +simplifyUn op e = SUn op e + +simplifyBin :: BinOp -> SRExpr -> SRExpr -> SRExpr +simplifyBin op (SConst a) (SConst b) = SConst (foldBin op a b) +simplifyBin SAdd a (SConst 0) = a +simplifyBin SAdd (SConst 0) b = b +simplifyBin SSub a (SConst 0) = a +simplifyBin SSub a b | a == b = SConst 0 +simplifyBin SMul _ (SConst 0) = SConst 0 +simplifyBin SMul (SConst 0) _ = SConst 0 +simplifyBin SMul a (SConst 1) = a +simplifyBin SMul (SConst 1) b = b +simplifyBin SDiv a (SConst 1) = a +simplifyBin SDiv a b | a == b = SConst 1 +simplifyBin op a b + | commutative op && a > b = SBin op b a + | otherwise = SBin op a b + +commutative :: BinOp -> Bool +commutative SAdd = True +commutative SMul = True +commutative _ = False + +foldBin :: BinOp -> Double -> Double -> Double +foldBin SAdd a b = a + b +foldBin SSub a b = a - b +foldBin SMul a b = a * b +foldBin SDiv a b = if abs b < 1e-9 then 1 else a / b + +foldUn :: UnOp -> Double -> Double +foldUn SNeg = negate +foldUn SSin = sin +foldUn SCos = cos +foldUn SExp = exp . min 50 +foldUn SLog = \x -> log (abs x + 1e-9) +foldUn SSqrt = sqrt . abs diff --git a/src/DataFrame/Boosting/AdaBoost.hs b/src/DataFrame/Boosting/AdaBoost.hs --- a/src/DataFrame/Boosting/AdaBoost.hs +++ b/src/DataFrame/Boosting/AdaBoost.hs @@ -4,6 +4,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} {- | AdaBoost (SAMME) over short, sample-weighted classification trees. The @@ -12,6 +13,7 @@ path is untouched. 'predict' is the arg-max of weighted votes. -} module DataFrame.Boosting.AdaBoost ( + module DataFrame.Model, AdaBoostConfig (..), defaultAdaBoostConfig, AdaBoostModel (..), @@ -35,7 +37,7 @@ import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..)) import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.==.)) data AdaBoostConfig = AdaBoostConfig @@ -55,10 +57,12 @@ } deriving (Show) -instance (Columnable a, Ord a) => Fit AdaBoostConfig (Expr a) (AdaBoostModel a) where +instance (Columnable a, Ord a) => Fit AdaBoostConfig (Expr a) where + type ModelOf AdaBoostConfig (Expr a) = (AdaBoostModel a) fit = fitAdaBoost -instance (Columnable a, Ord a) => Predict (AdaBoostModel a) a where +instance (Columnable a, Ord a) => Predict (AdaBoostModel a) where + type Prediction (AdaBoostModel a) = Expr a predict = adaBoostExpr -- | Fit an AdaBoost-SAMME classifier. diff --git a/src/DataFrame/Boosting/GBM.hs b/src/DataFrame/Boosting/GBM.hs --- a/src/DataFrame/Boosting/GBM.hs +++ b/src/DataFrame/Boosting/GBM.hs @@ -4,6 +4,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} {- | Gradient boosting of regression trees (Friedman). Trees are fitted to the negative gradient of the loss each round and accumulated with a shrinkage @@ -12,6 +13,7 @@ 'gbDecisionExpr' give the classification probability / decision. -} module DataFrame.Boosting.GBM ( + module DataFrame.Model, GBLoss (..), GBConfig (..), defaultGBConfig, @@ -37,7 +39,7 @@ import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..), getColumns) import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.>.)) -- | The boosting loss. @@ -77,10 +79,12 @@ } deriving (Show) -instance Fit GBConfig (Expr Double) GBModel where +instance Fit GBConfig (Expr Double) where + type ModelOf GBConfig (Expr Double) = GBModel fit = fitGBM -instance Predict GBModel Double where +instance Predict GBModel where + type Prediction GBModel = Expr Double predict = gbExpr -- | Fit a gradient-boosting ensemble predicting @target@ from the other columns. diff --git a/src/DataFrame/DBSCAN.hs b/src/DataFrame/DBSCAN.hs --- a/src/DataFrame/DBSCAN.hs +++ b/src/DataFrame/DBSCAN.hs @@ -3,6 +3,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} {- | Density-based clustering (DBSCAN). Brute-force @O(n²)@ region queries, no spatial index — suitable for the in-memory scales this library targets. DBSCAN @@ -11,6 +12,7 @@ interpretable decision-tree surrogate on the cluster labels instead. -} module DataFrame.DBSCAN ( + module DataFrame.Model, DBSCANConfig (..), defaultDBSCANConfig, DBSCANModel (..), @@ -35,7 +37,7 @@ import DataFrame.Internal.DataFrame (DataFrame, fromNamedColumns) import DataFrame.Internal.Expression (Expr) import DataFrame.LinearAlgebra (epsNeighbors) -import DataFrame.Model (Fit (..)) +import DataFrame.Model data DBSCANConfig = DBSCANConfig { dbEps :: !Double @@ -56,7 +58,8 @@ } deriving (Eq, Show) -instance Fit DBSCANConfig [Expr Double] DBSCANModel where +instance Fit DBSCANConfig [Expr Double] where + type ModelOf DBSCANConfig [Expr Double] = DBSCANModel fit = fitDBSCAN -- | Cluster the feature columns with DBSCAN. diff --git a/src/DataFrame/DecisionTree.hs b/src/DataFrame/DecisionTree.hs --- a/src/DataFrame/DecisionTree.hs +++ b/src/DataFrame/DecisionTree.hs @@ -1,29 +1,53 @@ -{- | Decision-tree training on DataFrames: a faithful CART tree refined by Tree -Alternating Optimization (TAO). This module re-exports the implementation, -which is split across the @DataFrame.DecisionTree.*@ modules. +{- | Interpretable decision trees on DataFrames (CART refined by TAO). The +curated public surface: classifier/regressor configs, fitted-model records, +and their 'Fit'\/'Predict' instances. -} module DataFrame.DecisionTree ( - module DataFrame.DecisionTree.Types, - module DataFrame.DecisionTree.CondVec, - module DataFrame.DecisionTree.Cart, - module DataFrame.DecisionTree.Numeric, - module DataFrame.DecisionTree.Categorical, - module DataFrame.DecisionTree.Pool, - module DataFrame.DecisionTree.Predict, - module DataFrame.DecisionTree.Linear, - module DataFrame.DecisionTree.Tao, - module DataFrame.DecisionTree.Prune, - module DataFrame.DecisionTree.Fit, + -- * Estimators + + {- | Fitted-model records with their @Fit@\/@Predict@ instances and the + estimator classes (via "DataFrame.Model"). + -} + module DataFrame.DecisionTree.Model, + + -- * Classifier configuration + TreeConfig (..), + defaultTreeConfig, + SynthConfig (..), + defaultSynthConfig, + ColumnOrdering (..), + orderable, + defaultColumnOrdering, + withOrdFrom, + + -- * Regressor configuration + RegTreeConfig (..), + defaultRegTreeConfig, + + -- * Fitted tree structure + Tree (..), + treeDepth, + + -- * Solver configuration (fills @TreeConfig.linearSolverConfig@) + SolverConfig (..), + defaultSolverConfig, ) where -import DataFrame.DecisionTree.Cart -import DataFrame.DecisionTree.Categorical -import DataFrame.DecisionTree.CondVec -import DataFrame.DecisionTree.Fit -import DataFrame.DecisionTree.Linear -import DataFrame.DecisionTree.Numeric -import DataFrame.DecisionTree.Pool -import DataFrame.DecisionTree.Predict -import DataFrame.DecisionTree.Prune -import DataFrame.DecisionTree.Tao -import DataFrame.DecisionTree.Types +import DataFrame.DecisionTree.Model +import DataFrame.DecisionTree.Regression ( + RegTreeConfig (..), + defaultRegTreeConfig, + ) +import DataFrame.DecisionTree.Types ( + ColumnOrdering (..), + SynthConfig (..), + Tree (..), + TreeConfig (..), + defaultColumnOrdering, + defaultSynthConfig, + defaultTreeConfig, + orderable, + treeDepth, + withOrdFrom, + ) +import DataFrame.LinearSolver (SolverConfig (..), defaultSolverConfig) diff --git a/src/DataFrame/DecisionTree/Cart.hs b/src/DataFrame/DecisionTree/Cart.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Cart.hs +++ /dev/null @@ -1,291 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | sklearn-faithful CART initializer used to seed TAO. One-hot encodes -categoricals, splits on exact (unsmoothed) Gini over midpoint thresholds -(@<=@ routes left), and emits a @Tree@ predicting identically to -@DecisionTreeClassifier(criterion='gini')@ on continuous features. --} -module DataFrame.DecisionTree.Cart ( - CartFeature (..), - CartNode (..), - sortIndicesByValue, - buildCartTree, - cartFeatures, - cartTargetLabels, -) where - -import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Internal.Types -import DataFrame.Operations.Core (nRows) -import DataFrame.Operators - -import Data.Either (fromRight) -import Data.Function (on) -import Data.List (foldl') -import qualified Data.Map.Strict as M -import qualified Data.Set as Set -import qualified Data.Text as T -import Data.Type.Equality (testEquality, (:~:) (..)) -import qualified Data.Vector as V -import qualified Data.Vector.Algorithms.Merge as VA -import qualified Data.Vector.Unboxed as VU -import Type.Reflection (typeRep) - -{- | A one-hot feature column: per-row Double values plus the sklearn LEFT -predicate (@x <= threshold@) over the ORIGINAL DataFrame. --} -data CartFeature = CartFeature - { cfValues :: !(VU.Vector Double) - , cfPred :: !(Double -> Expr Bool) - } - --- | Pre-'Tree' CART node: a leaf class id, or a split on feature @j@. -data CartNode = CLeaf !Int | CSplit !Int !Double !CartNode !CartNode - --- | Immutable per-fit context for the CART recursion. -data CartCtx = CartCtx - { ctxFeats :: !(V.Vector CartFeature) - , ctxNFeats :: !Int - , ctxCodes :: !(VU.Vector Int) - , ctxNClasses :: !Int - , ctxMaxDepth :: !Int - , ctxMinLeaf :: !Int - } - -{- | Indices @0..n-1@ stably sorted by their value (ascending), ties keeping -ascending index. In-place unboxed merge sort — no boxed-list allocation. --} -sortIndicesByValue :: VU.Vector Double -> VU.Vector Int -sortIndicesByValue vs = - VU.create $ do - mv <- VU.thaw (VU.enumFromN 0 (VU.length vs)) - VA.sortBy (compare `on` (vs VU.!)) mv - pure mv - -buildCartTree :: - forall a. (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> Tree a -buildCartTree cfg target df = - cartToTree feats classes (buildCartNode ctx 0 (VU.enumFromN 0 nAll) featSorted) - where - nAll = nRows df - feats = V.fromList (cartFeatures target df) - featSorted = V.map (sortIndicesByValue . cfValues) feats - labels = cartLabels @a df target - classes = cartClasses labels - ctx = - CartCtx - feats - (V.length feats) - (classCodes classes labels) - (V.length classes) - (maxTreeDepth cfg) - (max 1 (minLeafSize cfg)) - -cartLabels :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a -cartLabels df target = case interpret @a df (Col target) of - Right (TColumn column) -> fromRight err (toVector @a column) - _ -> err - where - err = error "buildCartTree: cannot interpret target column" - -cartClasses :: (Ord a) => V.Vector a -> V.Vector a -cartClasses = V.fromList . Set.toList . Set.fromList . V.toList - -classCodes :: (Ord a) => V.Vector a -> V.Vector a -> VU.Vector Int -classCodes classes labels = VU.generate (V.length labels) (\i -> M.findWithDefault 0 (labels V.! i) ix) - where - ix = M.fromList (zip (V.toList classes) [0 ..]) - -cartToTree :: V.Vector CartFeature -> V.Vector a -> CartNode -> Tree a -cartToTree feats classes = go - where - go (CLeaf cid) = Leaf (classes V.! cid) - go (CSplit fj thr l r) = Branch (cfPred (feats V.! fj) thr) (go l) (go r) - -classCounts :: CartCtx -> VU.Vector Int -> VU.Vector Int -classCounts ctx idxs = - VU.accumulate - (+) - (VU.replicate (ctxNClasses ctx) 0) - (VU.map (\i -> (ctxCodes ctx VU.! i, 1)) idxs) - -isPure :: VU.Vector Int -> Bool -isPure counts = VU.length (VU.filter (> 0) counts) <= 1 - -buildCartNode :: - CartCtx -> Int -> VU.Vector Int -> V.Vector (VU.Vector Int) -> CartNode -buildCartNode ctx depth idxs sortedByFeat - | VU.length idxs < 2 || depth >= ctxMaxDepth ctx || isPure counts = leaf - | otherwise = - maybe - leaf - (splitNode ctx depth idxs sortedByFeat) - (bestSplit ctx sortedByFeat counts n) - where - n = VU.length idxs - counts = classCounts ctx idxs - leaf = CLeaf (VU.maxIndex counts) - -splitNode :: - CartCtx -> - Int -> - VU.Vector Int -> - V.Vector (VU.Vector Int) -> - (Int, Double) -> - CartNode -splitNode ctx depth idxs sortedByFeat (fj, thr) = - CSplit fj thr (rec leftIdx leftSorted) (rec rightIdx rightSorted) - where - vals = cfValues (ctxFeats ctx V.! fj) - leftIdx = VU.filter (\i -> vals VU.! i <= thr) idxs - rightIdx = VU.filter (\i -> vals VU.! i > thr) idxs - leftSorted = V.map (VU.filter (\i -> vals VU.! i <= thr)) sortedByFeat - rightSorted = V.map (VU.filter (\i -> vals VU.! i > thr)) sortedByFeat - rec = buildCartNode ctx (depth + 1) - -{- | Minimum weighted-child-Gini @(feature, threshold)@; the first feature wins -ties; 'Nothing' when no feature has a leaf-size-respecting threshold. --} -bestSplit :: - CartCtx -> - V.Vector (VU.Vector Int) -> - VU.Vector Int -> - Int -> - Maybe (Int, Double) -bestSplit ctx sortedByFeat counts n = - fmap (\(_, j, t) -> (j, t)) (foldl' consider Nothing [0 .. ctxNFeats ctx - 1]) - where - total = VU.toList counts - consider acc fj = case sweepFeature ctx total (sortedByFeat V.! fj) (ctxFeats ctx V.! fj) n of - Just (g, thr) | maybe True (\(gB, _, _) -> g < gB) acc -> Just (g, fj, thr) - _ -> acc - -{- | Accumulator while sweeping a feature's sorted rows: best @(gini, thr)@ so -far, per-class left counts, rows moved left, and the previous value seen. --} -data Sweep = Sweep - { swBest :: !(Maybe (Double, Double)) - , swLeft :: ![Int] - , swMoved :: !Int - , swPrev :: !Double - } - -sweepFeature :: - CartCtx -> - [Int] -> - VU.Vector Int -> - CartFeature -> - Int -> - Maybe (Double, Double) -sweepFeature ctx total si feat n = - swBest - ( foldl' - step - (Sweep Nothing (replicate (ctxNClasses ctx) 0) 0 (0 / 0)) - [0 .. VU.length si - 1] - ) - where - vals = cfValues feat - step s k = advance ctx total n (vals VU.! i) (ctxCodes ctx VU.! i) s - where - i = si VU.! k - -advance :: CartCtx -> [Int] -> Int -> Double -> Int -> Sweep -> Sweep -advance ctx total n v c s = - Sweep - (considerThreshold ctx total n v s) - (bumpClass c (swLeft s)) - (swMoved s + 1) - v - -considerThreshold :: - CartCtx -> [Int] -> Int -> Double -> Sweep -> Maybe (Double, Double) -considerThreshold ctx total n v s - | swMoved s >= ctxMinLeaf ctx - , n - swMoved s >= ctxMinLeaf ctx - , v > swPrev s + 1e-7 = - keepBetter - (swBest s) - (weightedGini total (swLeft s) (swMoved s) n) - ((swPrev s + v) / 2) - | otherwise = swBest s - -keepBetter :: - Maybe (Double, Double) -> Double -> Double -> Maybe (Double, Double) -keepBetter best g thr = case best of - Just (wb, _) | wb <= g -> best - _ -> Just (g, thr) - -weightedGini :: [Int] -> [Int] -> Int -> Int -> Double -weightedGini total leftAcc nl n = - ( fromIntegral nl * giniImpurity leftAcc nl - + fromIntegral nr * giniImpurity rightAcc nr - ) - / fromIntegral n - where - nr = n - nl - rightAcc = zipWith (-) total leftAcc - --- | Gini impurity @1 - Σ (c/m)²@ of a class-count list of total @m@. -giniImpurity :: [Int] -> Int -> Double -giniImpurity _ 0 = 0 -giniImpurity cs m = 1 - sum [let p = fromIntegral c / fromIntegral m in p * p | c <- cs] - -bumpClass :: Int -> [Int] -> [Int] -bumpClass c = zipWith (\j x -> if j == c then x + 1 else x) [0 ..] - --- | One-hot features in @pd.get_dummies(drop_first=False)@ column order. -cartFeatures :: T.Text -> DataFrame -> [CartFeature] -cartFeatures target df = concatMap (featuresOfColumn df) (filter (/= target) (columnNames df)) - -featuresOfColumn :: DataFrame -> T.Text -> [CartFeature] -featuresOfColumn df c = case unsafeGetColumn c df of - UnboxedColumn _ (v :: VU.Vector b) -> numericFeature @b c v - BoxedColumn _ (v :: V.Vector b) -> oneHotFeatures @b (nRows df) c v - pt@(PackedText _ _) -> case materializePacked pt of - BoxedColumn _ (v :: V.Vector b) -> oneHotFeatures @b (nRows df) c v - _ -> [] - -numericFeature :: - forall b. (Columnable b, VU.Unbox b) => T.Text -> VU.Vector b -> [CartFeature] -numericFeature c v = case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [CartFeature v (\t -> F.col @Double c .<=. F.lit t)] - Nothing -> case sIntegral @b of - STrue -> - [ CartFeature (VU.map fromIntegral v) (\t -> F.toDouble (F.col @b c) .<=. F.lit t) - ] - SFalse -> [] - -oneHotFeatures :: - forall b. (Columnable b) => Int -> T.Text -> V.Vector b -> [CartFeature] -oneHotFeatures nAll c v = case testEquality (typeRep @b) (typeRep @T.Text) of - Just Refl -> [oneHot nAll c v cat | cat <- Set.toList (Set.fromList (V.toList v))] - Nothing -> [] - -oneHot :: Int -> T.Text -> V.Vector T.Text -> T.Text -> CartFeature -oneHot nAll c v cat = - CartFeature - (VU.generate nAll (\i -> if v V.! i == cat then 1 else 0)) - (const (F.col @T.Text c ./=. F.lit cat)) - --- | Target column as string labels (matches pandas @y.astype(str)@). -cartTargetLabels :: T.Text -> DataFrame -> V.Vector T.Text -cartTargetLabels target df = case unsafeGetColumn target df of - BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @T.Text) of - Just Refl -> v - Nothing -> V.map (T.pack . show) v - UnboxedColumn _ (v :: VU.Vector b) -> V.map (T.pack . show) (V.convert v) - pt@(PackedText _ _) -> case materializePacked pt of - BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @b) (typeRep @T.Text) of - Just Refl -> v - Nothing -> V.map (T.pack . show) v - _ -> V.empty diff --git a/src/DataFrame/DecisionTree/Categorical.hs b/src/DataFrame/DecisionTree/Categorical.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Categorical.hs +++ /dev/null @@ -1,381 +0,0 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Categorical split candidates: Breiman prefixes for binary targets, -subset/singleton enumeration otherwise, and cross-column equality. Each -value-list yields an OR-of-equalities condition (as an expression or a -directly-read membership truth vector). --} -module DataFrame.DecisionTree.Categorical ( - TargetInfo (..), - mkTargetInfo, - distinctValuesUpTo, - validBoxedValues, - orEqs, - subsetSplits, - subsetLists, - singletonSplits, - singletonLists, - breimanPrefixSplits, - breimanPrefixLists, - catValueLists, - membershipVec, - crossColumnConds, - discreteConditions, - discreteCondVecs, -) where - -import DataFrame.DecisionTree.CondVec (CondVec (..), materializeCondVec) -import DataFrame.DecisionTree.Types ( - ColumnOrdering, - SynthConfig (..), - TreeConfig (..), - withOrdFrom, - ) -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Internal.Types -import DataFrame.Operators - -import Data.Either (fromRight) -import Data.Function (on) -import Data.List (inits, sort, sortBy, subsequences) -import qualified Data.Map.Strict as M -import Data.Maybe (fromMaybe, mapMaybe) -import qualified Data.Set as Set -import qualified Data.Text as T -import Data.Type.Equality (testEquality, (:~:) (..)) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import Type.Reflection (typeRep) - --- | Valid-slot view of a nullable boxed column (null slots hold crash-thunks). -validBoxedValues :: Bitmap -> V.Vector a -> V.Vector a -validBoxedValues bm = V.ifilter (\i _ -> bitmapTestBit bm i) - -{- | Target-column summary driving the categorical generator: binary vs -multi-class, the deterministic positive class, and the raw label vector. --} -data TargetInfo target = TargetInfo - { tiIsBinary :: !Bool - , tiPositiveClass :: !(Maybe target) - , tiValues :: !(V.Vector target) - } - -{- | Compute 'TargetInfo' once per fit. The positive class for binary targets -is the lexicographically-first distinct value, for deterministic pools. --} -mkTargetInfo :: - forall target. - (Columnable target, Ord target) => - T.Text -> DataFrame -> Maybe (TargetInfo target) -mkTargetInfo target df = case interpret @target df (Col target) of - Right (TColumn column) -> - either (const Nothing) (Just . targetInfoFromValues) (toVector @target column) - _ -> Nothing - -targetInfoFromValues :: (Ord target) => V.Vector target -> TargetInfo target -targetInfoFromValues vals = TargetInfo isBinary posClass vals - where - distinct = Set.toAscList (Set.fromList (V.toList vals)) - isBinary = length distinct == 2 - posClass = case distinct of - (p : _) | isBinary -> Just p - _ -> Nothing - -{- | Distinct values, capped: @Right vs@ (sorted) under the cap, else @Left@ -the count-so-far so the caller routes to the high-cardinality path. --} -distinctValuesUpTo :: (Ord a) => Int -> V.Vector a -> Either Int [a] -distinctValuesUpTo cap values = go Set.empty 0 - where - n = V.length values - go !s !i - | i >= n = Right (Set.toAscList s) - | Set.size s > cap = Left (Set.size s) - | otherwise = go (Set.insert (V.unsafeIndex values i) s) (i + 1) - -{- | OR-of-equalities for a value-list, shared by the expression and -truth-vector discrete paths so they stay byte-identical. --} -orEqs :: (a -> Expr Bool) -> [a] -> Expr Bool -orEqs eqLit = foldr1 (.||.) . map eqLit - -subsetSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] -subsetSplits eqLit = map (orEqs eqLit) . subsetLists - --- | Proper non-empty, non-full subsets of the values. -subsetLists :: [a] -> [[a]] -subsetLists vs = drop 1 (init (subsequences vs)) - -singletonSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool] -singletonSplits = map - -singletonLists :: [a] -> [[a]] -singletonLists = map (: []) - -breimanPrefixSplits :: - (Ord a, Ord target) => - target -> - V.Vector a -> - V.Vector target -> - [a] -> - (a -> Expr Bool) -> - [Expr Bool] -breimanPrefixSplits pc values targetVals distinctVals eqLit = - map (orEqs eqLit) (breimanPrefixLists pc values targetVals distinctVals) - -{- | Breiman's binary-target split set: sort levels by Laplace-smoothed -positive rate, then take every contiguous non-trivial prefix. --} -breimanPrefixLists :: - (Ord a, Ord target) => target -> V.Vector a -> V.Vector target -> [a] -> [[a]] -breimanPrefixLists pc values targetVals distinctVals = - nonTrivialPrefixes (sortByRate (levelCounts pc values targetVals) distinctVals) - -levelCounts :: - (Ord a, Eq target) => - target -> V.Vector a -> V.Vector target -> M.Map a (Int, Int) -levelCounts pc values targetVals = V.ifoldl' add M.empty values - where - add acc i v = M.insertWith plus v (indicator (V.unsafeIndex targetVals i == pc), 1) acc - plus (p1, n1) (p2, n2) = (p1 + p2, n1 + n2) - indicator b = if b then 1 else 0 - -laplaceRate :: (Ord a) => M.Map a (Int, Int) -> a -> Double -laplaceRate counts v = case M.lookup v counts of - Nothing -> 0.5 - Just (pos, n) -> (fromIntegral pos + 1) / (fromIntegral n + 2) - -sortByRate :: (Ord a) => M.Map a (Int, Int) -> [a] -> [a] -sortByRate counts = sortBy (compare `on` (\v -> (laplaceRate counts v, v))) - -nonTrivialPrefixes :: [a] -> [[a]] -nonTrivialPrefixes = drop 1 . init . inits - -{- | Value-lists a categorical column contributes; shared by the expression and -truth-vector paths so both enumerate identical candidates in the same order. --} -catValueLists :: - (Ord a, Ord target) => - Bool -> Maybe target -> V.Vector target -> Int -> V.Vector a -> [[a]] -catValueLists isBinary posClass targetVals subsetCap values - | V.null values = [] - | isBinary, Just pc <- posClass = binaryLists pc targetVals values - | otherwise = multiclassLists subsetCap values - -binaryLists :: - (Ord a, Ord target) => target -> V.Vector target -> V.Vector a -> [[a]] -binaryLists pc targetVals values - | length distinct < 2 = [] - | otherwise = breimanPrefixLists pc values targetVals distinct - where - distinct = fromRight (ascDistinct values) (distinctValuesUpTo 64 values) - -multiclassLists :: (Ord a) => Int -> V.Vector a -> [[a]] -multiclassLists subsetCap values = case distinctValuesUpTo subsetCap values of - Right vs | length vs >= 2 -> subsetLists vs - Right _ -> [] - Left _ -> singletonLists (ascDistinct values) - -ascDistinct :: (Ord a) => V.Vector a -> [a] -ascDistinct = Set.toAscList . Set.fromList . V.toList - -{- | Truth vector of @col ∈ values@ read directly from the column; equal to -interpreting @orEqs (== v) values@ because the values are distinct. --} -membershipVec :: (Ord a) => V.Vector a -> [a] -> VU.Vector Bool -membershipVec colVals vs = - let !s = Set.fromList vs - in VU.generate (V.length colVals) (\i -> Set.member (colVals `V.unsafeIndex` i) s) - -{- | Per-fit categorical generation context bundling the target summary and -the column-ordering registry. --} -data CatCtx target = CatCtx - { ccBinary :: !Bool - , ccPos :: !(Maybe target) - , ccTargets :: !(V.Vector target) - , ccSubsetCap :: !Int - , ccOrds :: !ColumnOrdering - } - -catCtx :: TargetInfo target -> TreeConfig -> CatCtx target -catCtx ti cfg = - CatCtx - (tiIsBinary ti) - (tiPositiveClass ti) - (tiValues ti) - (maxCategoricalSubsetCardinality (synthConfig cfg)) - (columnOrdering cfg) - -catValueListsFor :: (Ord a, Ord target) => CatCtx target -> V.Vector a -> [[a]] -catValueListsFor ctx = catValueLists (ccBinary ctx) (ccPos ctx) (ccTargets ctx) (ccSubsetCap ctx) - --- | True for numeric columns (handled by the numeric pool, not here). -isNumericKind :: forall a. (Columnable a) => Bool -isNumericKind = case sFloating @a of - STrue -> True - SFalse -> case sIntegral @a of - STrue -> True - SFalse -> False - -{- | All equality-based candidate splits from non-numeric columns: per-column -categorical conditions plus cross-column equality/order conditions. --} -discreteConditions :: - forall target. - (Columnable target, Ord target) => - TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool] -discreteConditions targetInfo cfg df = - concatMap (columnConds (catCtx targetInfo cfg) df) (columnNames df) - ++ crossColumnConds cfg df - -columnConds :: - (Columnable target, Ord target) => - CatCtx target -> DataFrame -> T.Text -> [Expr Bool] -columnConds ctx df colName = case unsafeGetColumn colName df of - BoxedColumn Nothing (column :: V.Vector a) -> nonNullColConds ctx colName column - BoxedColumn (Just bm) (column :: V.Vector a) -> nullableColConds ctx colName bm column - UnboxedColumn _ (_ :: VU.Vector a) -> [] - pt@(PackedText _ _) -> case materializePacked pt of - BoxedColumn Nothing (column :: V.Vector a) -> nonNullColConds ctx colName column - BoxedColumn (Just bm) (column :: V.Vector a) -> nullableColConds ctx colName bm column - _ -> [] - -nonNullColConds :: - forall a target. - (Columnable a, Ord target) => - CatCtx target -> T.Text -> V.Vector a -> [Expr Bool] -nonNullColConds ctx colName column = - fromMaybe - [] - ( withOrdFrom @a - (ccOrds ctx) - (map (orEqs (eqExprFor @a colName)) (catValueListsFor ctx column)) - ) - -nullableColConds :: - forall a target. - (Columnable a, Ord target) => - CatCtx target -> T.Text -> Bitmap -> V.Vector a -> [Expr Bool] -nullableColConds ctx colName bm column - | isNumericKind @a || V.null valid = [] - | otherwise = - fromMaybe - [] - ( withOrdFrom @a - (ccOrds ctx) - (map (orEqs (eqJustFor @a colName)) (catValueListsFor ctx valid)) - ) - where - valid = validBoxedValues bm column - -eqExprFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool -eqExprFor colName v = Col @a colName .==. Lit v - -eqJustFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool -eqJustFor colName v = Col @(Maybe a) colName .==. Lit (Just v) - --- | Cross-column equality/order conditions over pairs of same-typed columns. -crossColumnConds :: TreeConfig -> DataFrame -> [Expr Bool] -crossColumnConds cfg df = concatMap (pairConds (columnOrdering cfg) df) (allowedPairs cfg df) - -allowedPairs :: TreeConfig -> DataFrame -> [(T.Text, T.Text)] -allowedPairs cfg df = - [ (l, r) - | l <- columnNames df - , r <- columnNames df - , l /= r - , not (isDisallowedPair cfg l r) - ] - -isDisallowedPair :: TreeConfig -> T.Text -> T.Text -> Bool -isDisallowedPair cfg l r = - any - (\(l', r') -> sort [l', r'] == sort [l, r]) - (disallowedCombinations (synthConfig cfg)) - -pairConds :: ColumnOrdering -> DataFrame -> (T.Text, T.Text) -> [Expr Bool] -pairConds ords df (l, r) = case ( materializePacked (unsafeGetColumn l df) - , materializePacked (unsafeGetColumn r df) - ) of - (BoxedColumn Nothing (_ :: V.Vector a), BoxedColumn Nothing (_ :: V.Vector b)) -> strictPairConds @a @b l r - (BoxedColumn (Just _) (_ :: V.Vector a), BoxedColumn (Just _) (_ :: V.Vector b)) -> nullablePairConds @a @b ords l r - _ -> [] - -strictPairConds :: - forall a b. (Columnable a, Columnable b) => T.Text -> T.Text -> [Expr Bool] -strictPairConds l r = case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> [Col @a l .==. Col @a r] - Nothing -> [] - -nullablePairConds :: - forall a b. - (Columnable a, Columnable b) => - ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] -nullablePairConds ords l r = case testEquality (typeRep @a) (typeRep @b) of - Nothing -> [] - Just Refl -> nullableEqOrLe @a ords l r - -nullableEqOrLe :: - forall a. (Columnable a) => ColumnOrdering -> T.Text -> T.Text -> [Expr Bool] -nullableEqOrLe ords l r - | isTextType @a = eqOnly - | otherwise = - maybe - eqOnly - (++ eqOnly) - (withOrdFrom @a ords [Col @(Maybe a) l .<=. Col @(Maybe a) r]) - where - eqOnly = [Col @(Maybe a) l .==. Col @(Maybe a) r] - -isTextType :: forall a. (Columnable a) => Bool -isTextType = case testEquality (typeRep @a) (typeRep @T.Text) of - Just Refl -> True - Nothing -> False - -{- | 'discreteConditions' materialized with shared per-column reads: the -non-nullable categorical path builds truth vectors directly from one read -per column; nullable and cross-column fall back to interpret. --} -discreteCondVecs :: - forall target. - (Columnable target, Ord target) => - TargetInfo target -> TreeConfig -> DataFrame -> [CondVec] -discreteCondVecs targetInfo cfg df = - concatMap (columnCondVecs (catCtx targetInfo cfg) df) (columnNames df) - ++ mapMaybe (materializeCondVec df) (crossColumnConds cfg df) - -columnCondVecs :: - (Columnable target, Ord target) => - CatCtx target -> DataFrame -> T.Text -> [CondVec] -columnCondVecs ctx df colName = case unsafeGetColumn colName df of - BoxedColumn Nothing (column :: V.Vector a) -> nonNullColCondVecs ctx colName column - BoxedColumn (Just bm) (column :: V.Vector a) -> mapMaybe (materializeCondVec df) (nullableColConds ctx colName bm column) - UnboxedColumn _ (_ :: VU.Vector a) -> [] - pt@(PackedText _ _) -> case materializePacked pt of - BoxedColumn Nothing (column :: V.Vector a) -> nonNullColCondVecs ctx colName column - BoxedColumn (Just bm) (column :: V.Vector a) -> mapMaybe (materializeCondVec df) (nullableColConds ctx colName bm column) - _ -> [] - -nonNullColCondVecs :: - forall a target. - (Columnable a, Ord target) => CatCtx target -> T.Text -> V.Vector a -> [CondVec] -nonNullColCondVecs ctx colName column = - fromMaybe - [] - ( withOrdFrom @a - (ccOrds ctx) - (map (membershipCondVec colName column) (catValueListsFor ctx column)) - ) - -membershipCondVec :: - forall a. (Columnable a, Ord a) => T.Text -> V.Vector a -> [a] -> CondVec -membershipCondVec colName column vs = CondVec (orEqs (eqExprFor @a colName) vs) (membershipVec column vs) diff --git a/src/DataFrame/DecisionTree/CondVec.hs b/src/DataFrame/DecisionTree/CondVec.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/CondVec.hs +++ /dev/null @@ -1,180 +0,0 @@ -{-# LANGUAGE GADTs #-} -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Cached condition truth vectors and the per-fit cache keyed by structural -form. A condition's truth over a fixed DataFrame is invariant for a whole -fit, so it is materialized once and reused. --} -module DataFrame.DecisionTree.CondVec ( - CondVec (..), - materializeCondVec, - CondCache, - condCacheKey, - condCacheFromVecs, - addTreeCondsToCache, - lookupCondVec, - partitionByVec, - countErrorsByVec, - consolidateThreshold, - combineAndVec, - combineOrVec, -) where - -import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..), Tree (..)) -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column (TypedColumn (..), toVector) -import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression ( - BinaryOp (binaryName), - Expr (..), - eqExpr, - normalize, - ) -import DataFrame.Internal.Interpreter (interpret) - -import qualified Data.Map.Strict as M -import Data.Maybe (fromMaybe) -import qualified Data.Text as T -import Data.Type.Equality (testEquality, (:~:) (..)) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import Type.Reflection (typeRep) - --- | A boolean condition paired with its truth vector over the full DataFrame. -data CondVec = CondVec - { cvExpr :: !(Expr Bool) - , cvVec :: !(VU.Vector Bool) - } - -{- | Interpret a condition once over the DataFrame; 'Nothing' on a -type/interpret failure so the candidate is silently dropped. --} -materializeCondVec :: DataFrame -> Expr Bool -> Maybe CondVec -materializeCondVec df cond = case interpret @Bool df cond of - Left _ -> Nothing - Right (TColumn column) -> CondVec cond <$> eitherToMaybe (toVector @Bool @VU.Vector column) - -eitherToMaybe :: Either e a -> Maybe a -eitherToMaybe = either (const Nothing) Just - -{- | Full-DataFrame truth vectors keyed by structural form, read-only once -built. Seeded for free from the candidate pool plus the initial tree so the -predict/loss passes index a vector instead of re-interpreting per node. --} -type CondCache = M.Map T.Text (VU.Vector Bool) - -{- | Structural key matching the candidate-dedup key, so a tree branch whose -condition came from the pool hits the cache (equal keys ⟹ equal vector). --} -condCacheKey :: Expr Bool -> T.Text -condCacheKey = T.pack . show . normalize - --- | Seed a cache from already-materialized candidate 'CondVec's (no interpret). -condCacheFromVecs :: [CondVec] -> CondCache -condCacheFromVecs cvs = M.fromList [(condCacheKey (cvExpr cv), cvVec cv) | cv <- cvs] - -{- | Add a tree's branch-condition vectors to a cache (one interpret per -distinct, not-yet-cached condition). --} -addTreeCondsToCache :: DataFrame -> Tree a -> CondCache -> CondCache -addTreeCondsToCache df = go - where - go (Leaf _) c = c - go (Branch cond l r) c = go r (go l (insertCond df cond c)) - -insertCond :: DataFrame -> Expr Bool -> CondCache -> CondCache -insertCond df cond c - | M.member k c = c - | otherwise = - maybe c (\cv -> M.insert k (cvVec cv) c) (materializeCondVec df cond) - where - k = condCacheKey cond - -{- | A condition's truth vector: a cache hit, else interpret over the -DataFrame. 'Nothing' mirrors the interpret-failure fallback (route left). --} -lookupCondVec :: CondCache -> DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) -lookupCondVec cache df cond = case M.lookup (condCacheKey cond) cache of - hit@(Just _) -> hit - Nothing -> cvVec <$> materializeCondVec df cond - --- | Partition row indices by a truth vector: @True@ → left, @False@ → right. -partitionByVec :: VU.Vector Bool -> V.Vector Int -> (V.Vector Int, V.Vector Int) -partitionByVec boolVals = V.partition (boolVals VU.!) - --- | Count care points the truth vector routes to the wrong child. -countErrorsByVec :: VU.Vector Bool -> [CarePoint] -> Int -countErrorsByVec boolVals = length . filter misrouted - where - misrouted cp = (boolVals VU.! cpIndex cp) /= (cpCorrectDir cp == GoLeft) - -{- | A same-column same-direction Double threshold comparison, with a rebuild -function to swap in a new threshold. --} -data ThreshCmp = ThreshCmp - { tcCol :: !T.Text - , tcName :: !T.Text - , tcThr :: !Double - , tcRebuild :: Double -> Expr Bool - } - -asDoubleThreshold :: Expr Bool -> Maybe ThreshCmp -asDoubleThreshold (Binary op (Col c :: Expr cc) (Lit (t :: tt))) = - case ( testEquality (typeRep @cc) (typeRep @Double) - , testEquality (typeRep @tt) (typeRep @Double) - ) of - (Just Refl, Just Refl) -> Just (ThreshCmp c (binaryName op) t (Binary op (Col c) . Lit)) - _ -> Nothing -asDoubleThreshold _ = Nothing - -directionalNames :: [T.Text] -directionalNames = ["lt", "leq", "gt", "geq"] - -{- | Tighter (AND) or looser (OR) of two same-direction thresholds: @<@/@<=@ -are left-half-spaces (AND = min), @>@/@>=@ are right-half-spaces (AND = max). --} -chooseThreshold :: Bool -> T.Text -> Double -> Double -> Double -chooseThreshold isAnd name t1 t2 - | leftDir = if isAnd then min t1 t2 else max t1 t2 - | otherwise = if isAnd then max t1 t2 else min t1 t2 - where - leftDir = name == "lt" || name == "leq" - -{- | Collapse two same-column same-direction strict-Double comparisons into one -comparison (the @True@ argument selects AND, @False@ OR); 'Nothing' otherwise. --} -consolidateThreshold :: Bool -> Expr Bool -> Expr Bool -> Maybe (Expr Bool) -consolidateThreshold isAnd ea eb = do - a <- asDoubleThreshold ea - b <- asDoubleThreshold eb - if tcCol a == tcCol b && tcName a == tcName b && tcName a `elem` directionalNames - then Just (tcRebuild a (chooseThreshold isAnd (tcName a) (tcThr a) (tcThr b))) - else Nothing - -{- | AND-combine two cached conditions: idempotence and threshold consolidation -first, else the generic @F.and@; the vector is always the elementwise AND. --} -combineAndVec :: CondVec -> CondVec -> CondVec -combineAndVec a b - | eqExpr (cvExpr a) (cvExpr b) = a - | otherwise = CondVec expr (VU.zipWith (&&) (cvVec a) (cvVec b)) - where - expr = - fromMaybe - (F.and (cvExpr a) (cvExpr b)) - (consolidateThreshold True (cvExpr a) (cvExpr b)) - -{- | OR-combine two cached conditions (see 'combineAndVec'; AND/OR direction -differs in 'consolidateThreshold'). --} -combineOrVec :: CondVec -> CondVec -> CondVec -combineOrVec a b - | eqExpr (cvExpr a) (cvExpr b) = a - | otherwise = CondVec expr (VU.zipWith (||) (cvVec a) (cvVec b)) - where - expr = - fromMaybe - (F.or (cvExpr a) (cvExpr b)) - (consolidateThreshold False (cvExpr a) (cvExpr b)) diff --git a/src/DataFrame/DecisionTree/Fit.hs b/src/DataFrame/DecisionTree/Fit.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Fit.hs +++ /dev/null @@ -1,213 +0,0 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Top-level fitting: assemble the candidate pool, seed from CART, run TAO, -and convert the result to an expression. Also the probability-tree variant -('fitProbTree') that annotates leaves with class distributions. --} -module DataFrame.DecisionTree.Fit ( - treeToExpr, - fitDecisionTree, - buildTree, - pruneTree, - partitionDataFrame, - calculateGini, - majorityValue, - getCounts, - percentile, - ProbTree, - probsFromIndices, - buildProbTree, - fitProbTree, - probExprs, -) where - -import DataFrame.DecisionTree.Cart (buildCartTree) -import DataFrame.DecisionTree.Categorical ( - TargetInfo (..), - discreteCondVecs, - discreteConditions, - mkTargetInfo, - ) -import DataFrame.DecisionTree.CondVec (CondVec) -import DataFrame.DecisionTree.Numeric (numericCondVecs, numericConditions) -import DataFrame.DecisionTree.Pool (dedupCVByExpr, nubByExpr) -import DataFrame.DecisionTree.Predict (partitionIndices) -import DataFrame.DecisionTree.Prune (pruneDead, pruneExpr) -import DataFrame.DecisionTree.Tao (taoOptimize, taoOptimizeCV) -import DataFrame.DecisionTree.Types (Tree (..), TreeConfig (..)) -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) -import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Operations.Core (nRows) -import DataFrame.Operations.Subset (exclude, filterWhere) - -import Control.Exception (throw) -import Data.Function (on) -import Data.List (foldl', maximumBy, nub, sort) -import qualified Data.Map.Strict as M -import Data.Maybe (fromMaybe) -import qualified Data.Text as T -import qualified Data.Vector as V - --- | Convert a fitted tree to a nested-conditional expression. -treeToExpr :: (Columnable a) => Tree a -> Expr a -treeToExpr (Leaf v) = Lit v -treeToExpr (Branch cond left right) = F.ifThenElse cond (treeToExpr left) (treeToExpr right) - --- | Fit a TAO decision tree (CART-seeded) and return it as an expression. -fitDecisionTree :: - forall a. (Columnable a, Ord a) => TreeConfig -> Expr a -> DataFrame -> Expr a -fitDecisionTree cfg (Col target) df = - pruneExpr - (treeToExpr (taoOptimizeCV @a cfg target condVecs df indices initialTree)) - where - condVecs = candidatePool @a cfg target df - initialTree = buildCartTree @a cfg target df - indices = V.enumFromN 0 (nRows df) -fitDecisionTree _ expr _ = error ("Cannot create tree for compound expression: " ++ show expr) - --- | The deduplicated numeric + discrete candidate pool for a target column. -candidatePool :: - forall a. - (Columnable a, Ord a) => TreeConfig -> T.Text -> DataFrame -> [CondVec] -candidatePool cfg target df = dedupCVByExpr (numericCVs ++ discreteCVs) - where - dfNoTarget = exclude [target] df - numericCVs = numericCondVecs cfg dfNoTarget df - discreteCVs = discreteCondVecs (targetInfoOrEmpty @a target df) cfg dfNoTarget - -targetInfoOrEmpty :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> TargetInfo a -targetInfoOrEmpty target df = fromMaybe (TargetInfo False Nothing V.empty) (mkTargetInfo @a target df) - --- | Fit a tree at a given depth from a raw condition list (CART + TAO + prune). -buildTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> Int -> T.Text -> [Expr Bool] -> DataFrame -> Expr a -buildTree cfg depth target conds df = - pruneExpr (treeToExpr (taoOptimize @a cfg target conds df indices tree)) - where - tree = buildCartTree @a cfg{maxTreeDepth = depth} target df - indices = V.enumFromN 0 (nRows df) - -pruneTree :: forall a. (Columnable a) => Expr a -> Expr a -pruneTree = pruneExpr - -partitionDataFrame :: Expr Bool -> DataFrame -> (DataFrame, DataFrame) -partitionDataFrame cond df = (filterWhere cond df, filterWhere (F.not cond) df) - --- | Laplace-smoothed Gini impurity of the target distribution. -calculateGini :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> Double -calculateGini target df - | n == 0 = 0 - | otherwise = 1 - sum (map (^ (2 :: Int)) probs) - where - n = fromIntegral (nRows df) - counts = getCounts @a target df - numClasses = fromIntegral (M.size counts) - probs = map (\c -> (fromIntegral c + 1) / (n + numClasses)) (M.elems counts) - -majorityValue :: forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> a -majorityValue target df - | M.null counts = error "Empty DataFrame in leaf" - | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) - where - counts = getCounts @a target df - -getCounts :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> M.Map a Int -getCounts target df = case interpret @a df (Col target) of - Left e -> throw e - Right (TColumn column) -> case toVector @a column of - Left e -> throw e - Right vals -> foldl' (\acc x -> M.insertWith (+) x 1 acc) M.empty (V.toList vals) - --- | The @p@-th percentile of an expression's values (@0@ on failure/empty). -percentile :: Int -> Expr Double -> DataFrame -> Double -percentile p expr df = case interpret @Double df expr of - Right (TColumn column) -> either (const 0) (percentileOfVec p) (toVector @Double column) - _ -> 0 - -percentileOfVec :: Int -> V.Vector Double -> Double -percentileOfVec p vals - | n == 0 = 0 - | otherwise = sorted V.! min (n - 1) (max 0 ((p * n) `div` 100)) - where - sorted = V.fromList (sort (V.toList vals)) - n = V.length sorted - --- | A tree whose leaves hold class-probability distributions. -type ProbTree a = Tree (M.Map a Double) - --- | Normalised class probabilities over a subset of training rows. -probsFromIndices :: - forall a. - (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> M.Map a Double -probsFromIndices target df indices = case interpret @a df (Col target) of - Right (TColumn column) -> either (const M.empty) (normaliseCounts indices) (toVector @a column) - _ -> M.empty - -normaliseCounts :: (Ord a) => V.Vector Int -> V.Vector a -> M.Map a Double -normaliseCounts indices vals = M.map (\c -> fromIntegral c / total) counts - where - counts = - V.foldl' - (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) - M.empty - indices - total = fromIntegral (V.length indices) :: Double - -{- | Re-label a fitted tree's leaves with class distributions, routing the -training data through the (unchanged) split conditions. --} -buildProbTree :: - forall a. - (Columnable a, Ord a) => - Tree a -> T.Text -> DataFrame -> V.Vector Int -> ProbTree a -buildProbTree (Leaf _) target df indices = Leaf (probsFromIndices @a target df indices) -buildProbTree (Branch cond left right) target df indices = - Branch - cond - (buildProbTree @a left target df l) - (buildProbTree @a right target df r) - where - (l, r) = partitionIndices cond df indices - --- | Fit a TAO tree and return one probability expression per class. -fitProbTree :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> Expr a -> DataFrame -> M.Map a (Expr Double) -fitProbTree cfg (Col target) df = probExprs (buildProbTree @a pruned target df indices) - where - conds = - nubByExpr - ( numericConditions cfg dfNoTarget - ++ discreteConditions (targetInfoOrEmpty @a target df) cfg dfNoTarget - ) - dfNoTarget = exclude [target] df - indices = V.enumFromN 0 (nRows df) - pruned = - pruneDead - (taoOptimize @a cfg target conds df indices (buildCartTree @a cfg target df)) -fitProbTree _ expr _ = error ("Cannot create prob tree for compound expression: " ++ show expr) - --- | Convert a 'ProbTree' into one @Expr Double@ per class. -probExprs :: - forall a. (Columnable a, Ord a) => ProbTree a -> M.Map a (Expr Double) -probExprs tree = M.fromList [(c, classExpr c tree) | c <- nub (allClasses tree)] - -allClasses :: ProbTree a -> [a] -allClasses (Leaf m) = M.keys m -allClasses (Branch _ l r) = allClasses l ++ allClasses r - -classExpr :: (Ord a) => a -> ProbTree a -> Expr Double -classExpr c (Leaf m) = Lit (M.findWithDefault 0.0 c m) -classExpr c (Branch cond l r) = F.ifThenElse cond (classExpr c l) (classExpr c r) diff --git a/src/DataFrame/DecisionTree/Linear.hs b/src/DataFrame/DecisionTree/Linear.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Linear.hs +++ /dev/null @@ -1,161 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} -{-# LANGUAGE TypeApplications #-} - -{- | Oblique split candidates: fit an L1-regularised logistic hyperplane to the -care points (class-balanced) and convert it to a boolean condition, rejecting -all-zero and degenerate (single-side) hyperplanes. --} -module DataFrame.DecisionTree.Linear ( - bestLinearCandidate, - fitLinearCandidate, - careRowsFromFeatures, - careLabels, - featName, - imputeMean, - materializeFeatureForCare, -) where - -import DataFrame.DecisionTree.Numeric (NumExpr (..), numericCols) -import DataFrame.DecisionTree.Types ( - CarePoint (..), - Direction (..), - TreeConfig (..), - ) -import DataFrame.Internal.Column (TypedColumn (..), toVector) -import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr, getColumns) -import DataFrame.Internal.Interpreter (interpret) -import qualified DataFrame.LinearSolver as LS - -import Data.Maybe (catMaybes, fromMaybe, mapMaybe) -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -{- | Best oblique candidate, or 'Nothing' when the linear path is disabled or -there are too few care points to fit on. --} -bestLinearCandidate :: - TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) -bestLinearCandidate cfg df carePoints - | not (useLinearSolver cfg) = Nothing - | length carePoints < minCarePointsForLinear cfg = Nothing - | otherwise = fitLinearCandidate cfg df carePoints - -{- | Fit an L1 logistic regression to the care points and convert the resulting -hyperplane to a condition, or 'Nothing' when no numeric features exist or the -fitted model is all-zero or degenerate. --} -fitLinearCandidate :: - TreeConfig -> DataFrame -> [CarePoint] -> Maybe (Expr Bool) -fitLinearCandidate cfg df carePoints = case materializedFeatures df carePoints of - [] -> Nothing - mats -> linearFromFeatures cfg carePoints mats - -materializedFeatures :: DataFrame -> [CarePoint] -> [(T.Text, VU.Vector Double)] -materializedFeatures df carePoints = mapMaybe (materializeFeatureForCare df carePoints) (numericCols df) - -linearFromFeatures :: - TreeConfig -> [CarePoint] -> [(T.Text, VU.Vector Double)] -> Maybe (Expr Bool) -linearFromFeatures cfg carePoints mats - | VU.all (== 0) weights = Nothing - | degenerateHyperplane rows weights (LS.lmIntercept model) = Nothing - | otherwise = Just (LS.modelToExpr model) - where - rows = careRowsFromFeatures (length carePoints) mats - labels = careLabels carePoints - model = - LS.fitL1Logistic - (solverConfigFor cfg labels) - rows - labels - (V.fromList (map fst mats)) - weights = LS.lmWeights model - -solverConfigFor :: TreeConfig -> VU.Vector Double -> LS.SolverConfig -solverConfigFor cfg labels = (linearSolverConfig cfg){LS.scSampleWeights = classBalancedWeights labels} - -{- | Class-balanced sklearn-form weights @w_i = N / (2 · N_class)@ (mean 1), or -'Nothing' in the degenerate one-class case (uniform weighting). --} -classBalancedWeights :: VU.Vector Double -> Maybe (VU.Vector Double) -classBalancedWeights labels - | nPos > 0 && nNeg > 0 = Just (VU.generate nCare weightAt) - | otherwise = Nothing - where - nCare = VU.length labels - nPos = VU.length (VU.filter (> 0) labels) - nNeg = nCare - nPos - weightAt i - | VU.unsafeIndex labels i > 0 = fromIntegral nCare / (2 * fromIntegral nPos) - | otherwise = fromIntegral nCare / (2 * fromIntegral nNeg) - -{- | A hyperplane is degenerate when every care row scores on the same side of -zero (equivalent to an invalid split, caught upstream). --} -degenerateHyperplane :: - V.Vector (VU.Vector Double) -> VU.Vector Double -> Double -> Bool -degenerateHyperplane rows weights bias = - nCare > 0 && (VU.minimum scores > 0 || VU.maximum scores < 0) - where - nCare = V.length rows - scores = - VU.generate - nCare - (\i -> VU.sum (VU.zipWith (*) weights (V.unsafeIndex rows i)) + bias) - -{- | Per-care-point feature rows from materialized columns (each of length -@nCare@, so indexing is in range). --} -careRowsFromFeatures :: - Int -> [(T.Text, VU.Vector Double)] -> V.Vector (VU.Vector Double) -careRowsFromFeatures nCare mats = - V.generate nCare (\i -> VU.generate nFeat (\j -> snd (matsVec V.! j) VU.! i)) - where - matsVec = V.fromList mats - nFeat = V.length matsVec - --- | Solver labels: @+1@ when 'GoLeft' is correct, @-1@ otherwise. -careLabels :: [CarePoint] -> VU.Vector Double -careLabels carePoints = - VU.fromList [if cpCorrectDir cp == GoLeft then 1.0 else -1.0 | cp <- carePoints] - --- | First column referenced by an expression, or a placeholder when none. -featName :: Expr b -> T.Text -featName expr = case getColumns expr of - (c : _) -> c - [] -> "" - -{- | Replace missing values with the mean of present ones; 'Nothing' when -nothing is present so the caller can drop the feature. --} -imputeMean :: [Maybe Double] -> Maybe (VU.Vector Double) -imputeMean careRaw = case catMaybes careRaw of - [] -> Nothing - present -> Just (VU.fromList [fromMaybe (mean present) mv | mv <- careRaw]) - where - mean xs = sum xs / fromIntegral (length xs) - -interpretDoubleVals :: DataFrame -> Expr Double -> Maybe (V.Vector Double) -interpretDoubleVals df expr = case interpret @Double df expr of - Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) - _ -> Nothing - -interpretMaybeDoubleVals :: - DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) -interpretMaybeDoubleVals df expr = case interpret @(Maybe Double) df expr of - Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) - _ -> Nothing - -{- | Materialize a 'NumExpr' over the care rows; 'Nothing' on interpret failure -or (nullable) when no care point has a present value, else mean-imputed. --} -materializeFeatureForCare :: - DataFrame -> [CarePoint] -> NumExpr -> Maybe (T.Text, VU.Vector Double) -materializeFeatureForCare df carePoints (NDouble expr) = do - vals <- interpretDoubleVals df expr - Just (featName expr, VU.fromList [vals V.! cpIndex cp | cp <- carePoints]) -materializeFeatureForCare df carePoints (NMaybeDouble expr) = do - vals <- interpretMaybeDoubleVals df expr - imputed <- imputeMean [vals V.! cpIndex cp | cp <- carePoints] - Just (featName expr, imputed) diff --git a/src/DataFrame/DecisionTree/Model.hs b/src/DataFrame/DecisionTree/Model.hs --- a/src/DataFrame/DecisionTree/Model.hs +++ b/src/DataFrame/DecisionTree/Model.hs @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} {- | sklearn-style standalone tree estimators returning inspectable records @@ -11,6 +12,7 @@ the classifier @Expr@. -} module DataFrame.DecisionTree.Model ( + module DataFrame.Model, DecisionTreeClassifier (..), DecisionTreeRegressor (..), ) where @@ -27,7 +29,7 @@ import DataFrame.Featurize.Internal (targetDoubles) import DataFrame.Internal.Column (Columnable) import DataFrame.Internal.Expression (Expr (..), getColumns) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model -- | A fitted classification tree with structural diagnostics. data DecisionTreeClassifier a = DecisionTreeClassifier @@ -48,7 +50,8 @@ } deriving (Show) -instance (Columnable a, Ord a) => Fit TreeConfig (Expr a) (DecisionTreeClassifier a) where +instance (Columnable a, Ord a) => Fit TreeConfig (Expr a) where + type ModelOf TreeConfig (Expr a) = (DecisionTreeClassifier a) fit cfg target df = DecisionTreeClassifier e @@ -58,10 +61,12 @@ where e = fitDecisionTree cfg target df -instance Predict (DecisionTreeClassifier a) a where +instance Predict (DecisionTreeClassifier a) where + type Prediction (DecisionTreeClassifier a) = Expr a predict = dtcExpr -instance Fit RegTreeConfig (Expr Double) DecisionTreeRegressor where +instance Fit RegTreeConfig (Expr Double) where + type ModelOf RegTreeConfig (Expr Double) = DecisionTreeRegressor fit cfg target df = DecisionTreeRegressor t @@ -82,7 +87,8 @@ ("fit @DecisionTreeRegressor: target must be a column, got " ++ show target) e = treeToExpr t -instance Predict DecisionTreeRegressor Double where +instance Predict DecisionTreeRegressor where + type Prediction DecisionTreeRegressor = Expr Double predict = dtrExpr usageCounts :: [T.Text] -> M.Map T.Text Int diff --git a/src/DataFrame/DecisionTree/Numeric.hs b/src/DataFrame/DecisionTree/Numeric.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Numeric.hs +++ /dev/null @@ -1,256 +0,0 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Numeric split candidates: per-column Double expressions, arithmetic -expansion, and threshold conditions. 'numericCondVecs' materializes the -pool with a single interpret per distinct expression, deriving every -threshold/operator truth vector by direct comparison. --} -module DataFrame.DecisionTree.Numeric ( - NumExpr (..), - numExprCols, - numExprEq, - combineNumExprs, - numericConditions, - generateNumericConds, - percentilesOf, - numericCondVecs, - numericExprsWithTerms, - numericCols, -) where - -import DataFrame.DecisionTree.CondVec (CondVec (..)) -import DataFrame.DecisionTree.Types (SynthConfig (..), TreeConfig (..)) -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column -import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn) -import DataFrame.Internal.Expression (Expr (..), eqExpr, getColumns, normalize) -import DataFrame.Internal.Interpreter (interpret) -import DataFrame.Internal.Types -import DataFrame.Operators - -import Data.List (sort) -import Data.Maybe (fromMaybe) -import qualified Data.Set as Set -import qualified Data.Text as T -import Data.Type.Equality (testEquality, (:~:) (..)) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import Type.Reflection (typeRep) - --- | A numeric feature expression, non-nullable or nullable. -data NumExpr - = NDouble !(Expr Double) - | NMaybeDouble !(Expr (Maybe Double)) - -numExprCols :: NumExpr -> [T.Text] -numExprCols (NDouble e) = getColumns e -numExprCols (NMaybeDouble e) = getColumns e - -numExprEq :: NumExpr -> NumExpr -> Bool -numExprEq (NDouble e1) (NDouble e2) = eqExpr e1 e2 -numExprEq (NMaybeDouble e1) (NMaybeDouble e2) = eqExpr e1 e2 -numExprEq _ _ = False - --- | Safe division: @0@ (or @Nothing@) where the divisor is zero. -safeDivD :: Expr Double -> Expr Double -> Expr Double -safeDivD a b = F.ifThenElse (b ./= F.lit (0 :: Double)) (a ./ b) (F.lit (0 :: Double)) - -safeDivMaybe :: Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double) -safeDivMaybe nonZero q = F.ifThenElse nonZero q (F.lit (Nothing :: Maybe Double)) - --- | Arithmetic combinations (@+@, @-@, @*@, safe @/@) of two numeric exprs. -combineNumExprs :: NumExpr -> NumExpr -> [NumExpr] -combineNumExprs (NDouble e1) (NDouble e2) = - map NDouble [e1 .+ e2, e1 .- e2, e1 .* e2, safeDivD e1 e2] -combineNumExprs (NDouble e1) (NMaybeDouble e2) = - map - NMaybeDouble - [ e1 .+ e2 - , e1 .- e2 - , e1 .* e2 - , safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2) - ] -combineNumExprs (NMaybeDouble e1) (NDouble e2) = - map - NMaybeDouble - [ e1 .+ e2 - , e1 .- e2 - , e1 .* e2 - , safeDivMaybe (e2 ./= F.lit (0 :: Double)) (e1 ./ e2) - ] -combineNumExprs (NMaybeDouble e1) (NMaybeDouble e2) = - map - NMaybeDouble - [ e1 .+ e2 - , e1 .- e2 - , e1 .* e2 - , safeDivMaybe (F.fromMaybe False (e2 ./= F.lit (0 :: Double))) (e1 ./ e2) - ] - -numericConditions :: TreeConfig -> DataFrame -> [Expr Bool] -numericConditions = generateNumericConds - -generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool] -generateNumericConds cfg df = do - expr <- numericExprsWithTerms (synthConfig cfg) df - threshold <- numericThresholds cfg df expr - condsFromExpr expr threshold - -numericThresholds :: TreeConfig -> DataFrame -> NumExpr -> [Double] -numericThresholds cfg df (NDouble e) = thresholdsForExpr cfg df e -numericThresholds cfg df (NMaybeDouble e) = thresholdsForExpr cfg df (F.fromMaybe 0 e) - -thresholdsForExpr :: TreeConfig -> DataFrame -> Expr Double -> [Double] -thresholdsForExpr cfg df e = - maybe [] (percentilesOf (percentiles cfg) . V.toList) (interpretDoubleCol df e) - -condsFromExpr :: NumExpr -> Double -> [Expr Bool] -condsFromExpr (NDouble e) t = [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] -condsFromExpr (NMaybeDouble e) t = - map - (F.fromMaybe False) - [e .<= F.lit t, e .>= F.lit t, e .< F.lit t, e .> F.lit t] - -{- | Percentile thresholds for a value list: sort once, index each percentile. -Shared by 'generateNumericConds' and 'numericCondVecs' for identical results. --} -percentilesOf :: [Int] -> [Double] -> [Double] -percentilesOf ps valsList - | n == 0 = [] - | otherwise = map (\p -> sortedV V.! min (n - 1) (max 0 (p * n `div` 100))) ps - where - !sortedV = V.fromList (sort valsList) - !n = V.length sortedV - -interpretDoubleCol :: DataFrame -> Expr Double -> Maybe (V.Vector Double) -interpretDoubleCol df e = case interpret @Double df e of - Right (TColumn column) -> either (const Nothing) Just (toVector @Double column) - _ -> Nothing - -interpretMaybeDoubleCol :: - DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double)) -interpretMaybeDoubleCol df e = case interpret @(Maybe Double) df e of - Right (TColumn column) -> either (const Nothing) Just (toVector @(Maybe Double) column) - _ -> Nothing - -{- | Materialize the numeric pool with one interpret per distinct expression, -deriving each threshold/operator truth vector by direct comparison. -Byte-identical to materializing 'numericConditions' one at a time, but -avoids re-interpreting each LHS per threshold and operator. --} -numericCondVecs :: TreeConfig -> DataFrame -> DataFrame -> [CondVec] -numericCondVecs cfg dfGen df = concatMap forExpr (numericExprsWithTerms (synthConfig cfg) dfGen) - where - forExpr (NDouble e) = maybe [] (condsForDouble cfg e) (interpretDoubleCol df e) - forExpr (NMaybeDouble e) = maybe [] (condsForMaybe cfg e) (interpretMaybeDoubleCol df e) - -condsForDouble :: TreeConfig -> Expr Double -> V.Vector Double -> [CondVec] -condsForDouble cfg e vals = concatMap (doubleCondsAt e vals (V.length vals)) ts - where - ts = percentilesOf (percentiles cfg) (V.toList vals) - -doubleCondsAt :: Expr Double -> V.Vector Double -> Int -> Double -> [CondVec] -doubleCondsAt e vals n t = - [ CondVec (e .<= F.lit t) (gen (<= t)) - , CondVec (e .>= F.lit t) (gen (>= t)) - , CondVec (e .< F.lit t) (gen (< t)) - , CondVec (e .> F.lit t) (gen (> t)) - ] - where - gen p = VU.generate n (\i -> p (vals V.! i)) - -condsForMaybe :: - TreeConfig -> Expr (Maybe Double) -> V.Vector (Maybe Double) -> [CondVec] -condsForMaybe cfg e mvals = concatMap (maybeCondsAt e mvals (V.length mvals)) ts - where - ts = percentilesOf (percentiles cfg) (map (fromMaybe 0) (V.toList mvals)) - -maybeCondsAt :: - Expr (Maybe Double) -> V.Vector (Maybe Double) -> Int -> Double -> [CondVec] -maybeCondsAt e mvals n t = - [ CondVec (F.fromMaybe False (e .<= F.lit t)) (gen (<= t)) - , CondVec (F.fromMaybe False (e .>= F.lit t)) (gen (>= t)) - , CondVec (F.fromMaybe False (e .< F.lit t)) (gen (< t)) - , CondVec (F.fromMaybe False (e .> F.lit t)) (gen (> t)) - ] - where - gen p = VU.generate n (\i -> maybe False p (mvals V.! i)) - -{- | Arithmetic candidate expansion, generated already-deduped: each round -combines @frontier × base@ and admits only normalized-novel candidates. -Produces @base@ plus @maxExprDepth-1@ combination rounds. --} -numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr] -numericExprsWithTerms cfg df - | not (enableArithOps cfg) = base - | otherwise = - base ++ expandRounds cfg base (max 0 (maxExprDepth cfg - 1)) base seen0 - where - base = numericCols df - seen0 = Set.fromList (map keyNum base) - -keyNum :: NumExpr -> String -keyNum (NDouble e) = show (normalize e) -keyNum (NMaybeDouble e) = show (normalize e) - -isDisallowed :: SynthConfig -> NumExpr -> NumExpr -> Bool -isDisallowed cfg e1 e2 = - any (\(l, r) -> l `elem` cols && r `elem` cols) (disallowedCombinations cfg) - where - cols = numExprCols e1 <> numExprCols e2 - -roundProducts :: SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr] -roundProducts cfg frontier base = - [ c - | e1 <- frontier - , e2 <- base - , not (numExprEq e1 e2) - , not (isDisallowed cfg e1 e2) - , c <- combineNumExprs e1 e2 - ] - -expandRounds :: - SynthConfig -> [NumExpr] -> Int -> [NumExpr] -> Set.Set String -> [NumExpr] -expandRounds _ _ 0 _ _ = [] -expandRounds cfg base d frontier seen - | null admitted = [] - | otherwise = admitted ++ expandRounds cfg base (d - 1) admitted seen' - where - (admitted, seen') = admitNovel seen (roundProducts cfg frontier base) - -admitNovel :: Set.Set String -> [NumExpr] -> ([NumExpr], Set.Set String) -admitNovel seen0 = go seen0 [] - where - go seen acc [] = (reverse acc, seen) - go seen acc (c : cs) - | keyNum c `Set.member` seen = go seen acc cs - | otherwise = go (Set.insert (keyNum c) seen) (c : acc) cs - -numericCols :: DataFrame -> [NumExpr] -numericCols df = concatMap (numExprsOfColumn df) (columnNames df) - -numExprsOfColumn :: DataFrame -> T.Text -> [NumExpr] -numExprsOfColumn df colName = case unsafeGetColumn colName df of - UnboxedColumn Nothing (_ :: VU.Vector b) -> strictNumeric @b colName - BoxedColumn (Just _) (_ :: V.Vector b) -> nullableNumeric @b colName - UnboxedColumn (Just _) (_ :: VU.Vector b) -> nullableNumeric @b colName - _ -> [] - -strictNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] -strictNumeric c = case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NDouble (Col c)] - Nothing -> case sIntegral @b of - STrue -> [NDouble (F.toDouble (Col @b c))] - SFalse -> [] - -nullableNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr] -nullableNumeric c = case testEquality (typeRep @b) (typeRep @Double) of - Just Refl -> [NMaybeDouble (Col @(Maybe b) c)] - Nothing -> case sIntegral @b of - STrue -> [NMaybeDouble (F.whenPresent (realToFrac @b @Double) (Col @(Maybe b) c))] - SFalse -> [] diff --git a/src/DataFrame/DecisionTree/Pool.hs b/src/DataFrame/DecisionTree/Pool.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Pool.hs +++ /dev/null @@ -1,224 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE OverloadedStrings #-} - -{- | Candidate-pool scoring and boolean expansion: penalized scoring, diverse -top-K selection, AND/OR saturation, and structural/truth-vector dedup. The -per-node scoring scans run in parallel chunks. --} -module DataFrame.DecisionTree.Pool ( - evalWithPenaltyVec, - primaryColExpr, - primaryColCV, - takeDiverse, - candidateParChunk, - bestDiscreteCandidate, - boolExprsVec, - DedupMode (..), - saturateCandidates, - roundProducts, - admitKeys, - admitVecs, - dedupCVByExpr, - nubByExpr, -) where - -import DataFrame.DecisionTree.CondVec ( - CondVec (..), - combineAndVec, - combineOrVec, - countErrorsByVec, - ) -import DataFrame.DecisionTree.Types ( - CarePoint, - SynthConfig (..), - TreeConfig (..), - ) -import DataFrame.Internal.Expression ( - Expr, - compareExpr, - eSize, - eqExpr, - getColumns, - normalize, - ) - -import Control.Parallel.Strategies (parListChunk, rdeepseq, using) -import Data.Function (on) -import Data.List (minimumBy, sortBy) -import qualified Data.Map.Strict as M -import qualified Data.Set as Set -import qualified Data.Text as T -import qualified Data.Vector.Unboxed as VU - -{- | Penalized score of a candidate: care-point errors plus a complexity -penalty, tie-broken by expression size. --} -evalWithPenaltyVec :: TreeConfig -> [CarePoint] -> CondVec -> (Int, Int) -evalWithPenaltyVec cfg carePoints cv = (countErrorsByVec (cvVec cv) carePoints + penalty, sz) - where - sz = eSize (cvExpr cv) - penalty = floor (complexityPenalty (synthConfig cfg) * fromIntegral sz) - -{- | First referenced column of a condition (a sentinel for literal-only ones), -used by 'takeDiverse' to enforce per-column diversity. --} -primaryColExpr :: Expr Bool -> T.Text -primaryColExpr e = case getColumns e of - [] -> "" - (c : _) -> c - -primaryColCV :: CondVec -> T.Text -primaryColCV = primaryColExpr . cvExpr - -{- | Keep the first @k@ of an already-sorted list, admitting at most @quota@ per -primary column (@Nothing@ disables the per-column cap). --} -takeDiverse :: Int -> Maybe Int -> (a -> T.Text) -> [a] -> [a] -takeDiverse k Nothing _ = take k -takeDiverse k (Just quota) primary = go M.empty 0 - where - go !_ !_ [] = [] - go !seen !n (x : xs) - | n >= k = [] - | M.findWithDefault 0 col seen >= quota = go seen n xs - | otherwise = x : go (M.insertWith (+) col 1 seen) (n + 1) xs - where - !col = primary x - -{- | Chunk size for the parallel per-node candidate scans; tuned by an -N -sweep, not correctness-affecting. --} -candidateParChunk :: Int -candidateParChunk = 64 - -{- | Decorate candidates with their penalty in parallel chunks, forcing only -the @(Int, Int)@ key so the order (hence later sorts/minima) is preserved. --} -decorate :: (CondVec -> (Int, Int)) -> [CondVec] -> [((Int, Int), CondVec)] -decorate penaltyCV xs = zip (map penaltyCV xs `using` parListChunk candidateParChunk rdeepseq) xs - --- | The diverse top-@expressionPairs@ valid candidates by penalty. -sortedTopK :: TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> [CondVec] -sortedTopK cfg penaltyCV validCondVecs = - map - snd - ( takeDiverse - (expressionPairs cfg) - (perColumnQuota (synthConfig cfg)) - (primaryColCV . snd) - sorted - ) - where - sorted = sortBy (compare `on` fst) (decorate penaltyCV validCondVecs) - --- | Lowest-penalty candidate after boolean saturation of the diverse top-K. -bestDiscreteCandidate :: - TreeConfig -> (CondVec -> (Int, Int)) -> [CondVec] -> Maybe CondVec -bestDiscreteCandidate _ _ [] = Nothing -bestDiscreteCandidate cfg penaltyCV validCondVecs = - case saturateCandidates - Structural - (boolExpansion (synthConfig cfg)) - (sortedTopK cfg penaltyCV validCondVecs) of - [] -> Nothing - xs -> Just (snd (minimumBy (compare `on` fst) (decorate penaltyCV xs))) - -{- | AND/OR expansion of cached conditions to depth @maxDepth@ (each -combination is a single vector op, not an interpret). --} -boolExprsVec :: [CondVec] -> [CondVec] -> Int -> Int -> [CondVec] -boolExprsVec baseExprs prevExprs depth maxDepth - | depth == 0 = - baseExprs ++ boolExprsVec baseExprs prevExprs (depth + 1) maxDepth - | depth >= maxDepth = [] - | otherwise = combined ++ boolExprsVec baseExprs combined (depth + 1) maxDepth - where - combined = roundProducts prevExprs baseExprs - -data DedupMode = Structural | TruthVector - deriving (Eq, Show) - -{- | Saturate the pool with AND/OR combinations, deduplicating structurally -(byte-identical, first occurrence kept) or by truth vector (opt-in). --} -saturateCandidates :: DedupMode -> Int -> [CondVec] -> [CondVec] -saturateCandidates Structural maxDepth base = base' ++ go 1 base' seen0 - where - (base', seen0) = admitKeys Set.empty base - go !depth frontier seen - | depth >= maxDepth || null frontier = [] - | otherwise = - let (admitted, seen') = admitKeys seen (roundProducts frontier base) - in admitted ++ go (depth + 1) admitted seen' -saturateCandidates TruthVector maxDepth base = M.elems (go 1 frontier0 reps0) - where - (reps0, frontier0) = admitVecs M.empty base - go !depth frontier reps - | depth >= maxDepth || null frontier = reps - | otherwise = - let (reps', admitted) = admitVecs reps (roundProducts frontier base) - in go (depth + 1) admitted reps' - -{- | One combination round: @frontier × base@ via AND then OR, skipping -self-pairs (mirrors 'boolExprsVec' for byte-identical structural output). --} -roundProducts :: [CondVec] -> [CondVec] -> [CondVec] -roundProducts frontier base = - [ c - | e1 <- frontier - , e2 <- base - , not (eqExpr (cvExpr e1) (cvExpr e2)) - , c <- [combineAndVec e1 e2, combineOrVec e1 e2] - ] - --- | Admit candidates with a not-yet-seen normalized form, preserving order. -admitKeys :: Set.Set String -> [CondVec] -> ([CondVec], Set.Set String) -admitKeys = go [] - where - go acc seen [] = (reverse acc, seen) - go acc !seen (c : cs) - | structuralKey c `Set.member` seen = go acc seen cs - | otherwise = go (c : acc) (Set.insert (structuralKey c) seen) cs - -structuralKey :: CondVec -> String -structuralKey = show . normalize . cvExpr - -{- | Admit candidates by distinct truth vector, keeping the smallest-expression -representative per vector. --} -admitVecs :: - M.Map (VU.Vector Bool) CondVec -> - [CondVec] -> - (M.Map (VU.Vector Bool) CondVec, [CondVec]) -admitVecs = go [] - where - go acc reps [] = (reps, reverse acc) - go acc !reps (c : cs) = case M.lookup (cvVec c) reps of - Nothing -> go (c : acc) (M.insert (cvVec c) c reps) cs - Just r -> go acc (M.insert (cvVec c) (smaller r c) reps) cs - -smaller :: CondVec -> CondVec -> CondVec -smaller a b = case compare (eSize (cvExpr a)) (eSize (cvExpr b)) of - LT -> a - GT -> b - EQ -> if compareExpr (cvExpr a) (cvExpr b) /= GT then a else b - --- | Deduplicate 'CondVec's by normalized 'cvExpr', keeping the first. -dedupCVByExpr :: [CondVec] -> [CondVec] -dedupCVByExpr = go Set.empty - where - go _ [] = [] - go seen (cv : cvs) - | structuralKey cv `Set.member` seen = go seen cvs - | otherwise = cv : go (Set.insert (structuralKey cv) seen) cvs - --- | Deduplicate expressions by normalized form, keeping the first. -nubByExpr :: [Expr Bool] -> [Expr Bool] -nubByExpr = go Set.empty - where - go _ [] = [] - go seen (e : es) - | k `Set.member` seen = go seen es - | otherwise = e : go (Set.insert k seen) es - where - k = show (normalize e) diff --git a/src/DataFrame/DecisionTree/Predict.hs b/src/DataFrame/DecisionTree/Predict.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Predict.hs +++ /dev/null @@ -1,216 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Prediction, care-point identification, node validity, and tree loss. The -batched, cache-aware variants resolve each branch condition's truth vector -once per call instead of once per row. --} -module DataFrame.DecisionTree.Predict ( - predictWithTree, - predictManyWithTree, - predictManyWithTreeCached, - identifyCarePoints, - identifyCarePointsCached, - countCarePointErrors, - partitionIndices, - partitionIndicesCached, - majorityValueFromIndices, - computeTreeLoss, - computeTreeLossCached, - isValidAtNode, -) where - -import DataFrame.DecisionTree.CondVec ( - CondCache, - countErrorsByVec, - lookupCondVec, - ) -import DataFrame.DecisionTree.Types ( - CarePoint (..), - Direction (..), - Tree (..), - TreeConfig (..), - ) -import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector) -import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Internal.Interpreter (interpret) - -import Control.Exception (throw) -import Control.Monad.ST (ST) -import Data.Function (on) -import Data.List (maximumBy) -import qualified Data.Map.Strict as M -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Mutable as VM -import qualified Data.Vector.Unboxed as VU - -{- | A condition's truth vector over the DataFrame, or 'Nothing' on a -type/interpret failure (callers default such rows to the left child). --} -branchBool :: DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) -branchBool df cond = case interpret @Bool df cond of - Right (TColumn column) -> either (const Nothing) Just (toVector @Bool @VU.Vector column) - _ -> Nothing - --- | The target column as a label vector, or 'Nothing' on failure. -interpretLabelCol :: - forall a. (Columnable a) => DataFrame -> T.Text -> Maybe (V.Vector a) -interpretLabelCol df target = case interpret @a df (Col target) of - Right (TColumn column) -> either (const Nothing) Just (toVector @a column) - _ -> Nothing - --- | Predict the label for a single row by walking a fixed tree (@True@ → left). -predictWithTree :: - forall a. (Columnable a) => T.Text -> DataFrame -> Int -> Tree a -> a -predictWithTree _ _ _ (Leaf v) = v -predictWithTree target df idx (Branch cond left right) = - predictWithTree @a target df idx (childFor cond left right idx df) - -childFor :: Expr Bool -> Tree a -> Tree a -> Int -> DataFrame -> Tree a -childFor cond left right idx df = case branchBool df cond of - Nothing -> left - Just boolVals -> if boolVals VU.! idx then left else right - -predictManyWithTree :: - forall a. (Columnable a) => Tree a -> DataFrame -> V.Vector Int -> V.Vector a -predictManyWithTree = predictManyWithTreeCached @a M.empty - -{- | 'predictManyWithTree' resolving each branch condition through a 'CondCache'. -Each condition is read at most once per call rather than once per row. --} -predictManyWithTreeCached :: - forall a. - (Columnable a) => CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a -predictManyWithTreeCached cache tree df indices = V.create $ do - mv <- VM.new (V.length indices) - fill mv (V.zip (V.enumFromN 0 (V.length indices)) indices) tree - pure mv - where - fill :: VM.MVector s a -> V.Vector (Int, Int) -> Tree a -> ST s () - fill mv prs (Leaf v) = V.mapM_ (\(p, _) -> VM.write mv p v) prs - fill mv prs (Branch cond left right) = case lookupCondVec cache df cond of - Nothing -> fill mv prs left - Just boolVals -> fillSplit mv (V.partition (\(_, i) -> boolVals VU.! i) prs) left right - - fillSplit :: - VM.MVector s a -> - (V.Vector (Int, Int), V.Vector (Int, Int)) -> - Tree a -> - Tree a -> - ST s () - fillSplit mv (leftPrs, rightPrs) left right = fill mv leftPrs left >> fill mv rightPrs right - -identifyCarePoints :: - forall a. - (Columnable a) => - T.Text -> DataFrame -> V.Vector Int -> Tree a -> Tree a -> [CarePoint] -identifyCarePoints = identifyCarePointsCached @a M.empty - -{- | Rows the parent must route to a specific child for the (fixed) subtrees to -classify correctly; a 'CondCache' avoids re-interpreting subtree conditions. --} -identifyCarePointsCached :: - forall a. - (Columnable a) => - CondCache -> - T.Text -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -> - [CarePoint] -identifyCarePointsCached cache target df indices leftTree rightTree = - maybe [] carePoints (interpretLabelCol @a df target) - where - leftPreds = predictManyWithTreeCached cache leftTree df indices - rightPreds = predictManyWithTreeCached cache rightTree df indices - carePoints targetVals = V.toList (V.imapMaybe (checkPoint targetVals leftPreds rightPreds) indices) - -checkPoint :: - (Eq a) => - V.Vector a -> V.Vector a -> V.Vector a -> Int -> Int -> Maybe CarePoint -checkPoint targetVals leftPreds rightPreds k idx = - case (leftPreds V.! k == trueLabel, rightPreds V.! k == trueLabel) of - (True, False) -> Just (CarePoint idx GoLeft) - (False, True) -> Just (CarePoint idx GoRight) - _ -> Nothing - where - trueLabel = targetVals V.! idx - --- | Care points a free condition misroutes (uncached; for the linear path). -countCarePointErrors :: Expr Bool -> DataFrame -> [CarePoint] -> Int -countCarePointErrors cond df carePoints = - maybe (length carePoints) (`countErrorsByVec` carePoints) (branchBool df cond) - -partitionIndices :: - Expr Bool -> DataFrame -> V.Vector Int -> (V.Vector Int, V.Vector Int) -partitionIndices = partitionIndicesCached M.empty - -{- | 'partitionIndices' resolving the condition through a 'CondCache'; a miss -routes every index left (matching the uncached fallback). --} -partitionIndicesCached :: - CondCache -> - Expr Bool -> - DataFrame -> - V.Vector Int -> - (V.Vector Int, V.Vector Int) -partitionIndicesCached cache cond df indices = case lookupCondVec cache df cond of - Nothing -> (indices, V.empty) - Just boolVals -> V.partition (boolVals VU.!) indices - --- | A split is valid at a node when both children keep at least 'minLeafSize'. -isValidAtNode :: TreeConfig -> DataFrame -> V.Vector Int -> Expr Bool -> Bool -isValidAtNode cfg df indices c = - V.length t >= minLeafSize cfg && V.length f >= minLeafSize cfg - where - (t, f) = partitionIndices c df indices - -majorityValueFromIndices :: - forall a. (Columnable a, Ord a) => T.Text -> DataFrame -> V.Vector Int -> a -majorityValueFromIndices target df indices = majorityOf (countLabels (labelColOrThrow @a df target) indices) - -labelColOrThrow :: forall a. (Columnable a) => DataFrame -> T.Text -> V.Vector a -labelColOrThrow df target = case interpret @a df (Col target) of - Left e -> throw e - Right (TColumn column) -> either throw id (toVector @a column) - -countLabels :: (Ord a) => V.Vector a -> V.Vector Int -> M.Map a Int -countLabels vals = V.foldl' (\acc i -> M.insertWith (+) (vals V.! i) (1 :: Int) acc) M.empty - -majorityOf :: M.Map a Int -> a -majorityOf counts - | M.null counts = error "Empty indices in majorityValueFromIndices" - | otherwise = fst (maximumBy (compare `on` snd) (M.toList counts)) - -computeTreeLoss :: - forall a. - (Columnable a) => T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double -computeTreeLoss = computeTreeLossCached @a M.empty - --- | 0/1 loss of a tree over @indices@, with a 'CondCache' for the predictions. -computeTreeLossCached :: - forall a. - (Columnable a) => - CondCache -> T.Text -> DataFrame -> V.Vector Int -> Tree a -> Double -computeTreeLossCached cache target df indices tree - | V.null indices = 0 - | otherwise = - maybe 1.0 (treeLoss cache tree df indices) (interpretLabelCol @a df target) - -treeLoss :: - (Columnable a) => - CondCache -> Tree a -> DataFrame -> V.Vector Int -> V.Vector a -> Double -treeLoss cache tree df indices targetVals = - fromIntegral (countMismatches targetVals indices preds) - / fromIntegral (V.length indices) - where - preds = predictManyWithTreeCached cache tree df indices - -countMismatches :: (Eq a) => V.Vector a -> V.Vector Int -> V.Vector a -> Int -countMismatches targetVals indices preds = - V.length - (V.ifilter (\k _ -> targetVals V.! (indices V.! k) /= preds V.! k) preds) diff --git a/src/DataFrame/DecisionTree/Prune.hs b/src/DataFrame/DecisionTree/Prune.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Prune.hs +++ /dev/null @@ -1,66 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE ScopedTypeVariables #-} - -{- | Post-convergence simplification of a fitted tree and its expression form: -drop branches forced by path-condition entailment, collapse identical -siblings, and fold redundant nested conditionals. --} -module DataFrame.DecisionTree.Prune ( - pruneDead, - treeEq, - pruneExpr, -) where - -import DataFrame.DecisionTree.Types (Tree (..)) -import DataFrame.Internal.Column (Columnable) -import DataFrame.Internal.Expression (Expr (..), eqExpr) -import DataFrame.Internal.Simplify (PredFact, entails, factFalse, factTrue) - -{- | Drop branches whose test is forced by the path conditions reaching them, -and collapse @Branch c t t@ to @t@. Sound for the decidable threshold subset; -other tests are left untouched. --} -pruneDead :: forall a. (Columnable a) => Tree a -> Tree a -pruneDead = go [] - where - go :: [PredFact] -> Tree a -> Tree a - go _ (Leaf v) = Leaf v - go facts (Branch cond left right) = case entails facts cond of - Just True -> go facts left - Just False -> go facts right - Nothing -> - reconcile - cond - (go (addFact (factTrue cond) facts) left) - (go (addFact (factFalse cond) facts) right) - -reconcile :: (Columnable a) => Expr Bool -> Tree a -> Tree a -> Tree a -reconcile cond left right - | treeEq left right = left - | otherwise = Branch cond left right - -addFact :: Maybe PredFact -> [PredFact] -> [PredFact] -addFact (Just f) fs = f : fs -addFact Nothing fs = fs - -treeEq :: (Columnable a) => Tree a -> Tree a -> Bool -treeEq (Leaf x) (Leaf y) = x == y -treeEq (Branch c1 l1 r1) (Branch c2 l2 r2) = eqExpr c1 c2 && treeEq l1 l2 && treeEq r1 r2 -treeEq _ _ = False - -{- | Recursively fold @If@ expressions whose branches coincide or nest the same -condition; leave other expressions structurally unchanged. --} -pruneExpr :: forall a. (Columnable a) => Expr a -> Expr a -pruneExpr (If cond t0 f0) = collapseIf cond (pruneExpr t0) (pruneExpr f0) -pruneExpr (Unary op e) = Unary op (pruneExpr e) -pruneExpr (Binary op l r) = Binary op (pruneExpr l) (pruneExpr r) -pruneExpr e = e - -collapseIf :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a -collapseIf cond t f - | eqExpr t f = t - | If ci ti _ <- t, eqExpr cond ci = If cond ti f - | If ci _ fi <- f, eqExpr cond ci = If cond t fi - | otherwise = If cond t f diff --git a/src/DataFrame/DecisionTree/Regression.hs b/src/DataFrame/DecisionTree/Regression.hs --- a/src/DataFrame/DecisionTree/Regression.hs +++ b/src/DataFrame/DecisionTree/Regression.hs @@ -1,24 +1,23 @@ {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} -{- | Variance-reduction (weighted-SSE) regression trees, reusing the CART -feature machinery. Leaves predict the (weighted) mean of their rows. The -matrix-level 'fitRegTreeOn' lets gradient boosting refit on residuals without -re-extracting features each round. +{- | Variance-reduction (weighted-SSE) regression trees over the CART feature +machinery; leaves predict the weighted mean of their rows. 'fitRegTreeOn' lets +gradient boosting refit on residuals without re-extracting features. -} module DataFrame.DecisionTree.Regression ( RegTreeConfig (..), defaultRegTreeConfig, + -- | Implementation verb used by the fit\/predict instances and boosting. fitRegTreeOn, ) where -import Data.Function (on) +import Control.Parallel (par, pseq) import Data.Maybe (maybeToList) import qualified Data.Vector as V -import qualified Data.Vector.Algorithms.Merge as VA import qualified Data.Vector.Unboxed as VU -import DataFrame.DecisionTree.Cart (CartFeature (..)) +import DataFrame.DecisionTree.Cart (CartFeature (..), sortIndicesByValue) import DataFrame.DecisionTree.Types (Tree (..)) -- | Stopping criteria for the regression tree. @@ -48,97 +47,95 @@ VU.Vector Double -> Maybe (VU.Vector Double) -> Tree Double -fitRegTreeOn cfg feats y mw = go 0 (VU.enumFromN 0 n) +fitRegTreeOn cfg feats y mw = buildNode 0 (VU.enumFromN 0 n) featSorted where n = VU.length y - wt i = maybe 1 (VU.! i) mw - go depth idxs - | VU.length idxs < rtMinSamplesSplit cfg - || depth >= rtMaxDepth cfg = - Leaf (nodeMean idxs) - | otherwise = case bestSplit idxs of - Nothing -> Leaf (nodeMean idxs) - Just (fj, thr) -> - let vals = cfValues (feats V.! fj) - (lefts, rights) = VU.partition (\i -> vals VU.! i <= thr) idxs - in if VU.null lefts || VU.null rights - then Leaf (nodeMean idxs) - else - Branch - (cfPred (feats V.! fj) thr) - (go (depth + 1) lefts) - (go (depth + 1) rights) - nodeMean idxs = - let (sw, sy) = VU.foldl' (\(!a, !b) i -> (a + wt i, b + wt i * (y VU.! i))) (0, 0) idxs - in if sw == 0 then 0 else sy / sw - bestSplit idxs = - let (totW, totSY, totSY2) = moments idxs - nodeSSE = totSY2 - safeDiv (totSY * totSY) totW - candidates = - [ (red, fj, thr) - | fj <- [0 .. V.length feats - 1] - , (thr, red) <- featureSplits idxs fj totW totSY totSY2 nodeSSE - ] - in case candidates of - [] -> Nothing - _ -> - let (red, fj, thr) = maximumByFst candidates - in if red >= rtMinImpurityDecrease cfg && red > 0 - then Just (fj, thr) - else Nothing - featureSplits idxs fj totW totSY totSY2 nodeSSE = - let vals = cfValues (feats V.! fj) - sorted = sortByVal vals idxs - in sweep sorted vals totW totSY totSY2 nodeSSE - sweep sorted vals totW totSY totSY2 nodeSSE = go0 0 0 0 0 Nothing + weightAt i = maybe 1 (VU.! i) mw + featSorted = V.map (sortIndicesByValue . cfValues) feats + + buildNode depth idxs sortedByFeat + | depth >= rtMaxDepth cfg || VU.length idxs < rtMinSamplesSplit cfg = leaf + | otherwise = + maybe leaf (splitNode depth idxs sortedByFeat) (bestSplit idxs sortedByFeat) where + leaf = Leaf (weightedMean idxs) + + splitNode depth idxs sortedByFeat (fj, thr) + | VU.null lefts || VU.null rights = Leaf (weightedMean idxs) + | otherwise = + forceTree l `par` (forceTree r `pseq` Branch (cfPred (feats V.! fj) thr) l r) + where + vals = cfValues (feats V.! fj) + goesLeft i = vals VU.! i <= thr + lefts = VU.filter goesLeft idxs + rights = VU.filter (not . goesLeft) idxs + l = buildNode (depth + 1) lefts (V.map (VU.filter goesLeft) sortedByFeat) + r = + buildNode (depth + 1) rights (V.map (VU.filter (not . goesLeft)) sortedByFeat) + + weightedMean idxs = + let (w, sy) = VU.foldl' step (0, 0) idxs + step (!a, !b) i = (a + weightAt i, b + weightAt i * (y VU.! i)) + in if w == 0 then 0 else sy / w + + bestSplit idxs sortedByFeat + | null candidates = Nothing + | red > 0 && red >= rtMinImpurityDecrease cfg = Just (fj, thr) + | otherwise = Nothing + where + (totW, totSY, totSY2) = moments idxs + nodeSSE = sse totSY totSY2 totW + candidates = + [ (red', fj', thr') + | fj' <- [0 .. V.length feats - 1] + , (thr', red') <- + bestThreshold fj' (sortedByFeat V.! fj') totW totSY totSY2 nodeSSE + ] + (red, fj, thr) = maximumByFst candidates + + bestThreshold fj sorted totW totSY totSY2 nodeSSE = maybeToList (go 0 0 0 0 Nothing) + where + vals = cfValues (feats V.! fj) m = VU.length sorted - go0 !k !wl !syl !syl2 best - | k >= m - 1 = maybeToList best - | otherwise = - let i = sorted VU.! k - wi = wt i - yi = y VU.! i - wl' = wl + wi - syl' = syl + wi * yi - syl2' = syl2 + wi * yi * yi - vCur = vals VU.! i - vNext = vals VU.! (sorted VU.! (k + 1)) - nl = k + 1 - nr = m - nl - wr = totW - wl' - valid = - vCur /= vNext - && nl >= rtMinLeafSize cfg - && nr >= rtMinLeafSize cfg - && wl' > 0 - && wr > 0 - red = - nodeSSE - - ( (syl2' - safeDiv (syl' * syl') wl') - + ( (totSY2 - syl2') - - safeDiv ((totSY - syl') * (totSY - syl')) wr - ) - ) - thr = (vCur + vNext) / 2 - best' = - if valid && maybe True (\(_, b) -> red > b) best - then Just (thr, red) - else best - in go0 (k + 1) wl' syl' syl2' best' - moments = - VU.foldl' - ( \(!w, !sy, !sy2) i -> - let wi = wt i; yi = y VU.! i - in (w + wi, sy + wi * yi, sy2 + wi * yi * yi) - ) - (0, 0, 0) + go !k !wl !syl !syl2 best + | k >= m - 1 = best + | otherwise = go (k + 1) wl' syl' syl2' best' + where + i = sorted VU.! k + next = sorted VU.! (k + 1) + wi = weightAt i + yi = y VU.! i + wl' = wl + wi + syl' = syl + wi * yi + syl2' = syl2 + wi * yi * yi + wr = totW - wl' + leafSizesOk = k + 1 >= rtMinLeafSize cfg && m - (k + 1) >= rtMinLeafSize cfg + splittable = vals VU.! i /= vals VU.! next && leafSizesOk && wl' > 0 && wr > 0 + reduction = nodeSSE - (sse syl' syl2' wl' + sse (totSY - syl') (totSY2 - syl2') wr) + best' + | splittable && maybe True ((reduction >) . snd) best = + Just ((vals VU.! i + vals VU.! next) / 2, reduction) + | otherwise = best + moments = VU.foldl' step (0, 0, 0) + where + step (!w, !sy, !sy2) i = + let wi = weightAt i; yi = y VU.! i + in (w + wi, sy + wi * yi, sy2 + wi * yi * yi) + safeDiv :: Double -> Double -> Double safeDiv a b = if b == 0 then 0 else a / b -sortByVal :: VU.Vector Double -> VU.Vector Int -> VU.Vector Int -sortByVal vals = VU.modify (VA.sortBy (compare `on` (vals VU.!))) +-- | Weighted SSE of a node from its Σy, Σy², and total weight: @Σy² − (Σy)²/w@. +sse :: Double -> Double -> Double -> Double +sse sumY sumSq w = sumSq - safeDiv (sumY * sumY) w + +{- | Force a subtree to WHNF throughout so the spark scoring the sibling has +substantial work to evaluate; pure and value-preserving (cf. 'Tao'). +-} +forceTree :: Tree Double -> () +forceTree (Leaf v) = v `seq` () +forceTree (Branch _ l r) = forceTree l `seq` forceTree r maximumByFst :: (Ord a) => [(a, b, c)] -> (a, b, c) maximumByFst = foldr1 (\x@(a, _, _) y@(b, _, _) -> if a >= b then x else y) diff --git a/src/DataFrame/DecisionTree/Tao.hs b/src/DataFrame/DecisionTree/Tao.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Tao.hs +++ /dev/null @@ -1,266 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Tree Alternating Optimization: hold the tree fixed and re-optimize one node -at a time, bottom-up, minimizing care-point misroutes. Sibling subtrees at a -depth level are independent and optimized in parallel. --} -module DataFrame.DecisionTree.Tao ( - taoOptimize, - taoOptimizeCV, - taoIteration, - taoIterationCV, - optimizeNode, - findBestSplitTAO, -) where - -import DataFrame.DecisionTree.CondVec -import DataFrame.DecisionTree.Linear (bestLinearCandidate) -import DataFrame.DecisionTree.Pool ( - bestDiscreteCandidate, - candidateParChunk, - evalWithPenaltyVec, - ) -import DataFrame.DecisionTree.Predict -import DataFrame.DecisionTree.Prune (pruneDead) -import DataFrame.DecisionTree.Types -import DataFrame.Internal.Column (Columnable) -import DataFrame.Internal.DataFrame (DataFrame) -import DataFrame.Internal.Expression (Expr) - -import Control.Parallel (par, pseq) -import Control.Parallel.Strategies (parListChunk, rdeepseq, using) -import Data.Function (on) -import Data.List (foldl', minimumBy) -import Data.Maybe (catMaybes, mapMaybe) -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -{- | The constant per-fit context threaded through the node-optimization -recursion (the cache is rebuilt each iteration). --} -data TaoEnv = TaoEnv - { teCache :: !CondCache - , teCfg :: !TreeConfig - , teTarget :: !T.Text - , teConds :: ![CondVec] - , teDf :: !DataFrame - } - --- | Public TAO entry point over raw conditions; materializes each once. -taoOptimize :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -taoOptimize cfg target conds df = - taoOptimizeCV @a cfg target (mapMaybe (materializeCondVec df) conds) df - -{- | TAO outer loop over pre-evaluated candidates: iterate until the iteration -budget or convergence tolerance is reached, then prune dead branches. --} -taoOptimizeCV :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [CondVec] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -taoOptimizeCV cfg target condVecs df rootIndices initialTree = - go 0 initialTree (lossWith baseCache initialTree) - where - baseCache = condCacheFromVecs condVecs - lossWith cache = computeTreeLossCached @a cache target df rootIndices - go iter tree prevLoss - | iter >= taoIterations cfg = pruneDead tree - | prevLoss - newLoss < taoConvergenceTol cfg = pruneDead tree' - | otherwise = go (iter + 1) tree' newLoss - where - cache = addTreeCondsToCache df tree baseCache - tree' = taoIterationCV @a cache cfg target condVecs df rootIndices tree - newLoss = lossWith cache tree' - --- | Public single-iteration entry point. -taoIteration :: - forall a. - (Columnable a, Ord a) => - TreeConfig -> - T.Text -> - [Expr Bool] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -taoIteration cfg target conds df rootIndices tree = - let condVecs = mapMaybe (materializeCondVec df) conds - cache = addTreeCondsToCache df tree (condCacheFromVecs condVecs) - in taoIterationCV @a cache cfg target condVecs df rootIndices tree - --- | One bottom-to-top sweep: re-optimize every node level by level. -taoIterationCV :: - forall a. - (Columnable a, Ord a) => - CondCache -> - TreeConfig -> - T.Text -> - [CondVec] -> - DataFrame -> - V.Vector Int -> - Tree a -> - Tree a -taoIterationCV cache cfg target condVecs df rootIndices tree = - foldl' - (optimizeDepthLevel env rootIndices) - tree - [treeDepth tree, treeDepth tree - 1 .. 0] - where - env = TaoEnv cache cfg target condVecs df - -optimizeDepthLevel :: - forall a. - (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Int -> Tree a -optimizeDepthLevel env rootIndices tree = optimizeAtDepth @a env rootIndices tree 0 - -optimizeAtDepth :: - forall a. - (Columnable a, Ord a) => - TaoEnv -> V.Vector Int -> Tree a -> Int -> Int -> Tree a -optimizeAtDepth env indices tree currentDepth targetDepth - | currentDepth == targetDepth = optimizeNode @a env indices tree - | otherwise = case tree of - Leaf v -> Leaf v - Branch cond left right -> optimizeChildren @a env indices cond left right currentDepth targetDepth - -{- | Optimize the two subtrees over their disjoint index sets, scoring the left -in parallel with the right (the cache is read-only, so this is pure). --} -optimizeChildren :: - forall a. - (Columnable a, Ord a) => - TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Int -> Int -> Tree a -optimizeChildren env indices cond left right currentDepth targetDepth = - forceTreeWork left' `par` (forceTreeWork right' `pseq` Branch cond left' right') - where - (indicesL, indicesR) = partitionIndicesCached (teCache env) cond (teDf env) indices - left' = optimizeAtDepth @a env indicesL left (currentDepth + 1) targetDepth - right' = optimizeAtDepth @a env indicesR right (currentDepth + 1) targetDepth - -{- | Force a subtree's optimization work to WHNF so the parallel scheduler has -something substantial to evaluate; pure and value-preserving. --} -forceTreeWork :: Tree a -> () -forceTreeWork (Leaf v) = v `seq` () -forceTreeWork (Branch c l r) = c `seq` forceTreeWork l `seq` forceTreeWork r - -{- | Re-optimize one node: pick its best split, or collapse to a leaf when the -node is empty or the chosen split underflows 'minLeafSize'. --} -optimizeNode :: - forall a. (Columnable a, Ord a) => TaoEnv -> V.Vector Int -> Tree a -> Tree a -optimizeNode env indices tree - | V.null indices = tree - | otherwise = case tree of - Leaf _ -> leaf - Branch oldCond left right -> rebuiltBranch env indices oldCond left right leaf - where - leaf = Leaf (majorityValueFromIndices @a (teTarget env) (teDf env) indices) - -rebuiltBranch :: - forall a. - (Columnable a, Ord a) => - TaoEnv -> V.Vector Int -> Expr Bool -> Tree a -> Tree a -> Tree a -> Tree a -rebuiltBranch env indices oldCond left right leaf - | underflows = leaf - | otherwise = Branch newCond left right - where - newCond = findBestSplitTAO @a env indices left right oldCond - (l, r) = partitionIndicesCached (teCache env) newCond (teDf env) indices - underflows = V.length l < minLeafSize (teCfg env) || V.length r < minLeafSize (teCfg env) - -{- | The lowest-penalty replacement condition for a node, falling back to the -current condition when no valid candidate beats it. --} -findBestSplitTAO :: - forall a. - (Columnable a) => - TaoEnv -> V.Vector Int -> Tree a -> Tree a -> Expr Bool -> Expr Bool -findBestSplitTAO env indices leftTree rightTree currentCond - | V.null indices || null carePoints = currentCond - | pureReplacementLinear cfg - , Just c <- linearCandidate - , isValidAtNode cfg (teDf env) indices c = - c - | otherwise = bestOfPool penaltyCV currentCond pool - where - cfg = teCfg env - carePoints = - identifyCarePointsCached @a - (teCache env) - (teTarget env) - (teDf env) - indices - leftTree - rightTree - penaltyCV = evalWithPenaltyVec cfg carePoints - linearCandidate = bestLinearCandidate cfg (teDf env) carePoints - valid = filterValidCandidates cfg indices (teConds env) - pool = - candidatePool - env - indices - currentCond - (bestDiscreteCandidate cfg penaltyCV valid) - linearCandidate - -bestOfPool :: (CondVec -> (Int, Int)) -> Expr Bool -> [CondVec] -> Expr Bool -bestOfPool _ currentCond [] = currentCond -bestOfPool penaltyCV _ pool = cvExpr (minimumBy (compare `on` penaltyCV) pool) - -{- | Validity-filtered candidates the node could split on: both children must -keep at least 'minLeafSize'. Scored in parallel chunks, order preserved. --} -filterValidCandidates :: TreeConfig -> V.Vector Int -> [CondVec] -> [CondVec] -filterValidCandidates cfg indices condVecs = map snd (filter fst (zip validity condVecs)) - where - validity = - map (validAtNode cfg indices) condVecs - `using` parListChunk candidateParChunk rdeepseq - -validAtNode :: TreeConfig -> V.Vector Int -> CondVec -> Bool -validAtNode cfg indices cv = nTrue >= minLeaf && (V.length indices - nTrue) >= minLeaf - where - minLeaf = minLeafSize cfg - nTrue = - V.foldl' - (\ !acc i -> if cvVec cv VU.! i then acc + 1 else acc) - (0 :: Int) - indices - -{- | The candidate pool to minimize over: the current condition, the best -discrete candidate, and the linear candidate, each kept only if valid. --} -candidatePool :: - TaoEnv -> - V.Vector Int -> - Expr Bool -> - Maybe CondVec -> - Maybe (Expr Bool) -> - [CondVec] -candidatePool env indices currentCond discreteCV linearCandidate = - filter - (isValidAtNode (teCfg env) (teDf env) indices . cvExpr) - (catMaybes [currentCV, discreteCV, linearCV]) - where - currentCV = CondVec currentCond <$> lookupCondVec (teCache env) (teDf env) currentCond - linearCV = linearCandidate >>= materializeCondVec (teDf env) diff --git a/src/DataFrame/DecisionTree/Types.hs b/src/DataFrame/DecisionTree/Types.hs deleted file mode 100644 --- a/src/DataFrame/DecisionTree/Types.hs +++ /dev/null @@ -1,190 +0,0 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Shared types, configuration and ordering machinery for the decision-tree -learner. Imported by every other @DataFrame.DecisionTree.*@ module. --} -module DataFrame.DecisionTree.Types ( - Tree (..), - treeDepth, - TreeConfig (..), - SynthConfig (..), - defaultTreeConfig, - defaultSynthConfig, - ColumnOrdering (..), - orderable, - defaultColumnOrdering, - withOrdFrom, - CarePoint (..), - Direction (..), - ttrace, -) where - -import DataFrame.Internal.Column (Columnable) -import DataFrame.Internal.Expression (Expr (..)) -import qualified DataFrame.LinearSolver as LS - -import Data.Int (Int16, Int32, Int64, Int8) -import qualified Data.Map.Strict as M -import Data.Proxy (Proxy (..)) -import qualified Data.Text as T -import Data.Type.Equality (testEquality, (:~:) (..)) -import Data.Word (Word16, Word32, Word64, Word8) -import qualified Debug.Trace as Trace -import System.Environment (lookupEnv) -import System.IO.Unsafe (unsafePerformIO) -import Type.Reflection (SomeTypeRep (..), typeRep) - -{- | A fitted tree: a leaf value, or an internal node testing a boolean -expression with @True@ routing left. --} -data Tree a - = Leaf !a - | Branch !(Expr Bool) !(Tree a) !(Tree a) - deriving (Show) - -treeDepth :: Tree a -> Int -treeDepth (Leaf _) = 0 -treeDepth (Branch _ l r) = 1 + max (treeDepth l) (treeDepth r) - -{- | A row the parent node must route to a specific child for the subtrees to -classify it correctly (the TAO objective is the count of misroutes). --} -data CarePoint = CarePoint - { cpIndex :: !Int - , cpCorrectDir :: !Direction - } - deriving (Eq, Show) - -data Direction = GoLeft | GoRight - deriving (Eq, Show) - -data TreeConfig = TreeConfig - { maxTreeDepth :: Int - , minSamplesSplit :: Int - , minLeafSize :: Int - , percentiles :: [Int] - , expressionPairs :: Int - , synthConfig :: SynthConfig - , taoIterations :: Int - , taoConvergenceTol :: Double - , columnOrdering :: ColumnOrdering - , useLinearSolver :: Bool - , linearSolverConfig :: LS.SolverConfig - , minCarePointsForLinear :: Int - , pureReplacementLinear :: Bool - } - -data SynthConfig = SynthConfig - { maxExprDepth :: Int - , boolExpansion :: Int - , disallowedCombinations :: [(T.Text, T.Text)] - , complexityPenalty :: Double - , enableStringOps :: Bool - , enableCrossCols :: Bool - , enableArithOps :: Bool - , maxCategoricalSubsetCardinality :: Int - , perColumnQuota :: Maybe Int - } - deriving (Eq, Show) - -defaultSynthConfig :: SynthConfig -defaultSynthConfig = - SynthConfig - { maxExprDepth = 2 - , boolExpansion = 2 - , disallowedCombinations = [] - , complexityPenalty = 0.05 - , enableStringOps = True - , enableCrossCols = True - , enableArithOps = True - , maxCategoricalSubsetCardinality = 4 - , perColumnQuota = Just 3 - } - -defaultTreeConfig :: TreeConfig -defaultTreeConfig = - TreeConfig - { maxTreeDepth = 4 - , minSamplesSplit = 5 - , minLeafSize = 1 - , percentiles = [0, 10 .. 100] - , expressionPairs = 10 - , synthConfig = defaultSynthConfig - , taoIterations = 10 - , taoConvergenceTol = 1e-6 - , columnOrdering = defaultColumnOrdering - , useLinearSolver = True - , linearSolverConfig = LS.defaultSolverConfig - , minCarePointsForLinear = 10 - , pureReplacementLinear = False - } - -{- | Which column types support ordering for splits. Register a type with -'orderable' and combine with @<>@. --} -newtype ColumnOrdering = ColumnOrdering (M.Map SomeTypeRep OrdDict) - -instance Semigroup ColumnOrdering where - ColumnOrdering a <> ColumnOrdering b = ColumnOrdering (a <> b) - -instance Monoid ColumnOrdering where - mempty = ColumnOrdering M.empty - --- | Register a type as orderable for decision-tree splits. -orderable :: forall a. (Columnable a, Ord a) => ColumnOrdering -orderable = ColumnOrdering (M.singleton (SomeTypeRep (typeRep @a)) (OrdDict (Proxy @a))) - --- | All standard numeric, text, and primitive types. -defaultColumnOrdering :: ColumnOrdering -defaultColumnOrdering = mconcat (numericOrderings ++ otherOrderings) - -numericOrderings :: [ColumnOrdering] -numericOrderings = - [ orderable @Int - , orderable @Int8 - , orderable @Int16 - , orderable @Int32 - , orderable @Int64 - , orderable @Word - , orderable @Word8 - , orderable @Word16 - , orderable @Word32 - , orderable @Word64 - , orderable @Integer - , orderable @Double - , orderable @Float - ] - -otherOrderings :: [ColumnOrdering] -otherOrderings = - [orderable @Bool, orderable @Char, orderable @T.Text, orderable @String] - --- | Existential @Ord@ dictionary keyed by type representation. -data OrdDict where - OrdDict :: (Columnable a, Ord a) => Proxy a -> OrdDict - -{- | Run @k@ with the @Ord a@ instance recovered from the ordering registry, -or 'Nothing' when @a@ is not registered. --} -withOrdFrom :: - forall a r. (Columnable a) => ColumnOrdering -> ((Ord a) => r) -> Maybe r -withOrdFrom (ColumnOrdering m) k = case M.lookup (SomeTypeRep (typeRep @a)) m of - Just (OrdDict (_ :: Proxy b)) -> case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> Just k - Nothing -> Nothing - Nothing -> Nothing - -{-# NOINLINE taoTraceEnabled #-} -taoTraceEnabled :: Bool -taoTraceEnabled = unsafePerformIO (fmap (== Just "1") (lookupEnv "TAO_TRACE")) - --- | Emit a trace line when @TAO_TRACE=1@; a no-op otherwise. -ttrace :: String -> a -> a -ttrace msg x - | taoTraceEnabled = Trace.trace ("[TAO] " ++ msg) x - | otherwise = x diff --git a/src/DataFrame/Featurize/Internal.hs b/src/DataFrame/Featurize/Internal.hs deleted file mode 100644 --- a/src/DataFrame/Featurize/Internal.hs +++ /dev/null @@ -1,146 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} - -{- | Shared internal helpers used across the model fitters: turning a 'DataFrame' -plus a target/feature 'Expr' into the numeric matrices the algorithms consume, -and the common expression builders (affine score, arg-max/arg-min over named -scores) that every linear model and classifier would otherwise re-implement. --} -module DataFrame.Featurize.Internal ( - -- * Supervised extraction - featureNames, - numericMatrix, - targetDoubles, - targetValues, - - -- * Unsupervised extraction - Features (..), - extractFeatures, - columnExprName, - materializeColumn, - - -- * Expression builders - affineExpr, - argMaxExpr, - argMinExpr, -) where - -import Control.Exception (throw) -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..)) -import qualified DataFrame.Functions as F -import DataFrame.Internal.Column (Columnable) -import DataFrame.Internal.DataFrame (DataFrame, columnNames) -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.LinearAlgebra (Matrix, transposeM) -import DataFrame.Operations.Core (columnAsDoubleVector, columnAsVector) -import DataFrame.Operators ((.&&.), (.*.), (.+.), (.<=.), (.>=.)) - --- | Every column name except the supervised target's. -featureNames :: Expr a -> DataFrame -> [T.Text] -featureNames (Col target) df = filter (/= target) (columnNames df) -featureNames _ df = columnNames df - -{- | The named columns as a row-major @n×d@ matrix of doubles (non-numeric -columns are coerced through 'columnAsDoubleVector'), paired with the names. --} -numericMatrix :: [T.Text] -> DataFrame -> (V.Vector T.Text, Matrix) -numericMatrix names df = (V.fromList names, transposeM colMajor) - where - colMajor = V.fromList (map column names) - column name = case columnAsDoubleVector (F.col @Double name) df of - Right v -> v - Left e -> throw (asFeatureError name e) - -asFeatureError :: T.Text -> DataFrameException -> DataFrameException -asFeatureError name (TypeMismatchException ctx) = - TypeMismatchException - ctx - { errorColumnName = Just (T.unpack name) - , callingFunctionName = - Just - "model fit (feature columns must be numeric Double; drop or encode non-numeric columns)" - } -asFeatureError _ e = e - --- | The target column as a vector of doubles. -targetDoubles :: Expr Double -> DataFrame -> VU.Vector Double -targetDoubles expr df = case columnAsDoubleVector expr df of - Right v -> v - Left e -> throw e - --- | The target column as a vector of its own type (for classifiers). -targetValues :: (Columnable a) => Expr a -> DataFrame -> V.Vector a -targetValues expr df = case columnAsVector expr df of - Right v -> v - Left e -> throw e - -{- | The extracted feature columns of an unsupervised fit, in the shapes the -algorithms need: names, column-major vectors, the row-major matrix, and the -@(n, d)@ dimensions. --} -data Features = Features - { ftNames :: ![T.Text] - , ftCols :: ![VU.Vector Double] - , ftRows :: !Matrix - , ftN :: !Int - , ftD :: !Int - } - --- | Extract the given feature columns once, in every shape the fitters use. -extractFeatures :: [Expr Double] -> DataFrame -> Features -extractFeatures features df = Features names cols rows n d - where - names = map columnExprName features - cols = map (materializeColumn df) features - n = case cols of - (x: _) -> VU.length x - _ -> 0 - d = length cols - rows = V.generate n (\i -> VU.generate d (\j -> (cols !! j) VU.! i)) - --- | The column name behind a @Col@ feature expression. -columnExprName :: Expr Double -> T.Text -columnExprName (Col n) = n -columnExprName e = error ("expected a column expression, got " ++ show e) - --- | Interpret a @Col@ (or numeric) expression to a @Double@ vector. -materializeColumn :: DataFrame -> Expr Double -> VU.Vector Double -materializeColumn df e = case columnAsDoubleVector e df of - Right v -> v - Left err -> throw err - -{- | An affine score @b + Σ wⱼ·colⱼ@ over named columns, dropping zero weights -(the shared core of linear/logistic/SVM margins). --} -affineExpr :: Double -> [(Double, T.Text)] -> Expr Double -affineExpr b terms = - foldr - (.+.) - (F.lit b) - [F.lit w .*. (Col n :: Expr Double) | (w, n) <- terms, w /= 0] - -{- | The class whose score is greatest, as a nested-@If@ expression; ties go to -the earlier class. --} -argMaxExpr :: (Columnable a) => [(a, Expr Double)] -> Expr a -argMaxExpr = argExtreme (.>=.) - --- | The class whose score is smallest (e.g. nearest cluster by distance). -argMinExpr :: (Columnable a) => [(a, Expr Double)] -> Expr a -argMinExpr = argExtreme (.<=.) - -argExtreme :: - (Columnable a) => - (Expr Double -> Expr Double -> Expr Bool) -> [(a, Expr Double)] -> Expr a -argExtreme _ [] = error "argExtreme: no classes" -argExtreme _ [(c, _)] = Lit c -argExtreme cmp ((c, sc) : rest) = - If - (foldr ((\o acc -> cmp sc o .&&. acc) . snd) (F.lit True) rest) - (Lit c) - (argExtreme cmp rest) diff --git a/src/DataFrame/GMM.hs b/src/DataFrame/GMM.hs --- a/src/DataFrame/GMM.hs +++ b/src/DataFrame/GMM.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {- | Gaussian mixture models fitted by EM. Full covariance by default (with a diagonal option and an automatic fall-back when a covariance is not positive @@ -11,6 +12,7 @@ log-densities are available via 'gmmLogDensityExprs'. -} module DataFrame.GMM ( + module DataFrame.Model, CovType (..), GMMConfig (..), defaultGMMConfig, @@ -33,7 +35,7 @@ import DataFrame.Internal.Expression (Expr (..)) import DataFrame.LinearAlgebra (Matrix, logSumExp) import DataFrame.LinearAlgebra.Solve (backSubst, cholesky, forwardSubst) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.-.)) import DataFrame.Random (mkGen, sampleIndices) @@ -74,10 +76,12 @@ } deriving (Eq, Show) -instance Fit GMMConfig [Expr Double] GMMModel where +instance Fit GMMConfig [Expr Double] where + type ModelOf GMMConfig [Expr Double] = GMMModel fit = fitGMM -instance Predict GMMModel Int where +instance Predict GMMModel where + type Prediction GMMModel = Expr Int predict = gmmAssignExpr -- | Fit a Gaussian mixture over the given feature columns. diff --git a/src/DataFrame/KMeans.hs b/src/DataFrame/KMeans.hs --- a/src/DataFrame/KMeans.hs +++ b/src/DataFrame/KMeans.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} {- | k-means clustering (Lloyd's algorithm with k-means++ seeding and multiple restarts). 'fit' trains a 'KMeansModel' (inspectable centres); 'predict' is the @@ -9,6 +10,7 @@ 'kmeansDistanceExprs' / 'kmeansTransform'. -} module DataFrame.KMeans ( + module DataFrame.Model, KMeansConfig (..), defaultKMeansConfig, KMeansModel (..), @@ -17,6 +19,7 @@ ) where import Data.List (minimumBy) +import Data.Maybe (fromMaybe, listToMaybe) import Data.Ord (comparing) import qualified Data.Text as T import qualified Data.Vector as V @@ -27,7 +30,7 @@ import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..), UExpr (..)) import DataFrame.LinearAlgebra (Matrix, nearestCenter, sqDist) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.-.)) import DataFrame.Random (Gen, mkGen, nextDouble, nextIntR, splitGen) import DataFrame.Transform (Transform (..)) @@ -55,10 +58,12 @@ } deriving (Eq, Show) -instance Fit KMeansConfig [Expr Double] KMeansModel where +instance Fit KMeansConfig [Expr Double] where + type ModelOf KMeansConfig [Expr Double] = KMeansModel fit = fitKMeans -instance Predict KMeansModel Int where +instance Predict KMeansModel where + type Prediction KMeansModel = Expr Int predict m = argMinExpr (zip [0 :: Int ..] (map snd (kmeansDistanceExprs m))) -- | Fit k-means over the given feature columns. @@ -110,7 +115,7 @@ meanOf :: [VU.Vector Double] -> VU.Vector Double meanOf vs = - let d = VU.length (head vs) + let d = VU.length (VU.empty `fromMaybe` listToMaybe vs) s = foldr (VU.zipWith (+)) (VU.replicate d 0) vs in VU.map (/ fromIntegral (length vs)) s diff --git a/src/DataFrame/Learn.hs b/src/DataFrame/Learn.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Learn.hs @@ -0,0 +1,51 @@ +{- | The one import an application needs for the @dataframe-learn@ estimators: +the @fit@\/@predict@ verbs, every model config and fitted-model record, and all +their 'Fit'\/'Predict' instances. + +@ +import DataFrame.Learn +... +fit defaultRegTreeConfig target df +@ +-} +module DataFrame.Learn ( + module DataFrame.Model, + module DataFrame.LinearModel, + module DataFrame.SVM, + module DataFrame.SVM.RFF, + module DataFrame.PCA, + module DataFrame.PCA.Kernel, + module DataFrame.KMeans, + module DataFrame.GMM, + module DataFrame.DBSCAN, + module DataFrame.Boosting, + module DataFrame.SymbolicRegression, + module DataFrame.Synthesis, + module DataFrame.Segmented, + module DataFrame.DecisionTree, + module DataFrame.Metrics, + module DataFrame.Metrics.Report, + module DataFrame.ModelSelection, + module DataFrame.Transform, + module DataFrame.Transform.Serialize, +) where + +import DataFrame.Boosting +import DataFrame.DBSCAN +import DataFrame.DecisionTree +import DataFrame.GMM +import DataFrame.KMeans +import DataFrame.LinearModel +import DataFrame.Metrics +import DataFrame.Metrics.Report +import DataFrame.Model +import DataFrame.ModelSelection +import DataFrame.PCA +import DataFrame.PCA.Kernel +import DataFrame.SVM +import DataFrame.SVM.RFF +import DataFrame.Segmented +import DataFrame.SymbolicRegression +import DataFrame.Synthesis +import DataFrame.Transform +import DataFrame.Transform.Serialize diff --git a/src/DataFrame/LinearAlgebra.hs b/src/DataFrame/LinearAlgebra.hs deleted file mode 100644 --- a/src/DataFrame/LinearAlgebra.hs +++ /dev/null @@ -1,122 +0,0 @@ -{-# LANGUAGE BangPatterns #-} - -{- | Dependency-free dense linear algebra over row-major matrices, shared by the -models in @dataframe-learn@. Basic vector/matrix operations plus stability and -distance helpers; solvers live in "DataFrame.LinearAlgebra.Solve" and -eigenproblems in "DataFrame.LinearAlgebra.Eigen". --} -module DataFrame.LinearAlgebra ( - Matrix, - dot, - axpy, - scaleV, - matVec, - tMatVec, - gram, - transposeM, - identityM, - logSumExp, - sqDist, - nearestCenter, - epsNeighbors, -) where - -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -{- | Row-major dense matrix: an outer boxed vector of equal-length rows. An -@n×d@ matrix has @n@ rows of length @d@. --} -type Matrix = V.Vector (VU.Vector Double) - --- | Inner product of two equal-length vectors. -dot :: VU.Vector Double -> VU.Vector Double -> Double -dot a b = VU.sum (VU.zipWith (*) a b) -{-# INLINE dot #-} - --- | @axpy a x y = a*x + y@. -axpy :: Double -> VU.Vector Double -> VU.Vector Double -> VU.Vector Double -axpy a = VU.zipWith (\xi yi -> a * xi + yi) -{-# INLINE axpy #-} - --- | Scalar-vector product. -scaleV :: Double -> VU.Vector Double -> VU.Vector Double -scaleV a = VU.map (* a) -{-# INLINE scaleV #-} - --- | @matVec A v@ for @A@ of shape @n×d@ and @v@ of length @d@; result length @n@. -matVec :: Matrix -> VU.Vector Double -> VU.Vector Double -matVec a v = VU.convert (V.map (`dot` v) a) - --- | @tMatVec A v = Aᵀ v@ for @A@ of shape @n×d@, @v@ of length @n@; result length @d@. -tMatVec :: Matrix -> VU.Vector Double -> VU.Vector Double -tMatVec a v - | V.null a = VU.empty - | otherwise = V.foldl' step (VU.replicate d 0) (V.zipWith (,) vBoxed a) - where - d = VU.length (V.head a) - vBoxed = V.generate (V.length a) (v VU.!) - step !acc (vi, row) = axpy vi row acc - --- | @gram A = Aᵀ A@, the @d×d@ symmetric matrix of column inner products. -gram :: Matrix -> Matrix -gram a - | V.null a = V.empty - | otherwise = - V.generate d $ \i -> - VU.generate d $ \j -> - V.foldl' (\ !acc row -> acc + (row VU.! i) * (row VU.! j)) 0 a - where - d = VU.length (V.head a) - --- | Transpose an @n×d@ matrix to @d×n@. -transposeM :: Matrix -> Matrix -transposeM a - | V.null a = V.empty - | otherwise = V.generate d $ \j -> VU.generate n $ \i -> (a V.! i) VU.! j - where - n = V.length a - d = VU.length (V.head a) - --- | @d×d@ identity matrix. -identityM :: Int -> Matrix -identityM d = V.generate d $ \i -> VU.generate d $ \j -> if i == j then 1 else 0 - --- | Numerically stable @log Σ exp xᵢ@. -logSumExp :: VU.Vector Double -> Double -logSumExp xs - | VU.null xs = negate (1 / 0) - | otherwise = m + log (VU.sum (VU.map (\x -> exp (x - m)) xs)) - where - m = VU.maximum xs - --- | Squared Euclidean distance. -sqDist :: VU.Vector Double -> VU.Vector Double -> Double -sqDist a b = VU.sum (VU.zipWith (\x y -> let z = x - y in z * z) a b) -{-# INLINE sqDist #-} - --- | Index of and squared distance to the nearest centre. -nearestCenter :: - V.Vector (VU.Vector Double) -> VU.Vector Double -> (Int, Double) -nearestCenter centers p = - V.ifoldl' - ( \(!bi, !bd) i c -> - let dd = sqDist c p in if dd < bd then (i, dd) else (bi, bd) - ) - (-1, 1 / 0) - centers - -{- | Indices @j@ (excluding @i@) within squared radius @eps²@ of row @i@, by -brute force; @O(n d)@ per query. --} -epsNeighbors :: Double -> Matrix -> Int -> VU.Vector Int -epsNeighbors eps rows i = - VU.fromList - [ j - | j <- [0 .. n - 1] - , j /= i - , sqDist (rows V.! i) (rows V.! j) <= eps2 - ] - where - n = V.length rows - eps2 = eps * eps diff --git a/src/DataFrame/LinearAlgebra/Eigen.hs b/src/DataFrame/LinearAlgebra/Eigen.hs deleted file mode 100644 --- a/src/DataFrame/LinearAlgebra/Eigen.hs +++ /dev/null @@ -1,121 +0,0 @@ -{-# LANGUAGE BangPatterns #-} - -{- | Symmetric eigenproblems in pure Haskell: cyclic Jacobi for full -decomposition (PCA covariance, @m×m@ kernels) and power iteration for the -dominant eigenpair (FISTA step sizes). Deterministic, sign-canonicalised output. --} -module DataFrame.LinearAlgebra.Eigen ( - jacobiEigenSym, - powerIterTop, -) where - -import Control.Monad (forM_, when) -import Control.Monad.ST (runST) -import Data.List (sortBy) -import Data.Ord (Down (..), comparing) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import DataFrame.LinearAlgebra (Matrix, dot, matVec, scaleV) - -{- | Cyclic Jacobi eigendecomposition of a symmetric matrix. Eigenvalues are -returned in descending order paired with eigenvectors as rows; each eigenvector -is sign-canonicalised (largest-magnitude component positive) so the output is -unique. --} -jacobiEigenSym :: Matrix -> (VU.Vector Double, Matrix) -jacobiEigenSym a0 - | V.null a0 = (VU.empty, V.empty) - | otherwise = runST $ do - a <- VUM.new (d * d) - forM_ [0 .. d - 1] $ \i -> - forM_ [0 .. d - 1] $ \j -> - VUM.write a (i * d + j) ((a0 V.! i) VU.! j) - v <- VUM.replicate (d * d) 0 - forM_ [0 .. d - 1] $ \i -> VUM.write v (i * d + i) 1 - sweep a v 0 - afrozen <- VU.freeze a - vmat <- VU.freeze v - let diag = VU.generate d (\i -> afrozen VU.! (i * d + i)) - vecs = - V.generate d $ \col -> - VU.generate d $ \row -> vmat VU.! (row * d + col) - paired = - sortBy - (comparing (Down . fst)) - (zip (VU.toList diag) (V.toList vecs)) - pure - ( VU.fromList (map fst paired) - , V.fromList (map (canonicalSign . snd) paired) - ) - where - d = V.length a0 - maxSweeps = 100 - tol = 1e-12 - sweep a v s - | s >= maxSweeps = pure () - | otherwise = do - off <- offNorm a - when (off >= tol) $ do - forM_ [0 .. d - 2] $ \p -> - forM_ [p + 1 .. d - 1] $ \q -> rotate a v p q - sweep a v (s + 1) - offNorm a = go 0 0 - where - go i !acc - | i >= d = pure acc - | otherwise = do - r <- goRow i (i + 1) acc - go (i + 1) r - goRow i j !acc - | j >= d = pure acc - | otherwise = do - x <- VUM.read a (i * d + j) - goRow i (j + 1) (acc + x * x) - rotate a v p q = do - apq <- VUM.read a (p * d + q) - when (abs apq > 1e-300) $ do - app <- VUM.read a (p * d + p) - aqq <- VUM.read a (q * d + q) - let theta = (aqq - app) / (2 * apq) - s' = if theta == 0 then 1 else signum theta - t = s' / (abs theta + sqrt (theta * theta + 1)) - c = 1 / sqrt (t * t + 1) - sn = t * c - forM_ [0 .. d - 1] $ \i -> do - aip <- VUM.read a (i * d + p) - aiq <- VUM.read a (i * d + q) - VUM.write a (i * d + p) (c * aip - sn * aiq) - VUM.write a (i * d + q) (sn * aip + c * aiq) - forM_ [0 .. d - 1] $ \j -> do - apj <- VUM.read a (p * d + j) - aqj <- VUM.read a (q * d + j) - VUM.write a (p * d + j) (c * apj - sn * aqj) - VUM.write a (q * d + j) (sn * apj + c * aqj) - forM_ [0 .. d - 1] $ \i -> do - vip <- VUM.read v (i * d + p) - viq <- VUM.read v (i * d + q) - VUM.write v (i * d + p) (c * vip - sn * viq) - VUM.write v (i * d + q) (sn * vip + c * viq) - -canonicalSign :: VU.Vector Double -> VU.Vector Double -canonicalSign vec = - let idx = VU.maxIndex (VU.map abs vec) - in if vec VU.! idx < 0 then VU.map negate vec else vec - -{- | Dominant eigenvalue and eigenvector of a symmetric PSD matrix via power -iteration with a deterministic all-ones start. --} -powerIterTop :: Int -> Matrix -> (Double, VU.Vector Double) -powerIterTop iters a - | V.null a = (0, VU.empty) - | otherwise = go iters (normalize (VU.replicate d 1)) - where - d = V.length a - normalize v = - let nrm = sqrt (dot v v) in if nrm == 0 then v else scaleV (1 / nrm) v - go 0 v = (dot v (matVec a v), v) - go k v = - let av = matVec a v - nrm = sqrt (dot av av) - in if nrm < 1e-300 then (0, v) else go (k - 1) (scaleV (1 / nrm) av) diff --git a/src/DataFrame/LinearAlgebra/Solve.hs b/src/DataFrame/LinearAlgebra/Solve.hs deleted file mode 100644 --- a/src/DataFrame/LinearAlgebra/Solve.hs +++ /dev/null @@ -1,213 +0,0 @@ -{-# LANGUAGE BangPatterns #-} - -{- | Householder QR (for ordinary least squares) and Cholesky factorisation (for -ridge normal equations and Gaussian log-densities). Pure, deterministic, no -LAPACK; sound at the @d@ ≤ low-hundreds scales this library targets. --} -module DataFrame.LinearAlgebra.Solve ( - qrLeastSquares, - cholesky, - choleskySolve, - logDetFromChol, - forwardSubst, - backSubst, -) where - -import Control.Monad (forM_) -import Control.Monad.ST (ST, runST) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import DataFrame.LinearAlgebra (Matrix) - -{- | Solve @min ‖A x − b‖₂@ for an @n×d@ matrix @A@ (@n ≥ d@) via Householder QR. -@Left cols@ reports rank deficiency (near-zero @R@ diagonal) with the offending -column indices; @Right x@ is the least-squares solution. --} -qrLeastSquares :: Matrix -> VU.Vector Double -> Either [Int] (VU.Vector Double) -qrLeastSquares a b - | V.null a = Right VU.empty - | n < d = Left [0 .. d - 1] - | otherwise = runST $ do - mat <- VUM.new (n * d) - forM_ [0 .. n - 1] $ \i -> - forM_ [0 .. d - 1] $ \j -> - VUM.write mat (j * n + i) ((a V.! i) VU.! j) - rhs <- VUM.new n - forM_ [0 .. n - 1] $ \i -> VUM.write rhs i (b VU.! i) - deficient <- householder mat rhs n d - if not (null deficient) - then pure (Left deficient) - else do - x <- VUM.new d - backSubstQR mat rhs n d x - Right <$> VU.freeze x - where - n = V.length a - d = VU.length (V.head a) - -householder :: - VUM.STVector s Double -> VUM.STVector s Double -> Int -> Int -> ST s [Int] -householder mat rhs n d = go 0 [] - where - tol = 1e-10 - go k acc - | k >= d = pure (reverse acc) - | otherwise = do - normSq <- sumSq k - let alphaMag = sqrt normSq - akk <- VUM.read mat (k * n + k) - let alpha = if akk > 0 then negate alphaMag else alphaMag - if alphaMag < tol - then go (k + 1) (k : acc) - else do - VUM.write mat (k * n + k) (akk - alpha) - vNormSq <- sumSq k - if vNormSq < tol * tol - then go (k + 1) acc - else do - forM_ [k + 1 .. d - 1] $ \j -> reflectColumn k j vNormSq - reflectRhs k vNormSq - VUM.write mat (k * n + k) alpha - go (k + 1) acc - sumSq k = foldRows k - where - foldRows i - | i >= n = pure 0 - | otherwise = do - x <- VUM.read mat (k * n + i) - rest <- foldRows (i + 1) - pure (x * x + rest) - reflectColumn k j vNormSq = do - dotv <- dotV k j k - let beta = 2 * dotv / vNormSq - forM_ [k .. n - 1] $ \i -> do - vi <- VUM.read mat (k * n + i) - aij <- VUM.read mat (j * n + i) - VUM.write mat (j * n + i) (aij - beta * vi) - reflectRhs k vNormSq = do - dotv <- dotRhs k k - let beta = 2 * dotv / vNormSq - forM_ [k .. n - 1] $ \i -> do - vi <- VUM.read mat (k * n + i) - bi <- VUM.read rhs i - VUM.write rhs i (bi - beta * vi) - dotV k j i - | i >= n = pure 0 - | otherwise = do - vi <- VUM.read mat (k * n + i) - aij <- VUM.read mat (j * n + i) - rest <- dotV k j (i + 1) - pure (vi * aij + rest) - dotRhs k i - | i >= n = pure 0 - | otherwise = do - vi <- VUM.read mat (k * n + i) - bi <- VUM.read rhs i - rest <- dotRhs k (i + 1) - pure (vi * bi + rest) - -backSubstQR :: - VUM.STVector s Double -> - VUM.STVector s Double -> - Int -> - Int -> - VUM.STVector s Double -> - ST s () -backSubstQR mat rhs n d x = forM_ [d - 1, d - 2 .. 0] $ \i -> do - bi <- VUM.read rhs i - s <- sumAbove i (i + 1) 0 - rii <- VUM.read mat (i * n + i) - VUM.write x i ((bi - s) / rii) - where - sumAbove i j !acc - | j >= d = pure acc - | otherwise = do - rij <- VUM.read mat (j * n + i) - xj <- VUM.read x j - sumAbove i (j + 1) (acc + rij * xj) - -{- | Cholesky factor @L@ (lower-triangular, @A = L Lᵀ@) of a symmetric -positive-definite matrix, or 'Nothing' if a non-positive pivot is hit. --} -cholesky :: Matrix -> Maybe Matrix -cholesky a - | V.null a = Just V.empty - | otherwise = runST $ do - l <- VUM.replicate (d * d) 0 - ok <- buildL l - if ok then Just <$> freezeLower l else pure Nothing - where - d = V.length a - buildL l = go 0 - where - go j - | j >= d = pure True - | otherwise = do - s <- sumLk l j j (j - 1) 0 - let ajj = (a V.! j) VU.! j - diag = ajj - s - if diag <= 0 - then pure False - else do - let ljj = sqrt diag - VUM.write l (j * d + j) ljj - forM_ [j + 1 .. d - 1] $ \i -> do - sij <- sumLk l i j (j - 1) 0 - let aij = (a V.! i) VU.! j - VUM.write l (i * d + j) ((aij - sij) / ljj) - go (j + 1) - sumLk l i j k !acc - | k < 0 = pure acc - | otherwise = do - lik <- VUM.read l (i * d + k) - ljk <- VUM.read l (j * d + k) - sumLk l i j (k - 1) (acc + lik * ljk) - freezeLower l = do - frozen <- VU.freeze l - pure $ V.generate d $ \i -> VU.slice (i * d) d frozen - --- | Solve @L y = b@ for lower-triangular @L@. -forwardSubst :: Matrix -> VU.Vector Double -> VU.Vector Double -forwardSubst l b = runST $ do - y <- VUM.new d - forM_ [0 .. d - 1] $ \i -> do - let row = l V.! i - s <- sumKnown y row i 0 0 - VUM.write y i ((b VU.! i - s) / (row VU.! i)) - VU.freeze y - where - d = V.length l - sumKnown y row i j !acc - | j >= i = pure acc - | otherwise = do - yj <- VUM.read y j - sumKnown y row i (j + 1) (acc + (row VU.! j) * yj) - --- | Solve @Lᵀ x = y@ for lower-triangular @L@. -backSubst :: Matrix -> VU.Vector Double -> VU.Vector Double -backSubst l y = runST $ do - x <- VUM.new d - forM_ [d - 1, d - 2 .. 0] $ \i -> do - s <- sumKnown x i (i + 1) 0 - VUM.write x i ((y VU.! i - s) / ((l V.! i) VU.! i)) - VU.freeze x - where - d = V.length l - sumKnown x i j !acc - | j >= d = pure acc - | otherwise = do - xj <- VUM.read x j - sumKnown x i (j + 1) (acc + ((l V.! j) VU.! i) * xj) - -{- | Solve the SPD system @A x = b@ via Cholesky; 'Nothing' when @A@ is not -positive-definite. --} -choleskySolve :: Matrix -> VU.Vector Double -> Maybe (VU.Vector Double) -choleskySolve a b = do - l <- cholesky a - pure (backSubst l (forwardSubst l b)) - --- | @log det A = 2 Σ log Lᵢᵢ@ from a Cholesky factor @L@. -logDetFromChol :: Matrix -> Double -logDetFromChol l = 2 * V.sum (V.imap (\i row -> log (row VU.! i)) l) diff --git a/src/DataFrame/LinearModel.hs b/src/DataFrame/LinearModel.hs --- a/src/DataFrame/LinearModel.hs +++ b/src/DataFrame/LinearModel.hs @@ -5,7 +5,17 @@ module DataFrame.LinearModel ( module DataFrame.LinearModel.Regression, module DataFrame.LinearModel.Logistic, + -- Re-exported from the internal solver so the public configs' and records' + -- @SolverConfig@\/@LinearModel@ fields stay usable. + SolverConfig (..), + defaultSolverConfig, + LinearModel (..), ) where import DataFrame.LinearModel.Logistic import DataFrame.LinearModel.Regression +import DataFrame.LinearSolver ( + LinearModel (..), + SolverConfig (..), + defaultSolverConfig, + ) diff --git a/src/DataFrame/LinearModel/Logistic.hs b/src/DataFrame/LinearModel/Logistic.hs --- a/src/DataFrame/LinearModel/Logistic.hs +++ b/src/DataFrame/LinearModel/Logistic.hs @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} {- | Logistic regression: binary and one-vs-rest multiclass over the FISTA @@ -9,6 +10,7 @@ auxiliary expressions. -} module DataFrame.LinearModel.Logistic ( + module DataFrame.Model, LogisticConfig (..), defaultLogisticConfig, LogisticModel (..), @@ -16,6 +18,7 @@ logisticProbExprs, ) where +import Control.Parallel.Strategies (Strategy, parList, rseq, using) import Data.List (sort) import qualified Data.Map.Strict as M import qualified Data.Vector as V @@ -38,7 +41,7 @@ defaultSolverConfig, fitL1Logistic, ) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model -- | Hyperparameters for logistic regression (the FISTA solver config). newtype LogisticConfig = LogisticConfig {lgSolver :: SolverConfig} @@ -56,10 +59,12 @@ } deriving (Eq, Show) -instance (Columnable a, Ord a) => Fit LogisticConfig (Expr a) (LogisticModel a) where +instance (Columnable a, Ord a) => Fit LogisticConfig (Expr a) where + type ModelOf LogisticConfig (Expr a) = (LogisticModel a) fit = fitLogistic -instance (Columnable a, Ord a) => Predict (LogisticModel a) a where +instance (Columnable a, Ord a) => Predict (LogisticModel a) where + type Prediction (LogisticModel a) = Expr a predict m = argMaxExpr (labelledMargins m) -- | Fit one-vs-rest logistic regression; the target column supplies the classes. @@ -67,7 +72,7 @@ (Columnable a, Ord a) => LogisticConfig -> Expr a -> DataFrame -> LogisticModel a fitLogistic (LogisticConfig cfg) target df = - LogisticModel (V.fromList classes) (V.fromList (map fitOne classes)) + LogisticModel (V.fromList classes) (V.fromList models) where names = featureNames target df (nameVec, mat) = numericMatrix names df @@ -78,6 +83,24 @@ let labels = VU.generate (V.length ys) (\i -> if ys V.! i == c then 1 else -1) in fitL1Logistic cfg mat labels nameVec + -- Binary one-vs-rest fits two sign-mirrored models; solve once and negate. + -- Multiclass fits the per-class binary problems in parallel. + models = case classes of + [c0, _c1] -> let m0 = fitOne c0 in [m0, negateModel m0] + _ -> map fitOne classes `using` parList forceModel + +{- | Sign-flip a fitted binary model (the other one-vs-rest class in a 2-class +problem is its exact mirror). +-} +negateModel :: LinearModel -> LinearModel +negateModel m = + m{lmWeights = VU.map negate (lmWeights m), lmIntercept = negate (lmIntercept m)} + +{- | Force a fitted model's solver work by evaluating its (unboxed, strict) +weight vector, so 'parList' actually runs the per-class fits in parallel. +-} +forceModel :: Strategy LinearModel +forceModel m = lmWeights m `seq` rseq m -- | The raw margin @Expr@ for each class. logisticMarginExprs :: diff --git a/src/DataFrame/LinearModel/Regression.hs b/src/DataFrame/LinearModel/Regression.hs --- a/src/DataFrame/LinearModel/Regression.hs +++ b/src/DataFrame/LinearModel/Regression.hs @@ -1,17 +1,18 @@ +{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} -{- | Linear regression with the standard penalties: ordinary least squares -(Householder QR), ridge (Cholesky on the regularized normal equations), and -lasso / elastic net (FISTA). 'fit' produces a 'LinearRegressor' record; -'predict' compiles it to an @Expr Double@ over the raw feature columns. +{- | Linear regression with the standard penalties: OLS (QR), ridge (Cholesky), +and lasso\/elastic net (FISTA). 'fit' produces a 'LinearRegressor'; 'predict' +compiles it to an @Expr Double@ over the raw feature columns. -} module DataFrame.LinearModel.Regression ( + module DataFrame.Model, Penalty (..), LinearConfig (..), defaultLinearConfig, LinearRegressor (..), - predictLinear, ) where import qualified Data.Text as T @@ -34,7 +35,7 @@ fitProx, ) import DataFrame.LinearSolver.Loss (squaredLoss) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model -- | Regularization choice. @alpha@ is the penalty strength; @l1Ratio@ mixes L1/L2. data Penalty @@ -65,7 +66,9 @@ } deriving (Eq, Show) -instance Fit LinearConfig (Expr Double) LinearRegressor where +instance Fit LinearConfig (Expr Double) where + type ModelOf LinearConfig (Expr Double) = LinearRegressor + type FrameReq LinearConfig (Expr Double) = 'AllDoubleFrame fit (LinearConfig penalty cfg) target df = case penalty of OLS -> closedForm (olsSolve mat y) @@ -84,7 +87,8 @@ m = fitProx squaredLoss proxCfg mat y nameVec in LinearRegressor (lmWeights m) (lmIntercept m) nameVec penalty -instance Predict LinearRegressor Double where +instance Predict LinearRegressor where + type Prediction LinearRegressor = Expr Double predict m = affineExpr (regIntercept m) diff --git a/src/DataFrame/LinearSolver.hs b/src/DataFrame/LinearSolver.hs deleted file mode 100644 --- a/src/DataFrame/LinearSolver.hs +++ /dev/null @@ -1,472 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE ScopedTypeVariables #-} - -{- | Proximal-gradient (FISTA) solver for L1/L2-regularized generalized linear -models. 'fitL1Logistic' is the binary logistic split solver used by -'DataFrame.DecisionTree'; 'fitProx' generalizes it to any 'SmoothLoss' -(squared loss for lasso/elastic-net, squared hinge for LinearSVC). Features are -standardized internally and weights de-standardized, so the model applies to -raw column values. --} -module DataFrame.LinearSolver ( - -- * Model - LinearModel (..), - - -- * Configuration - SolverConfig (..), - defaultSolverConfig, - - -- * Solvers - fitL1Logistic, - fitProx, - - -- * Expr conversion - modelToExpr, - - -- * Internals (exposed for testing) - standardize, - columnStats, - softThreshold, - sigmoid, - dotProduct, -) where - -import qualified DataFrame.Functions as F -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.LinearAlgebra (Matrix, gram, scaleV) -import DataFrame.LinearAlgebra.Eigen (powerIterTop) -import DataFrame.LinearSolver.Loss ( - SmoothLoss (..), - logisticLoss, - sigmoid, - ) -import DataFrame.Operators ((.*.), (.+.), (.>.)) - -import Control.Monad.ST (ST, runST) -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM - -{- | A fitted linear classifier: predicts the positive class when -@sum (weights .* features) + intercept > 0@. Weights of exactly @0@ mark -features dropped by the L1 penalty (filtered out by 'modelToExpr'). --} -data LinearModel = LinearModel - { lmWeights :: !(VU.Vector Double) - , lmIntercept :: !Double - , lmFeatureNames :: !(V.Vector T.Text) - } - deriving (Eq, Show) - --- | Hyper-parameters for the FISTA solver. -data SolverConfig = SolverConfig - { scL1Lambda :: !Double - -- ^ Strength of the L1 penalty on weights (intercept is not regularized). - , scL2Lambda :: !Double - {- ^ Strength of the L2 penalty @(λ₂/2)·|w|²@ (Elastic Net; Zou & Hastie - 2005). Combined with @scL1Lambda@ this gives the standard elastic-net - objective @λ₁·|w|₁ + (λ₂/2)·|w|²@. At @scL2Lambda = 0@ the solver - reduces to pure L1 (the original behaviour). The Friedman/Hastie/ - Tibshirani 2010 glmnet proximal step under step size @1/L@ is - @softThreshold(z, λ₁/L) / (1 + λ₂/L)@ with @L = (d+1)/4 + λ₂@. - -} - , scMaxIter :: !Int - -- ^ Maximum number of FISTA iterations. - , scTol :: !Double - -- ^ Convergence tolerance on the weight delta (L-inf norm). - , scSampleWeights :: !(Maybe (VU.Vector Double)) - {- ^ Optional per-row sample weights, length @n@. @Nothing@ is uniform - weight 1 (legacy behaviour, A1-A18 path). The 1/N gradient - normalisation is preserved by convention: weights should have mean - 1 (i.e. @Σ w_i = N@) so the existing Lipschitz bound stays valid. - See 'fitLinearCandidate' in 'DataFrame.DecisionTree' for the - class-balanced construction @w_i = N / (2 · N_class(label_i))@. - -} - } - deriving (Eq, Show) - -defaultSolverConfig :: SolverConfig -defaultSolverConfig = - SolverConfig - { scL1Lambda = 0.005 - , scL2Lambda = 0.005 - , scMaxIter = 200 - , scTol = 1.0e-4 - , scSampleWeights = Nothing - } - -{- | Fit L1-regularized binary logistic regression by FISTA. Rows are feature -vectors of equal length; labels are in @{\-1,+1}@. Features are standardized -internally and weights de-standardized, so the model applies to raw values. --} -fitL1Logistic :: - SolverConfig -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - V.Vector T.Text -> - LinearModel -{-# INLINEABLE fitL1Logistic #-} -fitL1Logistic = runFista logisticLoss logisticLipschitz - where - logisticLipschitz _ keepN = fromIntegral (keepN + 1) / 4 - -{- | Fit any 'SmoothLoss' with the elastic-net proximal-gradient engine. The -Lipschitz constant uses the spectral norm of the standardized Gram matrix -(power iteration), which is tight for squared and squared-hinge losses where -the logistic trace bound would be far too loose. --} -fitProx :: - SmoothLoss -> - SolverConfig -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - V.Vector T.Text -> - LinearModel -fitProx loss = runFista loss specNormLipschitz - where - specNormLipschitz xKept _ = - let n = V.length xKept - gramN = V.map (scaleV (1 / fromIntegral n)) (gram xKept) - (specNorm, _) = powerIterTop 50 gramN - in slCurvBound loss * (specNorm + 1) - -{- | Shared FISTA scaffolding: standardize, drop near-constant columns, run the -inner loop, de-standardize. @lipschitzOf@ receives the standardized kept-feature -matrix and the number of kept columns and returns the smooth-part Lipschitz -bound; the L2 penalty contribution @λ₂@ is added here. --} -runFista :: - SmoothLoss -> - (Matrix -> Int -> Double) -> - SolverConfig -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - V.Vector T.Text -> - LinearModel -runFista loss lipschitzOf cfg rows labels featureNames - | n == 0 || d == 0 = zeroModel - | otherwise = - let (!means, !stds, !variances) = columnStats rows - !keep = keptIndices variances - in if VU.null keep - then zeroModel - else - let !meansKept = gatherBy keep means - !stdsKept = gatherBy keep stds - !xKept = V.map (standardizeRowKept keep means stds) rows - !lipschitz = - lipschitzOf xKept (VU.length keep) + scL2Lambda cfg - (!wStdKept, !bStd) = - fistaLoop - loss - (scL1Lambda cfg) - (scL2Lambda cfg) - lipschitz - (scMaxIter cfg) - (scTol cfg) - (scSampleWeights cfg) - xKept - labels - (VU.replicate (VU.length keep) 0) - 0 - !wRawKept = VU.zipWith (/) wStdKept stdsKept - !bRaw = bStd - VU.sum (VU.zipWith (*) wRawKept meansKept) - in LinearModel (expandWeights d keep wRawKept) bRaw featureNames - where - !n = V.length rows - !d = V.length featureNames - zeroModel = LinearModel (VU.replicate d 0) 0 featureNames - -{- | Indices of columns whose variance clears the near-constant threshold. -Columns below it are dropped before fitting; their weight ends up @0@. --} -keptIndices :: VU.Vector Double -> VU.Vector Int -keptIndices variances = - VU.fromList - [ j - | j <- [0 .. VU.length variances - 1] - , VU.unsafeIndex variances j >= 1.0e-12 - ] - -{- | Gather the entries of @v@ at @idxs@, preserving order. unsafeIndex is -safe: every index in @idxs@ is in range by construction. --} -gatherBy :: VU.Vector Int -> VU.Vector Double -> VU.Vector Double -gatherBy idxs v = VU.map (VU.unsafeIndex v) idxs - -{- | Standardize one row to the kept columns only (subtract column mean, divide -by column std). unsafeIndex is safe: rows share the column layout. --} -standardizeRowKept :: - VU.Vector Int -> - VU.Vector Double -> - VU.Vector Double -> - VU.Vector Double -> - VU.Vector Double -standardizeRowKept keep means stds row = VU.map standardizeAt keep - where - standardizeAt j = - (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j - -{- | Scatter kept-column weights back into a full-width vector, with @0@ for -the dropped (near-constant) columns. --} -expandWeights :: Int -> VU.Vector Int -> VU.Vector Double -> VU.Vector Double -expandWeights d keep wKept = VU.create $ do - mv <- VUM.replicate d 0 - VU.iforM_ keep $ \k j -> VUM.unsafeWrite mv j (VU.unsafeIndex wKept k) - pure mv - -{- | Convert a fitted model to an 'Expr Bool' over its feature columns, -dropping zero-weight features. With no non-zero weights it returns the -constant @Lit (intercept > 0)@. --} -modelToExpr :: LinearModel -> Expr Bool -modelToExpr m = - case nonZero of - [] -> F.lit (b > 0) - (w0, n0) : rest -> score rest (term w0 n0) .>. F.lit (0 :: Double) - where - b = lmIntercept m - nonZero = - [ (w, n) - | (w, n) <- zip (VU.toList (lmWeights m)) (V.toList (lmFeatureNames m)) - , w /= 0 - ] - term w n = F.lit w .*. (Col n :: Expr Double) - score rest first = foldl (\acc (w, n) -> acc .+. term w n) first rest .+. F.lit b - -{- | Per-column @(means, stds, variances)@ of a feature matrix. Cheaper than -'standardize' when only the statistics are needed. unsafeIndex within is -safe: all rows share width @d@. --} -columnStats :: - V.Vector (VU.Vector Double) -> - (VU.Vector Double, VU.Vector Double, VU.Vector Double) -columnStats x - | V.null x = (VU.empty, VU.empty, VU.empty) - | otherwise = - let !d = VU.length (V.unsafeHead x) - !invN = 1 / fromIntegral (V.length x) - !means = columnMeans d invN x - !variances = columnVariances d invN means x - !stds = VU.map (\v -> if v < 1e-12 then 1 else sqrt v) variances - in (means, stds, variances) - --- | Mean of each of the @d@ columns; @invN@ is @1 / nRows@. -columnMeans :: Int -> Double -> V.Vector (VU.Vector Double) -> VU.Vector Double -columnMeans d invN x = runST $ do - acc <- VUM.replicate d 0 - V.forM_ x $ \row -> - VU.iforM_ row $ \j v -> VUM.unsafeModify acc (+ v) j - scaleInPlace invN acc - VU.unsafeFreeze acc - --- | Variance of each of the @d@ columns about the supplied @means@. -columnVariances :: - Int -> - Double -> - VU.Vector Double -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -columnVariances d invN means x = runST $ do - acc <- VUM.replicate d 0 - V.forM_ x $ \row -> - VU.iforM_ row $ \j v -> - let !c = v - VU.unsafeIndex means j in VUM.unsafeModify acc (+ c * c) j - scaleInPlace invN acc - VU.unsafeFreeze acc - --- | Multiply every element of a mutable vector by @factor@ in place. -scaleInPlace :: Double -> VUM.MVector s Double -> ST s () -scaleInPlace factor mv = go 0 - where - go !j - | j >= VUM.length mv = pure () - | otherwise = VUM.unsafeModify mv (* factor) j >> go (j + 1) - -{- | Standardize each column to zero mean and unit variance, also returning -@(means, stds, variances)@. Near-constant columns get std @1@; callers use -the raw variances to detect and drop them (see 'fitL1Logistic'). --} -standardize :: - V.Vector (VU.Vector Double) -> - ( V.Vector (VU.Vector Double) - , VU.Vector Double - , VU.Vector Double - , VU.Vector Double - ) -standardize x - | V.null x = (x, VU.empty, VU.empty, VU.empty) - | otherwise = - let (!means, !stds, !variances) = columnStats x - !d = VU.length (V.unsafeHead x) - standardizeRow row = - VU.generate d $ \j -> - (VU.unsafeIndex row j - VU.unsafeIndex means j) / VU.unsafeIndex stds j - in (V.map standardizeRow x, means, stds, variances) - -{- | Proximal operator for the L1 norm: shrink @v@ toward zero by @lambda@, -clamping at zero. --} -softThreshold :: Double -> Double -> Double -softThreshold lambda v - | v > lambda = v - lambda - | v < -lambda = v + lambda - | otherwise = 0 - -{- | Dot product of two unboxed vectors. Caller must ensure equal length; -lengths are not checked. --} -dotProduct :: VU.Vector Double -> VU.Vector Double -> Double -dotProduct u v = VU.sum (VU.zipWith (*) u v) - -{- | Gradient of the average loss at @(w, b)@. Returns @(gradW, gradB)@. - -Sample-weighted variant: when @sampleWeights@ is @Just ws@ the per-row -contribution is multiplied by @ws[i]@. With weights of mean 1 -(i.e. @Σ w_i = N@; the class-balanced convention used by -'fitLinearCandidate'), the @1/N@ normalisation is preserved exactly. --} -lossGradient :: - SmoothLoss -> - Maybe (VU.Vector Double) -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - VU.Vector Double -> - Double -> - (VU.Vector Double, Double) -lossGradient loss sampleWeights features labels w b = (gradW, gradB) - where - !invN = 1 / fromIntegral (V.length features) - !coeffs = rowCoeffs loss sampleWeights features labels w b invN - !gradW = accumulateGradW (VU.length w) features coeffs - !gradB = VU.sum coeffs - -{- | Per-row loss coefficient @c_i = ℓ'(y_i, z_i) / N@ at margin -@z_i = w·x_i + b@, optionally scaled by @ws[i]@. - -unsafeIndex is safe: @i@ ranges over @[0,n-1]@ and @labels@ / -@sampleWeights@ both have length @n@ by construction. --} -rowCoeffs :: - SmoothLoss -> - Maybe (VU.Vector Double) -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - VU.Vector Double -> - Double -> - Double -> - VU.Vector Double -rowCoeffs loss sampleWeights features labels w b invN = - VU.generate (V.length features) $ \i -> - let !yi = VU.unsafeIndex labels i - !row = V.unsafeIndex features i - !z = dotProduct w row + b - !base = slGradZ loss yi z * invN - in case sampleWeights of - Nothing -> base - Just ws -> base * VU.unsafeIndex ws i - -{- | Accumulate the weight gradient in one pass over every (row, feature) -pair, scattering into a length-@d@ mutable vector. --} -accumulateGradW :: - Int -> V.Vector (VU.Vector Double) -> VU.Vector Double -> VU.Vector Double -accumulateGradW d features coeffs = runST $ do - mv <- VUM.replicate d 0 - V.iforM_ features $ \i row -> - let !c = VU.unsafeIndex coeffs i - in VU.iforM_ row $ \j v -> VUM.unsafeModify mv (+ c * v) j - VU.unsafeFreeze mv - -{- | Inner FISTA loop over standardized features. Returns the final @(w, b)@; -the caller is responsible for de-standardization. - -@lambda1@ and @lambda2@ are the L1 / L2 penalty strengths; @lp@ is the -Lipschitz constant of the smooth part @(d+1)/4 + λ₂@. The Elastic-Net -proximal step is applied per FHT 2010 glmnet §2.6: -@prox(z) = softThreshold(z, λ₁/lp) / (1 + λ₂/lp)@. --} -fistaLoop :: - SmoothLoss -> - Double -> - Double -> - Double -> - Int -> - Double -> - Maybe (VU.Vector Double) -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - VU.Vector Double -> - Double -> - (VU.Vector Double, Double) -fistaLoop loss lambda1 lambda2 lp maxIter tol sampleWeights features labels w0 b0 = - go 0 w0 b0 w0 b0 1.0 - where - !shrink = lambda1 / lp - !ridgeDenom = 1 + lambda2 / lp - !stepInv = 1 / lp - proxStep = fistaProxStep loss sampleWeights features labels shrink ridgeDenom stepInv - go !iter !xWPrev !xBPrev !yW !yB !t - | iter >= maxIter = (xWPrev, xBPrev) - | iter > 0 && delta < tol = (xW, xB) - | otherwise = go (iter + 1) xW xB yWNew yBNew tNew - where - (!xW, !xB) = proxStep yW yB - !delta = if VU.null xW then 0 else deltaInf xWPrev xW - (!yWNew, !yBNew, !tNew) = fistaMomentum t xWPrev xBPrev xW xB - -{- | One fused FISTA prox step: gradient step plus the Elastic-Net -proximal operator (soft-threshold then ridge shrinkage), without -materializing the intermediate trial weights. - -The Elastic-Net prox of @g(w) = λ₁·|w|₁ + (λ₂/2)·|w|²@ at step @1/lp@ is -@softThreshold(z, λ₁/lp) / (1 + λ₂/lp)@ (FHT 2010 glmnet §2.6; Beck & -Teboulle 2009 §4). The intercept is unregularised (no L1 or L2 applied). --} -fistaProxStep :: - SmoothLoss -> - Maybe (VU.Vector Double) -> - V.Vector (VU.Vector Double) -> - VU.Vector Double -> - Double -> - Double -> - Double -> - VU.Vector Double -> - Double -> - (VU.Vector Double, Double) -fistaProxStep loss sampleWeights features labels shrink ridgeDenom stepInv yW yB = - let (gW, gB) = lossGradient loss sampleWeights features labels yW yB - !wNew = - VU.zipWith - (\yi gi -> softThreshold shrink (yi - gi * stepInv) / ridgeDenom) - yW - gW - !bNew = yB - gB * stepInv - in (wNew, bNew) - -{- | Nesterov momentum extrapolation: new look-ahead point @(yW, yB)@ and the -updated step size @t@. --} -fistaMomentum :: - Double -> - VU.Vector Double -> - Double -> - VU.Vector Double -> - Double -> - (VU.Vector Double, Double, Double) -fistaMomentum t xWPrev xBPrev xW xB = - let !tNew = (1 + sqrt (1 + 4 * t * t)) / 2 - !mom = (t - 1) / tNew - !yW = VU.zipWith (\new old -> new + mom * (new - old)) xW xWPrev - !yB = xB + mom * (xB - xBPrev) - in (yW, yB, tNew) - -{- | L-inf norm of the weight delta. unsafeIndex is safe: both vectors share -the same length by construction. --} -{-# INLINE deltaInf #-} -deltaInf :: VU.Vector Double -> VU.Vector Double -> Double -deltaInf xWPrev = VU.ifoldl' (\acc i x -> max acc (abs (x - VU.unsafeIndex xWPrev i))) 0 diff --git a/src/DataFrame/LinearSolver/Loss.hs b/src/DataFrame/LinearSolver/Loss.hs deleted file mode 100644 --- a/src/DataFrame/LinearSolver/Loss.hs +++ /dev/null @@ -1,47 +0,0 @@ -{-# LANGUAGE OverloadedStrings #-} - -{- | Smooth losses for the proximal-gradient engine. Each carries its -derivative @∂ℓ/∂z@ at @z = w·x + b@ and a global bound on the curvature -@∂²ℓ/∂z²@ (used for the FISTA step size). --} -module DataFrame.LinearSolver.Loss ( - SmoothLoss (..), - sigmoid, - logisticLoss, - squaredLoss, - sqHingeLoss, -) where - -import qualified Data.Text as T - -{- | A convex, @C¹@ per-sample loss @ℓ(y, z)@. 'slGradZ' is @∂ℓ/∂z@; -'slCurvBound' bounds @∂²ℓ/∂z²@ over all @(y, z)@. --} -data SmoothLoss = SmoothLoss - { slName :: !T.Text - , slGradZ :: Double -> Double -> Double - , slCurvBound :: !Double - } - --- | Numerically stable logistic sigmoid. -sigmoid :: Double -> Double -sigmoid z - | z >= 0 = 1 / (1 + exp (-z)) - | otherwise = let ez = exp z in ez / (1 + ez) - --- | Binary logistic loss for labels in @{\-1,+1}@: @ℓ = log(1 + exp(-y z))@. -logisticLoss :: SmoothLoss -logisticLoss = - SmoothLoss "logistic" (\y z -> negate (y * sigmoid (negate (y * z)))) 0.25 - --- | Squared error for regression: @ℓ = ½ (z - y)²@. -squaredLoss :: SmoothLoss -squaredLoss = SmoothLoss "squared" (flip (-)) 1.0 - --- | Squared hinge for classification (LinearSVC default), labels @{\-1,+1}@. -sqHingeLoss :: SmoothLoss -sqHingeLoss = - SmoothLoss - "squared_hinge" - (\y z -> let m = 1 - y * z in if m > 0 then negate (2 * y * m) else 0) - 2.0 diff --git a/src/DataFrame/Model.hs b/src/DataFrame/Model.hs --- a/src/DataFrame/Model.hs +++ b/src/DataFrame/Model.hs @@ -1,5 +1,12 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} {- | The two verbs every model speaks. Instead of a per-model @fitX@ / @xExpr@ zoo, every estimator is an instance of these classes: @@ -20,50 +27,127 @@ functions — they are not the one canonical prediction, so they are not forced through 'predict'. -Supervised 'fit' treats every non-target column as a feature; use -'selectFeatures' to restrict to an explicit set when the frame also carries ids, -timestamps, or a second candidate target. +'fit' is a single, frame-polymorphic verb: it accepts an untyped 'DataFrame' or a +phantom-typed 'DataFrame.Typed.Types.TypedDataFrame', via 'ToDataFrame'. When a +model declares a schema requirement (via 'FrameReq'), that requirement is checked +at /compile time/ for a typed frame and is a no-op for an untyped one. Linear +regression, for instance, requires an all-'Double' frame ('AllDoubleFrame'), so +@fit@ on a typed frame with a non-'Double' column is a compile error; on an +untyped frame the same mistake surfaces as a fit-time error. -} module DataFrame.Model ( Fit (..), + ToDataFrame (..), + Fitted (..), + FitResult, + FrameFor, + FrameKind (..), + CheckFrame, + AllDouble, Predict (..), - selectFeatures, + AsTExpr, + ToTExpr (..), ) where -import qualified Data.Text as T +import Data.Kind (Constraint, Type) import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Operations.Subset (select) +import DataFrame.Typed.Freeze (ToDataFrame (..), thaw) +import DataFrame.Typed.Schema (AllDouble) +import DataFrame.Typed.Types (AsTExpr, TExpr (..), ToTExpr (..), TypedDataFrame) -{- | Train a model. @cfg@ is the hyperparameter config; @input@ is the supervised -target @Expr a@ or the unsupervised feature list @[Expr Double]@. The config and -input together determine the model, so @fit cfg target df@ needs no annotation -(a classifier's label type comes from its @Expr a@ target). +{- | A model trained on a typed frame, carrying the schema @cols@ as a phantom so +its 'predict' yields a typed 'TExpr'. Use 'fittedModel' to recover the bare model +record (coefficients, etc.). -} -class Fit cfg input model | cfg input -> model where - fit :: cfg -> input -> DataFrame -> model +newtype Fitted (cols :: [Type]) model = Fitted {fittedModel :: model} -{- | Compile a fitted model's canonical prediction to an expression over the raw -columns. The result type @r@ is determined by the model. +{- | The type 'fit' returns for a given frame source: the bare @model@ for an +untyped 'DataFrame', or a schema-tagged 'Fitted' for a 'TypedDataFrame'. Training +on an untyped frame is therefore unchanged. -} -class Predict model r | model -> r where - predict :: model -> Expr r +type family FitResult (f :: Type) (model :: Type) :: Type where + FitResult DataFrame model = model + FitResult (TypedDataFrame cols) model = Fitted cols model -{- | Restrict @df@ to exactly the named feature columns plus the supervised -target (when the target is a column), so a following 'fit' trains on those -features only. +{- | The frame a given training @input@ is fit against: an untyped target/feature +input pairs with a plain 'DataFrame'; a typed input ('TExpr' / @['TExpr']@) pairs +with a 'TypedDataFrame' over the /same/ schema. So the target expression and the +frame are forced to share their columns at compile time. +-} +type family FrameFor (input :: Type) :: Type where + FrameFor (Expr a) = DataFrame + FrameFor (TExpr cols a) = TypedDataFrame cols + FrameFor [Expr Double] = DataFrame + FrameFor [TExpr cols Double] = TypedDataFrame cols -Supervised 'fit' otherwise uses /every/ non-target column as a feature — -convenient on a clean frame, a leakage hazard when the frame carries ids, -timestamps, or a second candidate target. This mirrors the explicit -@[Expr Double]@ feature list the unsupervised fitters already take: +{- | The schema requirement a model places on its training frame. 'AnyFrame' +imposes nothing (the default); 'AllDoubleFrame' demands every column be 'Double'. +-} +data FrameKind = AnyFrame | AllDoubleFrame -> model = fit defaultLinearConfig target (selectFeatures ["age", "income"] target df) +{- | Turn a model's 'FrameKind' requirement into a constraint on the actual frame +type. Untyped 'DataFrame's are never constrained (they are runtime-checked); a +typed frame must satisfy the requirement at compile time. -} -selectFeatures :: [T.Text] -> Expr a -> DataFrame -> DataFrame -selectFeatures cols target = select (cols ++ targetCols target) - where - targetCols :: Expr b -> [T.Text] - targetCols (Col n) = [n] - targetCols _ = [] +type family CheckFrame (req :: FrameKind) (f :: Type) :: Constraint where + CheckFrame _ DataFrame = () + CheckFrame 'AnyFrame _ = () + CheckFrame 'AllDoubleFrame (TypedDataFrame cols) = AllDouble cols + +{- | Train a model. @cfg@ is the hyperparameter config; @input@ is the supervised +target @Expr a@ or the unsupervised feature list @[Expr Double]@; the frame is +any 'ToDataFrame' source (an untyped 'DataFrame' or a 'TypedDataFrame'). The +config and input together determine the model, so no annotation is needed (a +classifier's label type comes from its @Expr a@ target). + +A model overrides 'FrameReq' to demand a schema shape (e.g. 'AllDoubleFrame'), +which 'fit' enforces at compile time on a typed frame; the default is 'AnyFrame'. + +The frame is determined by the @input@ (via 'FrameFor'): an untyped target pairs +with a 'DataFrame' and yields the bare model; a typed target ('TExpr' over @cols@) +pairs with a @TypedDataFrame cols@ and yields a @Fitted cols model@, so 'predict' +gives a typed 'TExpr'. A non-'Double' typed column is then a compile error. +-} +class Fit cfg input where + {- | The model this config + input trains, computed as a type family so it is + visible in the instance head (and so the typed-lift instances below can + forward it) without needing a result annotation at the call site. + -} + type ModelOf cfg input :: Type + + type FrameReq cfg input :: FrameKind + type FrameReq cfg input = 'AnyFrame + fit :: + (CheckFrame (FrameReq cfg input) (FrameFor input)) => + cfg -> + input -> + FrameFor input -> + FitResult (FrameFor input) (ModelOf cfg input) + +{- | Lift any model fittable on an untyped target @Expr a@ to a typed target +@TExpr cols a@ over a @TypedDataFrame cols@, returning a schema-tagged 'Fitted'. +-} +instance (Fit cfg (Expr a)) => Fit cfg (TExpr cols a) where + type ModelOf cfg (TExpr cols a) = ModelOf cfg (Expr a) + type FrameReq cfg (TExpr cols a) = FrameReq cfg (Expr a) + fit cfg (TExpr e) tdf = Fitted (fit cfg e (thaw tdf)) + +-- | The same lift for the unsupervised feature-list inputs. +instance (Fit cfg [Expr Double]) => Fit cfg [TExpr cols Double] where + type ModelOf cfg [TExpr cols Double] = ModelOf cfg [Expr Double] + type FrameReq cfg [TExpr cols Double] = FrameReq cfg [Expr Double] + fit cfg feats tdf = Fitted (fit cfg (map unTExpr feats) (thaw tdf)) + +{- | Compile a fitted model's canonical prediction to an expression over the raw +columns. The 'Prediction' type tracks the model: a bare model gives @Expr r@; a +'Fitted' model (trained on a typed frame) gives @TExpr cols r@. +-} +class Predict model where + type Prediction model :: Type + predict :: model -> Prediction model + +instance (Predict model, ToTExpr cols (Prediction model)) => Predict (Fitted cols model) where + type Prediction (Fitted cols model) = AsTExpr cols (Prediction model) + predict (Fitted m) = toTExpr @cols (predict m) diff --git a/src/DataFrame/ModelSelection.hs b/src/DataFrame/ModelSelection.hs --- a/src/DataFrame/ModelSelection.hs +++ b/src/DataFrame/ModelSelection.hs @@ -2,10 +2,9 @@ fitters have heterogeneous types, so these helpers are parameterized by a user-supplied @train -> test -> score@ closure; the search maximizes the mean cross-validated score (use a negated error metric to minimize). Splitting reuses -the deterministic 'kFolds' / 'randomSplit' from @dataframe-operations@. +the deterministic 'kFolds' from @dataframe-operations@. -} module DataFrame.ModelSelection ( - trainTestSplit, crossValScore, crossValidate, GridSearchResult (..), @@ -20,13 +19,7 @@ import DataFrame.Internal.Expression (Expr) import DataFrame.Metrics (Metric, evaluate) import DataFrame.Operations.Merge () -import DataFrame.Operations.Subset (kFolds, randomSplit) - -{- | Split into @(train, test)@ with the given training fraction and seed -(deterministic). --} -trainTestSplit :: Double -> Int -> DataFrame -> (DataFrame, DataFrame) -trainTestSplit trainFrac seed = randomSplit (mkStdGen seed) trainFrac +import DataFrame.Operations.Subset (kFolds) {- | Per-fold scores from k-fold cross-validation. @scoreFn train test@ fits on the training rows and returns a score on the held-out fold. diff --git a/src/DataFrame/PCA.hs b/src/DataFrame/PCA.hs --- a/src/DataFrame/PCA.hs +++ b/src/DataFrame/PCA.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} {- | Principal component analysis via the symmetric Jacobi eigensolver on the covariance of the (optionally standardized) feature columns. 'fit' trains a @@ -9,6 +10,7 @@ 'pcaExprs' / 'pcaTransform' (PCA is a transformer, so it has no 'Predict'). -} module DataFrame.PCA ( + module DataFrame.Model, NComponents (..), PCAConfig (..), defaultPCAConfig, @@ -27,7 +29,7 @@ import DataFrame.Internal.Expression (Expr (..), UExpr (..)) import DataFrame.LinearAlgebra (gram) import DataFrame.LinearAlgebra.Eigen (jacobiEigenSym) -import DataFrame.Model (Fit (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.-.)) import DataFrame.Transform (Transform (..)) @@ -58,7 +60,8 @@ } deriving (Eq, Show) -instance Fit PCAConfig [Expr Double] PCAModel where +instance Fit PCAConfig [Expr Double] where + type ModelOf PCAConfig [Expr Double] = PCAModel fit = fitPCA -- | Fit PCA on the given feature columns (each must be a @Col@). diff --git a/src/DataFrame/PCA/Kernel.hs b/src/DataFrame/PCA/Kernel.hs --- a/src/DataFrame/PCA/Kernel.hs +++ b/src/DataFrame/PCA/Kernel.hs @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} {- | Kernel PCA with an RBF kernel, solved on a set of landmark points (Nyström). Exact kernel PCA when the landmark count covers every row, a principled @@ -8,6 +9,7 @@ 'kernelPCAExprs' / 'kernelPcaTransform' (a transformer, so no 'Predict'). -} module DataFrame.PCA.Kernel ( + module DataFrame.Model, KernelPCAConfig (..), defaultKernelPCAConfig, KernelPCAModel (..), @@ -25,7 +27,7 @@ import DataFrame.Internal.Expression (Expr (..), UExpr (..)) import DataFrame.LinearAlgebra (sqDist) import DataFrame.LinearAlgebra.Eigen (jacobiEigenSym) -import DataFrame.Model (Fit (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.-.)) import DataFrame.Random (mkGen, sampleIndices) import DataFrame.Transform (Transform (..)) @@ -60,7 +62,8 @@ } deriving (Eq, Show) -instance Fit KernelPCAConfig [Expr Double] KernelPCAModel where +instance Fit KernelPCAConfig [Expr Double] where + type ModelOf KernelPCAConfig [Expr Double] = KernelPCAModel fit = fitKernelPCA -- | Fit kernel PCA over the given feature columns. diff --git a/src/DataFrame/Random.hs b/src/DataFrame/Random.hs deleted file mode 100644 --- a/src/DataFrame/Random.hs +++ /dev/null @@ -1,121 +0,0 @@ -{-# LANGUAGE CPP #-} - -{- | Deterministic, platform-independent random sampling for the stochastic -fitters. Built on @random@'s SplitMix 'StdGen' (only 'genWord64' and a -version-bridged split are used, the stable surface across @random@ versions); -the distributions here are our own so a seeded fit is bit-reproducible on -Linux, macOS, and Windows. --} -module DataFrame.Random ( - Gen, - mkGen, - splitGen, - nextWord64, - nextDouble, - nextIntR, - gaussianPair, - gaussianVector, - shuffleInts, - sampleIndices, -) where - -import Control.Monad (forM_) -import Control.Monad.ST (runST) -import Data.Bits (shiftR) -import qualified Data.Vector.Unboxed as VU -import qualified Data.Vector.Unboxed.Mutable as VUM -import Data.Word (Word64) -import System.Random (StdGen, genWord64, mkStdGen) -import qualified System.Random as R - --- | The pure splittable generator. A fit is a function of @(seed, data)@. -type Gen = StdGen - --- | Seed a generator from an 'Int'. -mkGen :: Int -> Gen -mkGen = mkStdGen - --- | Split into two independent generators. -splitGen :: Gen -> (Gen, Gen) -#if MIN_VERSION_random(1,3,0) -splitGen = R.splitGen -#else -splitGen = R.split -#endif - --- | Raw 64-bit draw. -nextWord64 :: Gen -> (Word64, Gen) -nextWord64 = genWord64 - --- | Uniform 'Double' in @[0, 1)@ from the top 53 bits (exact mantissa). -nextDouble :: Gen -> (Double, Gen) -nextDouble g = - let (w, g') = genWord64 g - d = fromIntegral (w `shiftR` 11) * (1 / 9007199254740992) - in (d, g') - -{- | Uniform 'Int' in the inclusive range @[lo, hi]@ by rejection sampling -(unbiased). Returns @lo@ when @hi <= lo@. --} -nextIntR :: (Int, Int) -> Gen -> (Int, Gen) -nextIntR (lo, hi) g - | hi <= lo = (lo, g) - | otherwise = loop g - where - range = fromIntegral (hi - lo + 1) :: Word64 - threshold = negate range `mod` range - loop gg = - let (w, gg') = genWord64 gg - in if w >= threshold - then (lo + fromIntegral (w `mod` range), gg') - else loop gg' - -{- | A pair of independent standard normals via Box-Muller, consuming exactly two -uniforms so stream offsets stay data-independent. --} -gaussianPair :: Gen -> ((Double, Double), Gen) -gaussianPair g = - let (u1, g1) = nextDouble g - (u2, g2) = nextDouble g1 - u1' = if u1 <= 0 then 2.220446049250313e-16 else u1 - r = sqrt (-(2 * log u1')) - a = 2 * pi * u2 - in ((r * cos a, r * sin a), g2) - --- | A length-@n@ vector of standard normals. -gaussianVector :: Int -> Gen -> (VU.Vector Double, Gen) -gaussianVector n g0 = go n g0 [] - where - go k g acc - | k <= 0 = (VU.fromList (take n (reverse acc)), g) - | otherwise = - let ((z0, z1), g') = gaussianPair g - in go (k - 2) g' (z1 : z0 : acc) - -{- | A uniformly random permutation of @[0 .. n-1]@ (Fisher-Yates), threading the -generator purely. --} -shuffleInts :: Int -> Gen -> (VU.Vector Int, Gen) -shuffleInts n g0 - | n <= 1 = (VU.enumFromN 0 (max 0 n), g0) - | otherwise = - let (swaps, g1) = genSwaps (n - 1) g0 [] - v = runST $ do - m <- VU.thaw (VU.enumFromN 0 n) - forM_ swaps $ uncurry (VUM.swap m) - VU.freeze m - in (v, g1) - where - genSwaps i g acc - | i < 1 = (reverse acc, g) - | otherwise = - let (j, g') = nextIntR (0, i) g - in genSwaps (i - 1) g' ((i, j) : acc) - -{- | @sampleIndices k n@ draws @k@ distinct indices from @[0 .. n-1]@ (the first -@k@ of a full shuffle); returns all @n@ when @k >= n@. --} -sampleIndices :: Int -> Int -> Gen -> (VU.Vector Int, Gen) -sampleIndices k n g = - let (perm, g') = shuffleInts n g - in (VU.take (min k n) perm, g') diff --git a/src/DataFrame/SVM.hs b/src/DataFrame/SVM.hs --- a/src/DataFrame/SVM.hs +++ b/src/DataFrame/SVM.hs @@ -1,18 +1,21 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} -{- | Linear support vector classification: L2-regularized squared hinge fitted -with the FISTA engine (sklearn's LinearSVC default loss). 'fit' trains a -one-vs-rest 'LinearSVCModel'; 'predict' is the arg-max class margin. There is no -@predict_proba@, matching sklearn's LinearSVC. +{- | Linear support vector classification: L2-regularized squared hinge via +FISTA (sklearn's LinearSVC default). 'fit' trains a one-vs-rest 'LinearSVCModel'; +'predict' is the arg-max class margin (no @predict_proba@, as in sklearn). -} module DataFrame.SVM ( + module DataFrame.Model, LinearSVCModel (..), SVCConfig (..), defaultSVCConfig, svcMarginExprs, + -- | Surfaced by @LinearSVCModel.svcModels@. + LinearModel (..), ) where import Data.List (sort) @@ -32,7 +35,7 @@ import DataFrame.Internal.Expression (Expr) import DataFrame.LinearSolver (LinearModel (..), SolverConfig (..), fitProx) import DataFrame.LinearSolver.Loss (sqHingeLoss) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model -- | Hyper-parameters. @svcC@ is the inverse regularization strength (sklearn @C@). data SVCConfig = SVCConfig @@ -52,10 +55,12 @@ } deriving (Eq, Show) -instance (Columnable a, Ord a) => Fit SVCConfig (Expr a) (LinearSVCModel a) where +instance (Columnable a, Ord a) => Fit SVCConfig (Expr a) where + type ModelOf SVCConfig (Expr a) = (LinearSVCModel a) fit = fitLinearSVC -instance (Columnable a, Ord a) => Predict (LinearSVCModel a) a where +instance (Columnable a, Ord a) => Predict (LinearSVCModel a) where + type Prediction (LinearSVCModel a) = Expr a predict m = argMaxExpr (labelledMargins m) -- | Fit a one-vs-rest linear SVC. diff --git a/src/DataFrame/SVM/RFF.hs b/src/DataFrame/SVM/RFF.hs --- a/src/DataFrame/SVM/RFF.hs +++ b/src/DataFrame/SVM/RFF.hs @@ -3,6 +3,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} {- | Approximate RBF-kernel SVM via Random Fourier Features (Rahimi & Recht): map @@ -11,6 +12,7 @@ @Σ_r β_r·cos(…)@ expression of size @O(D·d)@, independent of the row count. -} module DataFrame.SVM.RFF ( + module DataFrame.Model, RFFConfig (..), defaultRFFConfig, RFFSVMModel (..), @@ -29,7 +31,7 @@ import DataFrame.LinearAlgebra (dot) import DataFrame.LinearSolver (LinearModel (..), SolverConfig (..), fitProx) import DataFrame.LinearSolver.Loss (sqHingeLoss) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operators ((.*.), (.+.), (.>.)) import DataFrame.Random (Gen, gaussianVector, mkGen, nextDouble) @@ -69,10 +71,12 @@ } deriving (Show) -instance (Columnable a, Ord a) => Fit RFFConfig (Expr a) (RFFSVMModel a) where +instance (Columnable a, Ord a) => Fit RFFConfig (Expr a) where + type ModelOf RFFConfig (Expr a) = (RFFSVMModel a) fit = fitRFFSVM -instance (Columnable a) => Predict (RFFSVMModel a) a where +instance (Columnable a) => Predict (RFFSVMModel a) where + type Prediction (RFFSVMModel a) = Expr a predict m = If (margin .>. F.lit 0) (Lit (rffPosClass m)) (Lit (rffNegClass m)) where diff --git a/src/DataFrame/Segmented.hs b/src/DataFrame/Segmented.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Segmented.hs @@ -0,0 +1,359 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# 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 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.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 + { segBase :: !cfg + , segOn :: !(Maybe [T.Text]) + , segMaxCard :: !Int + , segMinRows :: !Int + , segPool :: !Double + } + deriving (Eq, Show) + +-- | Wrap a base config with the defaults: auto-detect, cap 32, min 30 rows, no pooling. +segmented :: cfg -> Segmented cfg +segmented base = Segmented base Nothing 32 30 0 + +-- | Segment only on the named columns (each must be Text), overriding auto-detect. +segmentOn :: Segmented cfg -> [T.Text] -> Segmented cfg +segmentOn s cols = s{segOn = Just cols} + +-- | Set the pooling strength @λ@ (shrink segments toward the reference). +pooled :: Segmented cfg -> Double -> Segmented cfg +pooled s lam = s{segPool = lam} + +-- | One fitted segment: its categorical key, row count, and base model. +data Segment model = Segment + { segKey :: ![T.Text] + , segN :: !Int + , segModel :: !model + } + deriving (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 + { smCatCols :: ![T.Text] + , smSegments :: ![Segment model] + , smFellBack :: ![([T.Text], Int)] + , smFallback :: !model + , smExpr :: !(Expr a) + } + deriving (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 lam target dfs + | lam == 0 = map (fit cfg target) dfs + | otherwise = + error + "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 = fitSegmented + +instance Predict (SegmentedModel a model) where + type Prediction (SegmentedModel a model) = Expr a + predict = 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 (Segmented base mcols maxCard minRows lam) target df = + seq (guardNumeric df textCols target) result + where + mTarget = case target of + Col n -> Just n + _ -> Nothing + feats = featureNames target df + textCols = [c | c <- feats, isTextCol df c] + catCols = resolveCatCols df mTarget textCols mcols maxCard + numericFrame = exclude textCols + globalM = fit base target (numericFrame df) + result + | null catCols = + SegmentedModel [] [] [] globalM (predict globalM) + | otherwise = + let d = length (feats \\ textCols) + floor' = max minRows (d + 1) + grouped = groupByKey df catCols + (qualifying, undersized) = + span' (\(_, ixs) -> VU.length ixs >= floor') grouped + qualFrames = + [numericFrame (rowsAtIndices ixs df) | (_, ixs) <- qualifying] + segModels = fitSegments base lam target qualFrames + segments = + zipWith + (\(k, ixs) m -> Segment k (VU.length ixs) m) + qualifying + segModels + fellBack = [(k, VU.length ixs) | (k, ixs) <- undersized] + expr = buildExpr catCols segments globalM + in SegmentedModel catCols segments fellBack globalM expr + +-- | 'span' over a predicate that need not hold contiguously (a filter partition). +span' :: (b -> Bool) -> [b] -> ([b], [b]) +span' p xs = (filter p xs, filter (not . p) 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 catCols segs fallback = + foldr + (\(Segment key _ m) acc -> If (keyCond catCols key) (predict m) acc) + (predict fallback) + segs + +-- | @col1 == v1 && col2 == v2 && ...@ for a segment's key. +keyCond :: [T.Text] -> [T.Text] -> Expr Bool +keyCond catCols vals = + foldr1 (.&&.) [(Col c :: Expr T.Text) .==. Lit v | (c, v) <- zip catCols 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 df mTarget textCols mcols maxCard = case mcols of + Just cols -> + let bad = filter (\c -> not (isTextCol df c) || Just c == mTarget) cols + in if null bad + then cols + else + error + ( "Segmented: segmentOn columns must be Text features (not the target); invalid: " + ++ show bad + ) + Nothing -> [c | c <- textCols, distinctCount df c <= maxCard] + +isTextCol :: DataFrame -> T.Text -> Bool +isTextCol df c = hasElemType @T.Text (unsafeGetColumn c df) + +distinctCount :: DataFrame -> T.Text -> Int +distinctCount df c = + Set.size (Set.fromList (V.toList (columnToTextVec (unsafeGetColumn c 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 df textCols target = + case problems of + [] -> () + ps -> + error + ( "Segmented: unusable feature column(s):\n" + ++ unlines (map fmt ps) + ) + where + problems = + [ (c, r) + | c <- featureNames target df \\ textCols + , Just r <- [reason (unsafeGetColumn c df)] + ] + fmt (c, r) = " " ++ T.unpack c ++ ": " ++ r + reason col + | isJust (columnBitmap col) = + Just + "has missing values — drop them (filterJust / filterAllJust) or model missingness explicitly; imputing risks train/inference skew" + | isIntegralCol col = + Just "is an integer column — cast to Double with F.toDouble" + | not (hasElemType @Double col) = + Just + "is not Double — convert to Double (numeric) or segment on it (categorical)" + | otherwise = Nothing + isIntegralCol col = case col of + UnboxedColumn _ (_ :: VU.Vector b) -> case sIntegral @b of + STrue -> True + _ -> False + _ -> 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 df catCols = + map (\(k, is) -> (k, VU.fromList (reverse is))) (M.toAscList grouped) + where + n = nRows df + cols = map (`unsafeGetColumn` df) catCols + textVecs = map columnToTextVec cols + bitmaps = map columnBitmap cols + validRow i = all (maybe True (`bitmapTestBit` i)) bitmaps + keyOf i = [tv V.! i | tv <- textVecs] + grouped = + foldl' + (\m i -> if validRow i then M.insertWith (++) (keyOf i) [i] m else m) + M.empty + [0 .. n - 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 cfg lam target dfs + | lam == 0 = map (fit cfg target) dfs + | null dfs = [] + | otherwise = shrinkLinear cfg lam target dfs + +shrinkLinear :: + LinearConfig -> Double -> Expr Double -> [DataFrame] -> [LinearRegressor] +shrinkLinear cfg lam target dfs = map toReg dsegs + where + names = case dfs of + (d0 : _) -> featureNames target d0 + [] -> error "shrinkLinear: no segments" + d = length names + mats = [snd (numericMatrix names dframe) | dframe <- dfs] + ys = [targetDoubles target dframe | dframe <- dfs] + ns = map V.length mats + pooledRows = V.concat mats + nP = V.length pooledRows + means = + VU.generate d $ \j -> + sum [(pooledRows V.! i) VU.! j | i <- [0 .. nP - 1]] / fromIntegral nP + sds = + VU.generate d $ \j -> + let m = means VU.! j + var = + sum [sq ((pooledRows V.! i) VU.! j - m) | i <- [0 .. nP - 1]] + / fromIntegral nP + in sqrt var + stdValue j x = let s = sds VU.! j in if s == 0 then 0 else (x - means VU.! j) / s + augStd row = VU.cons 1 (VU.imap stdValue row) + zMats = [V.map augStd m | m <- mats] + olsAll = zipWith fitOLS zMats ys + fitOLS z y = either (const Nothing) Just (qrLeastSquares z y) + good = [(n, sol) | (n, Just sol) <- zip ns olsAll] + totW = fromIntegral (sum (map fst good)) :: Double + dref + | null good = VU.replicate (d + 1) 0 + | otherwise = + VU.generate (d + 1) $ \k -> + sum [fromIntegral n * (sol VU.! k) | (n, sol) <- good] / totW + dsegs = zipWith solveSeg zMats ys + solveSeg z y = + let r = VU.zipWith (-) y (matVec z dref) + a = addDiagonal lam (gram z) + rhs = tMatVec z r + in case choleskySolve a rhs of + Just eta -> VU.zipWith (+) dref eta + Nothing -> dref + toReg dseg = + let bStd = dseg VU.! 0 + wStd = VU.drop 1 dseg + rawCoef = VU.imap (\j w -> let s = sds VU.! j in if s == 0 then 0 else w / s) wStd + adj = + sum + [ let s = sds VU.! j + in if s == 0 then 0 else (wStd VU.! j) * (means VU.! j) / s + | j <- [0 .. d - 1] + ] + in LinearRegressor rawCoef (bStd - adj) (V.fromList names) (lcPenalty cfg) + +sq :: Double -> Double +sq x = x * x + +-- | Add @lam@ to the diagonal of a square matrix. +addDiagonal :: Double -> Matrix -> Matrix +addDiagonal lam = V.imap (\i row -> row VU.// [(i, (row VU.! i) + lam)]) diff --git a/src/DataFrame/SymbolicRegression.hs b/src/DataFrame/SymbolicRegression.hs --- a/src/DataFrame/SymbolicRegression.hs +++ b/src/DataFrame/SymbolicRegression.hs @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {- | Symbolic regression by genetic programming (modelled on the @symbolic-regression@ library, ported dependency-light: no e-graphs, no NLOPT). @@ -8,6 +9,7 @@ accuracy-vs-complexity Pareto front. Deterministic given the seed. -} module DataFrame.SymbolicRegression ( + module DataFrame.Model, UnOp (..), SRConfig (..), defaultSRConfig, @@ -21,7 +23,7 @@ import DataFrame.Featurize.Internal (featureNames, targetDoubles) import DataFrame.Internal.DataFrame (DataFrame) import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operations.Core (columnAsDoubleVector) import DataFrame.Random (mkGen) import DataFrame.SymbolicRegression.Expr ( @@ -71,10 +73,12 @@ , srGenerationsRun :: !Int } -instance Fit SRConfig (Expr Double) SRModel where +instance Fit SRConfig (Expr Double) where + type ModelOf SRConfig (Expr Double) = SRModel fit = fitSymbolicRegression -instance Predict SRModel Double where +instance Predict SRModel where + type Prediction SRModel = Expr Double predict = srBest -- | Search for an expression predicting @target@ from the other columns. diff --git a/src/DataFrame/SymbolicRegression/Expr.hs b/src/DataFrame/SymbolicRegression/Expr.hs deleted file mode 100644 --- a/src/DataFrame/SymbolicRegression/Expr.hs +++ /dev/null @@ -1,122 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} - -{- | The symbolic-regression expression tree: a small first-order ADT (no hegg, -no 'Fix'). Vectorized evaluation over a feature matrix and a total translation -to a dataframe 'Expr Double' (the SR result IS a dataframe expression). Division, -log, and sqrt are protected so evaluation never produces @NaN@. --} -module DataFrame.SymbolicRegression.Expr ( - SRExpr (..), - BinOp (..), - UnOp (..), - evalSR, - toDataFrameExpr, - srSize, - constants, - setConstants, - allBinOps, - allUnOps, -) where - -import qualified Data.Text as T -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -import qualified DataFrame.Functions as F -import DataFrame.Internal.Expression (Expr (..)) -import DataFrame.Operators ((.*.), (.+.), (.-.), (./.)) - -data BinOp = SAdd | SSub | SMul | SDiv - deriving (Eq, Ord, Show, Enum, Bounded) - -data UnOp = SNeg | SSin | SCos | SExp | SLog | SSqrt - deriving (Eq, Ord, Show, Enum, Bounded) - --- | A symbolic-regression expression over feature variables and constants. -data SRExpr - = SVar !Int - | SConst !Double - | SUn !UnOp SRExpr - | SBin !BinOp SRExpr SRExpr - deriving (Eq, Ord, Show) - -allBinOps :: [BinOp] -allBinOps = [minBound .. maxBound] - -allUnOps :: [UnOp] -allUnOps = [minBound .. maxBound] - -{- | Evaluate over a feature matrix given column-major (@feats ! j@ is feature -@j@ across all rows). Protected operators keep results finite. --} -evalSR :: V.Vector (VU.Vector Double) -> Int -> SRExpr -> VU.Vector Double -evalSR feats n = go - where - go (SVar j) - | j < V.length feats = feats V.! j - | otherwise = VU.replicate n 0 - go (SConst c) = VU.replicate n c - go (SUn op e) = VU.map (unFn op) (go e) - go (SBin op a b) = VU.zipWith (binFn op) (go a) (go b) - -binFn :: BinOp -> Double -> Double -> Double -binFn SAdd a b = a + b -binFn SSub a b = a - b -binFn SMul a b = a * b -binFn SDiv a b = if abs b < 1e-9 then 1 else a / b - -unFn :: UnOp -> Double -> Double -unFn SNeg = negate -unFn SSin = sin -unFn SCos = cos -unFn SExp = exp . min 50 -unFn SLog = \x -> log (abs x + 1e-9) -unFn SSqrt = sqrt . abs - --- | Translate to a dataframe expression over the named feature columns. -toDataFrameExpr :: V.Vector T.Text -> SRExpr -> Expr Double -toDataFrameExpr names = go - where - go (SVar j) - | j < V.length names = Col (names V.! j) - | otherwise = F.lit 0 - go (SConst c) = F.lit c - go (SUn op e) = unExpr op (go e) - go (SBin op a b) = binExpr op (go a) (go b) - unExpr SNeg = negate - unExpr SSin = sin - unExpr SCos = cos - unExpr SExp = exp - unExpr SLog = log - unExpr SSqrt = sqrt - binExpr SAdd = (.+.) - binExpr SSub = (.-.) - binExpr SMul = (.*.) - binExpr SDiv = (./.) - -srSize :: SRExpr -> Int -srSize (SVar _) = 1 -srSize (SConst _) = 1 -srSize (SUn _ e) = 1 + srSize e -srSize (SBin _ a b) = 1 + srSize a + srSize b - --- | The constant values in left-to-right traversal order. -constants :: SRExpr -> [Double] -constants (SConst c) = [c] -constants (SVar _) = [] -constants (SUn _ e) = constants e -constants (SBin _ a b) = constants a ++ constants b - --- | Replace the constants in traversal order; extra values are ignored. -setConstants :: [Double] -> SRExpr -> SRExpr -setConstants vals e = fst (go vals e) - where - go vs (SConst _) = case vs of - (v : rest) -> (SConst v, rest) - [] -> (SConst 0, []) - go vs (SVar j) = (SVar j, vs) - go vs (SUn op a) = let (a', vs') = go vs a in (SUn op a', vs') - go vs (SBin op a b) = - let (a', vs') = go vs a - (b', vs'') = go vs' b - in (SBin op a' b', vs'') diff --git a/src/DataFrame/SymbolicRegression/GP.hs b/src/DataFrame/SymbolicRegression/GP.hs deleted file mode 100644 --- a/src/DataFrame/SymbolicRegression/GP.hs +++ /dev/null @@ -1,207 +0,0 @@ -{- | A compact generational genetic-programming search over 'SRExpr': ramped -random initialization, tournament selection, subtree crossover and mutation, -elitism, and a complexity-keyed Pareto archive. Deterministic given the seed -(the splitmix generator from "DataFrame.Random"). --} -module DataFrame.SymbolicRegression.GP ( - GPParams (..), - runGP, -) where - -import Data.List (foldl', minimumBy, sortBy) -import qualified Data.Map.Strict as M -import Data.Ord (comparing) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -import DataFrame.Random (Gen, nextDouble, nextIntR) -import DataFrame.SymbolicRegression.Expr -import DataFrame.SymbolicRegression.Optimize ( - meanSquaredError, - optimizeConstants, - ) -import DataFrame.SymbolicRegression.Simplify (simplify) - --- | GP hyper-parameters resolved from the public config. -data GPParams = GPParams - { gpFeats :: !(V.Vector (VU.Vector Double)) - , gpN :: !Int - , gpTarget :: !(VU.Vector Double) - , gpNVars :: !Int - , gpUnOps :: ![UnOp] - , gpPopSize :: !Int - , gpGenerations :: !Int - , gpMaxSize :: !Int - , gpTournament :: !Int - , gpCrossoverP :: !Double - , gpMutationP :: !Double - , gpOptimizeP :: !Double - , gpParsimony :: !Double - } - -type Scored = (SRExpr, Double) - --- | Run the search; returns @(best, pareto front, generations run)@. -runGP :: GPParams -> Gen -> (SRExpr, [(Int, Double, SRExpr)], Int) -runGP p g0 = - let (pop0, g1) = initPop p g0 - scored0 = map (scoreOf p) pop0 - arch0 = foldl' (archiveInsert p) M.empty scored0 - (_, finalArch, gN, _) = - iterate' 0 scored0 arch0 g1 - best = bestOfArchive finalArch - front = - [ (sz, mse, e) - | (sz, (mse, e)) <- M.toList finalArch - ] - in (snd3 best, sortBy (comparing fst3) front, gN) - where - iterate' gen pop arch g - | gen >= gpGenerations p = (pop, arch, gen, g) - | otherwise = - let (pop', g') = nextGen p pop g - arch' = foldl' (archiveInsert p) arch pop' - in iterate' (gen + 1) pop' arch' g' - fst3 (a, _, _) = a - snd3 (_, b, _) = b - bestOfArchive arch = - case M.toList arch of - [] -> (0 :: Int, SConst 0, 1 / 0) - xs -> - let (sz, (mse, e)) = minimumBy (comparing (fst . snd)) xs - in (sz, e, mse) - -scoreOf :: GPParams -> SRExpr -> Scored -scoreOf p e = (e, meanSquaredError (gpFeats p) (gpN p) (gpTarget p) e) - -fitness :: GPParams -> Scored -> Double -fitness p (e, mse) = mse + gpParsimony p * fromIntegral (srSize e) - -archiveInsert :: - GPParams -> M.Map Int (Double, SRExpr) -> Scored -> M.Map Int (Double, SRExpr) -archiveInsert _ arch (e, mse) - | isNaN mse || isInfinite mse = arch - | otherwise = - let key = srSize (simplify e) - in M.insertWith better key (mse, e) arch - where - better newv@(m1, _) oldv@(m2, _) = if m1 < m2 then newv else oldv - -initPop :: GPParams -> Gen -> ([SRExpr], Gen) -initPop p = go (gpPopSize p) [] - where - go 0 acc g = (acc, g) - go k acc g = - let (depth, g1) = nextIntR (1, 4) g - (e, g2) = randomExpr p depth g1 - in go (k - 1) (e : acc) g2 - -randomExpr :: GPParams -> Int -> Gen -> (SRExpr, Gen) -randomExpr p depth g - | depth <= 1 = randomLeaf p g - | otherwise = - let (r, g1) = nextDouble g - in if r < 0.3 - then randomLeaf p g1 - else - let (isUn, g2) = nextDouble g1 - in if isUn < 0.3 && not (null (gpUnOps p)) - then - let (oi, g3) = nextIntR (0, length (gpUnOps p) - 1) g2 - (e, g4) = randomExpr p (depth - 1) g3 - in (SUn (gpUnOps p !! oi) e, g4) - else - let (oi, g3) = nextIntR (0, length allBinOps - 1) g2 - (a, g4) = randomExpr p (depth - 1) g3 - (b, g5) = randomExpr p (depth - 1) g4 - in (SBin (allBinOps !! oi) a b, g5) - -randomLeaf :: GPParams -> Gen -> (SRExpr, Gen) -randomLeaf p g = - let (r, g1) = nextDouble g - in if r < 0.6 && gpNVars p > 0 - then let (j, g2) = nextIntR (0, gpNVars p - 1) g1 in (SVar j, g2) - else let (c, g2) = nextDouble g1 in (SConst (c * 4 - 2), g2) - -nextGen :: GPParams -> [Scored] -> Gen -> ([Scored], Gen) -nextGen p pop g0 = - let elite = minimumBy (comparing (fitness p)) pop - (rest, g1) = go (gpPopSize p - 1) [] g0 - in (elite : rest, g1) - where - go 0 acc g = (acc, g) - go k acc g = - let (child, g') = breed p pop g - scored = optimizeMaybe p child g' - in go (k - 1) (fst scored : acc) (snd scored) - -optimizeMaybe :: GPParams -> SRExpr -> Gen -> (Scored, Gen) -optimizeMaybe p e g = - let (r, g1) = nextDouble g - e' = - if r < gpOptimizeP p - then optimizeConstants (gpFeats p) (gpN p) (gpTarget p) 15 e - else e - in (scoreOf p e', g1) - -breed :: GPParams -> [Scored] -> Gen -> (SRExpr, Gen) -breed p pop g0 = - let (pa, g1) = tournament p pop g0 - (doX, g2) = nextDouble g1 - (child, g3) = - if doX < gpCrossoverP p - then - let (pb, g2') = tournament p pop g2 - (c, g3') = crossover pa pb g2' - in (c, g3') - else (pa, g2) - (doM, g4) = nextDouble g3 - (child', g5) = - if doM < gpMutationP p then mutate p child g4 else (child, g4) - capped = if srSize child' > gpMaxSize p then pa else child' - in (simplify capped, g5) - -tournament :: GPParams -> [Scored] -> Gen -> (SRExpr, Gen) -tournament p pop g0 = - let (picks, g1) = pickN (gpTournament p) g0 - chosen = map (pop !!) picks - in (fst (minimumBy (comparing (fitness p)) chosen), g1) - where - n = length pop - pickN 0 g = ([], g) - pickN k g = - let (i, g') = nextIntR (0, n - 1) g - (is, g'') = pickN (k - 1) g' - in (i : is, g'') - -crossover :: SRExpr -> SRExpr -> Gen -> (SRExpr, Gen) -crossover a b g0 = - let (ia, g1) = nextIntR (0, srSize a - 1) g0 - (ib, g2) = nextIntR (0, srSize b - 1) g1 - sub = subtreeAt ib b - in (replaceAt ia a sub, g2) - -mutate :: GPParams -> SRExpr -> Gen -> (SRExpr, Gen) -mutate p e g0 = - let (i, g1) = nextIntR (0, srSize e - 1) g0 - (depth, g2) = nextIntR (1, 3) g1 - (newSub, g3) = randomExpr p depth g2 - in (replaceAt i e newSub, g3) - -subtreeAt :: Int -> SRExpr -> SRExpr -subtreeAt 0 e = e -subtreeAt i (SUn _ e) = subtreeAt (i - 1) e -subtreeAt i (SBin _ a b) = - let sa = srSize a - in if i <= sa then subtreeAt (i - 1) a else subtreeAt (i - 1 - sa) b -subtreeAt _ e = e - -replaceAt :: Int -> SRExpr -> SRExpr -> SRExpr -replaceAt 0 _ new = new -replaceAt i (SUn op e) new = SUn op (replaceAt (i - 1) e new) -replaceAt i (SBin op a b) new = - let sa = srSize a - in if i <= sa - then SBin op (replaceAt (i - 1) a new) b - else SBin op a (replaceAt (i - 1 - sa) b new) -replaceAt _ e _ = e diff --git a/src/DataFrame/SymbolicRegression/Optimize.hs b/src/DataFrame/SymbolicRegression/Optimize.hs deleted file mode 100644 --- a/src/DataFrame/SymbolicRegression/Optimize.hs +++ /dev/null @@ -1,68 +0,0 @@ -{-# LANGUAGE BangPatterns #-} - -{- | Constant optimization for symbolic-regression candidates: finite-difference -gradient descent with backtracking line search on the embedded constants. Pure -and dependency-free; effective at the one-to-few constants a tree carries, which -is where random restarts plus a quasi-Newton step would otherwise be used. --} -module DataFrame.SymbolicRegression.Optimize ( - optimizeConstants, - meanSquaredError, -) where - -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as VU - -import DataFrame.SymbolicRegression.Expr ( - SRExpr, - constants, - evalSR, - setConstants, - ) - --- | Mean squared error of an expression's predictions against the target. -meanSquaredError :: - V.Vector (VU.Vector Double) -> Int -> VU.Vector Double -> SRExpr -> Double -meanSquaredError feats n target e = - let pred = evalSR feats n e - diff = VU.zipWith (-) pred target - in VU.sum (VU.map (\x -> x * x) diff) / fromIntegral (max 1 n) - --- | Refine an expression's constants to reduce MSE (no-op when constant-free). -optimizeConstants :: - V.Vector (VU.Vector Double) -> - Int -> - VU.Vector Double -> - Int -> - SRExpr -> - SRExpr -optimizeConstants feats n target iters expr - | null theta0 = expr - | otherwise = setConstants (descend iters theta0) expr - where - theta0 = constants expr - eps = 1e-6 - mseAt theta = meanSquaredError feats n target (setConstants theta expr) - descend 0 theta = theta - descend k theta = - let f0 = mseAt theta - g = numGrad theta - gn = sqrt (sum (map (\x -> x * x) g)) - in if gn < 1e-10 - then theta - else - let theta' = lineSearch theta g f0 - in if theta' == theta then theta else descend (k - 1) theta' - numGrad theta = - [ (mseAt (bump i eps theta) - mseAt (bump i (negate eps) theta)) / (2 * eps) - | i <- [0 .. length theta - 1] - ] - bump i delta theta = - [if j == i then t + delta else t | (j, t) <- zip [0 ..] theta] - lineSearch theta g f0 = go (1.0 :: Double) - where - go !step - | step < 1e-8 = theta - | otherwise = - let theta' = zipWith (\t gi -> t - step * gi) theta g - in if mseAt theta' < f0 then theta' else go (step / 2) diff --git a/src/DataFrame/SymbolicRegression/Simplify.hs b/src/DataFrame/SymbolicRegression/Simplify.hs deleted file mode 100644 --- a/src/DataFrame/SymbolicRegression/Simplify.hs +++ /dev/null @@ -1,65 +0,0 @@ -{- | A fuel-bounded, deterministic algebraic simplifier — the dependency-light -stand-in for equality saturation. It is used as a canonical dedup key for the -Pareto archive and to tidy reported expressions; no confluence is claimed, only -that it is a total, size-non-increasing, idempotent function whose rewrites -preserve evaluation. --} -module DataFrame.SymbolicRegression.Simplify ( - simplify, -) where - -import DataFrame.SymbolicRegression.Expr (BinOp (..), SRExpr (..), UnOp (..)) - --- | Simplify to a fixed point (bounded by a fuel counter). -simplify :: SRExpr -> SRExpr -simplify = go (10 :: Int) - where - go 0 e = e - go fuel e = - let e' = step e - in if e' == e then e else go (fuel - 1) e' - -step :: SRExpr -> SRExpr -step (SUn op e) = simplifyUn op (step e) -step (SBin op a b) = simplifyBin op (step a) (step b) -step e = e - -simplifyUn :: UnOp -> SRExpr -> SRExpr -simplifyUn SNeg (SUn SNeg e) = e -simplifyUn op (SConst c) = SConst (foldUn op c) -simplifyUn op e = SUn op e - -simplifyBin :: BinOp -> SRExpr -> SRExpr -> SRExpr -simplifyBin op (SConst a) (SConst b) = SConst (foldBin op a b) -simplifyBin SAdd a (SConst 0) = a -simplifyBin SAdd (SConst 0) b = b -simplifyBin SSub a (SConst 0) = a -simplifyBin SSub a b | a == b = SConst 0 -simplifyBin SMul _ (SConst 0) = SConst 0 -simplifyBin SMul (SConst 0) _ = SConst 0 -simplifyBin SMul a (SConst 1) = a -simplifyBin SMul (SConst 1) b = b -simplifyBin SDiv a (SConst 1) = a -simplifyBin SDiv a b | a == b = SConst 1 -simplifyBin op a b - | commutative op && a > b = SBin op b a - | otherwise = SBin op a b - -commutative :: BinOp -> Bool -commutative SAdd = True -commutative SMul = True -commutative _ = False - -foldBin :: BinOp -> Double -> Double -> Double -foldBin SAdd a b = a + b -foldBin SSub a b = a - b -foldBin SMul a b = a * b -foldBin SDiv a b = if abs b < 1e-9 then 1 else a / b - -foldUn :: UnOp -> Double -> Double -foldUn SNeg = negate -foldUn SSin = sin -foldUn SCos = cos -foldUn SExp = exp . min 50 -foldUn SLog = \x -> log (abs x + 1e-9) -foldUn SSqrt = sqrt . abs diff --git a/src/DataFrame/Synthesis.hs b/src/DataFrame/Synthesis.hs --- a/src/DataFrame/Synthesis.hs +++ b/src/DataFrame/Synthesis.hs @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} {- | Feature synthesis by bottom-up enumerative search with observational equivalence — the canonical enumerative method from Solar-Lezama's @@ -39,6 +40,7 @@ for very large frames, and piecewise (condition-abduction) features. -} module DataFrame.Synthesis ( + module DataFrame.Model, LossFunction (..), SynthesisConfig (..), defaultSynthesisConfig, @@ -67,7 +69,7 @@ percentile', variance', ) -import DataFrame.Model (Fit (..), Predict (..)) +import DataFrame.Model import DataFrame.Operations.Core (columnAsDoubleVector) -- | How a candidate's output column is scored against the target (higher is better). @@ -110,10 +112,12 @@ , sfFeatures :: ![(Expr Double, Double)] } -instance Fit SynthesisConfig (Expr Double) SynthesizedFeature where +instance Fit SynthesisConfig (Expr Double) where + type ModelOf SynthesisConfig (Expr Double) = SynthesizedFeature fit = synthesizeFeatures -instance Predict SynthesizedFeature Double where +instance Predict SynthesizedFeature where + type Prediction SynthesizedFeature = Expr Double predict = sfExpr -- | A candidate's evaluated column over the example rows. diff --git a/src/DataFrame/Transform/Serialize.hs b/src/DataFrame/Transform/Serialize.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Transform/Serialize.hs @@ -0,0 +1,42 @@ +{- | Persist and reload a fitted 'Transform'. + +A 'Transform' is just an ordered list of named output expressions, so it +serializes through the same JSON wire format as any pipeline (see +"DataFrame.Expr.Serialize"). Save a fitted preprocessing/feature transform in one +process and reload it in another with 'applyTransform' to run inference: + +> Right t <- loadTransformFromFile "scaler.json" +> let scored = applyTransform t newData +-} +module DataFrame.Transform.Serialize ( + encodeTransform, + decodeTransform, + saveTransformToFile, + loadTransformFromFile, +) where + +import qualified Data.Aeson as Aeson + +import DataFrame.Expr.Serialize ( + decodeNamedExprs, + encodeNamedExprs, + loadPipelineFromFile, + savePipelineToFile, + ) +import DataFrame.Transform (Transform (..)) + +-- | Encode a transform's output expressions to JSON. +encodeTransform :: Transform -> Either String Aeson.Value +encodeTransform = encodeNamedExprs . transformOutputs + +-- | Decode a transform produced by 'encodeTransform'. +decodeTransform :: Aeson.Value -> Either String Transform +decodeTransform = fmap Transform . decodeNamedExprs + +-- | Encode a transform and write it to a file. No file is written on failure. +saveTransformToFile :: FilePath -> Transform -> IO (Either String ()) +saveTransformToFile fp = savePipelineToFile fp . transformOutputs + +-- | Load a transform produced by 'saveTransformToFile'. +loadTransformFromFile :: FilePath -> IO (Either String Transform) +loadTransformFromFile fp = fmap (fmap Transform) (loadPipelineFromFile fp) diff --git a/tests-internal/Cart.hs b/tests-internal/Cart.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Cart.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | Agreement tests: 'buildCartTree' must predict identically to sklearn's +@DecisionTreeClassifier(random_state=0, max_depth=4)@ on shared folds (golden +fixtures from @bench/export_cart_fixtures.py@); missing fixtures SKIP. +-} +module Cart (tests) where + +import Control.Exception (SomeException, try) +import Data.Aeson (FromJSON (..), eitherDecode, withObject, (.:)) +import qualified Data.ByteString.Lazy as BL +import qualified Data.Text as T +import qualified Data.Vector as V +import Test.HUnit + +import DataFrame.DecisionTree ( + TreeConfig (..), + defaultTreeConfig, + ) +import DataFrame.DecisionTree.Cart (buildCartTree) +import DataFrame.DecisionTree.Predict (predictManyWithTree) +import qualified DataFrame.Operations.Subset as DSub +import qualified DataFrameApi as D + +data Fold = Fold ![Int] ![Int] +instance FromJSON Fold where + parseJSON = withObject "fold" $ \o -> Fold <$> o .: "train" <*> o .: "test" + +newtype Folds = Folds [Fold] +instance FromJSON Folds where + parseJSON = withObject "folds" $ \o -> Folds <$> o .: "folds" + +data Fixture = Fixture ![Int] ![T.Text] +instance FromJSON Fixture where + parseJSON = withObject "fixture" $ \o -> Fixture <$> o .: "test_index" <*> o .: "test_pred" + +-- sklearn cart_d4 params: max_depth 4, min_samples_leaf 1 (min_samples_split +-- is fixed at 2 inside buildCartTree). +cartCfg :: TreeConfig +cartCfg = defaultTreeConfig{maxTreeDepth = 4, minLeafSize = 1} + +cartCases :: [(String, Int)] +cartCases = + [("wine", i) | i <- [0 .. 4]] ++ [("bcw", i) | i <- [0 .. 4]] ++ [("adult", 0)] + +tests :: [Test] +tests = + [ TestLabel ("cart: " ++ n ++ " fold " ++ show i) (TestCase (runCase n i)) + | (n, i) <- cartCases + ] + +readJson :: (FromJSON a) => FilePath -> IO (Either String a) +readJson fp = do + e <- try (BL.readFile fp) :: IO (Either SomeException BL.ByteString) + pure $ case e of + Left _ -> Left "missing" + Right raw -> eitherDecode raw + +-- wine is tie-free so sklearn is deterministic (exact match); bcw/adult have +-- equal-gain ties sklearn breaks via a seeded feature permutation, so we only +-- report the match fraction rather than chase its RNG. +runCase :: String -> Int -> IO () +runCase name i = do + efx <- readJson ("tests/fixtures/cart/" ++ name ++ "_fold" ++ show i ++ ".json") + case efx of + Left "missing" -> + putStrLn + ( " [skip] cart " + ++ name + ++ " fold " + ++ show i + ++ ": fixture missing (run bench/export_cart_fixtures.py)" + ) + Left e -> assertFailure ("fixture parse (" ++ name ++ "): " ++ e) + Right (Fixture _ predExpected) -> do + efolds <- readJson ("data/folds/" ++ name ++ ".json") + case efolds of + Left e -> assertFailure ("folds parse (" ++ name ++ "): " ++ e) + Right (Folds fs) -> do + df <- D.readCsv ("data/uci/" ++ name ++ "_clean.csv") + let Fold trainIdx testIdx = fs !! i + trainDf = DSub.selectRows trainIdx df + tree = buildCartTree @Int cartCfg "target" trainDf + preds = + map + (T.pack . show) + (V.toList (predictManyWithTree tree df (V.fromList testIdx))) + if name == "wine" + then assertEqual ("cart " ++ name ++ " fold " ++ show i) predExpected preds + else do + let n = length predExpected + m = length (filter id (zipWith (==) predExpected preds)) + putStrLn + ( " [diagnostic] cart " + ++ name + ++ " fold " + ++ show i + ++ ": " + ++ show m + ++ "/" + ++ show n + ++ " predictions match sklearn(random_state=0) (remainder = sklearn's seeded equal-gain tie-break)" + ) diff --git a/tests-internal/DataFrameApi.hs b/tests-internal/DataFrameApi.hs new file mode 100644 --- /dev/null +++ b/tests-internal/DataFrameApi.hs @@ -0,0 +1,22 @@ +{- | The umbrella subset these internal tests use. The suite cannot depend on +the meta @dataframe@ package (that would be a package-level cycle), so the +names the tests reach for are re-exported under one alias here. +-} +module DataFrameApi ( + DataFrame, + fromNamedColumns, + fromUnnamedColumns, + dimensions, + nRows, + rename, + exclude, + randomSplit, + derive, + readCsv, +) where + +import DataFrame.Core (DataFrame, fromNamedColumns) +import DataFrame.IO.CSV (readCsv) +import DataFrame.Operations.Core (dimensions, fromUnnamedColumns, nRows, rename) +import DataFrame.Operations.Subset (exclude, randomSplit) +import DataFrame.Operations.Transformations (derive) diff --git a/tests-internal/DecisionTree.hs b/tests-internal/DecisionTree.hs new file mode 100644 --- /dev/null +++ b/tests-internal/DecisionTree.hs @@ -0,0 +1,1360 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module DecisionTree where + +import DataFrame.DecisionTree +import DataFrame.DecisionTree.Cart (buildCartTree) +import DataFrame.DecisionTree.Categorical ( + TargetInfo, + discreteConditions, + mkTargetInfo, + ) +import DataFrame.DecisionTree.CondVec ( + CondVec (..), + combineAndVec, + combineOrVec, + materializeCondVec, + ) +import DataFrame.DecisionTree.Fit ( + ProbTree, + buildProbTree, + fitDecisionTree, + fitProbTree, + probExprs, + probsFromIndices, + ) +import DataFrame.DecisionTree.Numeric ( + NumExpr (NMaybeDouble), + generateNumericConds, + numericCols, + numericExprsWithTerms, + ) +import DataFrame.DecisionTree.Predict ( + computeTreeLoss, + countCarePointErrors, + identifyCarePoints, + majorityValueFromIndices, + partitionIndices, + predictWithTree, + ) +import DataFrame.DecisionTree.Tao (taoIteration, taoOptimize) +import DataFrame.DecisionTree.Types (CarePoint (..), Direction (..)) +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr (..), eqExpr, getColumns) +import DataFrame.Internal.Interpreter (interpret) +import qualified DataFrame.LinearSolver +import DataFrame.Operators +import qualified DataFrameApi as D + +import Data.Function (on) +import Data.List (maximumBy, sort) +import qualified Data.Map.Strict as M +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Test.HUnit + +------------------------------------------------------------------------ +-- Shared fixtures +------------------------------------------------------------------------ + +{- | Build a 'TargetInfo' or fail loudly; the test fixtures always satisfy +'mkTargetInfo', so a 'Nothing' here is a broken test, not a runtime case. +-} +requireTargetInfo :: T.Text -> D.DataFrame -> TargetInfo T.Text +requireTargetInfo target df = case mkTargetInfo @T.Text target df of + Just ti -> ti + Nothing -> error ("requireTargetInfo: no target info for " <> T.unpack target) + +-- 4 rows: label = ["A","B","A","C"], x = [1.0,2.0,3.0,4.0] +fixtureDF :: D.DataFrame +fixtureDF = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "B", "A", "C"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + +allIndices :: V.Vector Int +allIndices = V.fromList [0, 1, 2, 3] + +leftTree :: Tree T.Text +leftTree = Leaf "A" + +rightTree :: Tree T.Text +rightTree = Leaf "B" + +-- x <= 2.5: True for idx 0,1 (→ left); False for idx 2,3 (→ right) +splitCond :: Expr Bool +splitCond = F.col @Double "x" .<= F.lit (2.5 :: Double) + +-- Pre-computed care points for the full fixture +carePoints3 :: [CarePoint] +carePoints3 = + identifyCarePoints @T.Text "label" fixtureDF allIndices leftTree rightTree + +------------------------------------------------------------------------ +-- Unit tests: identifyCarePoints +------------------------------------------------------------------------ + +carePointsBothWrong :: Test +carePointsBothWrong = + TestCase $ + assertBool + "idx 3 (label=C, neither A nor B) should not be a care point" + (3 `notElem` map cpIndex carePoints3) + +carePointsLeftCorrect :: Test +carePointsLeftCorrect = TestCase $ do + let cp0 = filter ((== 0) . cpIndex) carePoints3 + case cp0 of + (c : _) -> + assertEqual + "idx 0 (label=A matches left Leaf A) should route GoLeft" + GoLeft + (cpCorrectDir c) + [] -> assertFailure "idx 0 should be a care point" + +carePointsRightCorrect :: Test +carePointsRightCorrect = TestCase $ do + let cp1 = filter ((== 1) . cpIndex) carePoints3 + case cp1 of + (c : _) -> + assertEqual + "idx 1 (label=B matches right Leaf B) should route GoRight" + GoRight + (cpCorrectDir c) + [] -> assertFailure "idx 1 should be a care point" + +carePointsMixed :: Test +carePointsMixed = TestCase $ do + assertEqual "exactly 3 care points" 3 (length carePoints3) + let idxs = map cpIndex carePoints3 + assertBool "idx 0 present" (0 `elem` idxs) + assertBool "idx 1 present" (1 `elem` idxs) + assertBool "idx 2 present" (2 `elem` idxs) + assertBool "idx 3 absent" (3 `notElem` idxs) + +carePointsBothCorrect :: Test +carePointsBothCorrect = TestCase $ do + let df2 = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0] :: [Double])) + ] + cps = + identifyCarePoints @T.Text + "label" + df2 + (V.fromList [0, 1]) + (Leaf "A") + (Leaf "A") + assertEqual "no care points when both subtrees agree" 0 (length cps) + +------------------------------------------------------------------------ +-- Unit tests: majorityValueFromIndices +------------------------------------------------------------------------ + +majorityVoteTest :: Test +majorityVoteTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["cat", "dog", "cat", "cat"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + assertEqual + "majority is cat (3 votes)" + "cat" + (majorityValueFromIndices @T.Text "label" df (V.fromList [0, 1, 2, 3])) + +majorityVoteSubset :: Test +majorityVoteSubset = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["cat", "dog", "cat", "cat"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + result = majorityValueFromIndices @T.Text "label" df (V.fromList [0, 1, 3]) + assertEqual "majority from subset [0,1,3] is cat" "cat" result + +------------------------------------------------------------------------ +-- Unit tests: computeTreeLoss +------------------------------------------------------------------------ + +computeLossZero :: Test +computeLossZero = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + loss = computeTreeLoss @T.Text "label" df (V.fromList [0, 1, 2, 3]) stump + assertEqual "perfect stump has zero loss" 0.0 loss + +computeLossHalf :: Test +computeLossHalf = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + constTree = Leaf "A" :: Tree T.Text + loss = computeTreeLoss @T.Text "label" df (V.fromList [0, 1, 2, 3]) constTree + assertEqual "constant leaf misclassifies half of balanced data" 0.5 loss + +------------------------------------------------------------------------ +-- Unit tests: partitionIndices +------------------------------------------------------------------------ + +partitionDisjoint :: Test +partitionDisjoint = TestCase $ do + let (lft, rgt) = partitionIndices splitCond fixtureDF allIndices + leftSet = V.toList lft + rightSet = V.toList rgt + intersection = filter (`elem` rightSet) leftSet + assertEqual "left and right partitions are disjoint" [] intersection + +partitionUnion :: Test +partitionUnion = TestCase $ do + let (lft, rgt) = partitionIndices splitCond fixtureDF allIndices + combined = sort (V.toList lft ++ V.toList rgt) + assertEqual + "union of partitions equals the original index set" + [0, 1, 2, 3] + combined + +------------------------------------------------------------------------ +-- Unit tests: countCarePointErrors +------------------------------------------------------------------------ + +countErrorsAllCorrect :: Test +countErrorsAllCorrect = TestCase $ do + let cps = [CarePoint 0 GoLeft, CarePoint 1 GoRight] + cond = F.col @Double "x" .<= F.lit (1.5 :: Double) + errs = countCarePointErrors cond fixtureDF cps + assertEqual "condition routes all care points correctly" 0 errs + +countErrorsAllWrong :: Test +countErrorsAllWrong = TestCase $ do + let cps = [CarePoint 0 GoLeft, CarePoint 1 GoRight] + cond = F.col @Double "x" .> F.lit (1.5 :: Double) + errs = countCarePointErrors cond fixtureDF cps + assertEqual "reversed condition misroutes all care points" 2 errs + +------------------------------------------------------------------------ +-- Unit tests: predictWithTree +------------------------------------------------------------------------ + +predictLeaf :: Test +predictLeaf = + TestCase $ + assertEqual + "leaf prediction ignores row index" + "Z" + (predictWithTree @T.Text "label" fixtureDF 0 (Leaf "Z")) + +predictBranch :: Test +predictBranch = TestCase $ do + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + assertEqual + "idx 0 (x=1.0 <= 2.5) routes left -> A" + "A" + (predictWithTree @T.Text "label" fixtureDF 0 stump) + assertEqual + "idx 3 (x=4.0 > 2.5) routes right -> B" + "B" + (predictWithTree @T.Text "label" fixtureDF 3 stump) + +------------------------------------------------------------------------ +-- Integration tests +------------------------------------------------------------------------ + +-- 20-row, linearly separable: x in [1..10] -> "pos", x in [11..20] -> "neg" +sepDF :: D.DataFrame +sepDF = + let xs = map fromIntegral [1 .. 20 :: Int] :: [Double] + labels = map (\x -> if x <= 10.0 then "pos" else "neg") xs :: [T.Text] + in D.fromNamedColumns + [ ("label", DI.fromList labels) + , ("x", DI.fromList xs) + ] + +-- Candidate conditions that bracket the decision boundary +sepConds :: [Expr Bool] +sepConds = + [ F.col @Double "x" .<= F.lit (10.5 :: Double) + , F.col @Double "x" .> F.lit (10.5 :: Double) + ] + +testCfg :: TreeConfig +testCfg = + defaultTreeConfig + { taoIterations = 5 + , expressionPairs = 4 + , minLeafSize = 1 + } + +-- Initial tree deliberately wrong: routes "pos" rows to the "neg" leaf +wrongStump :: Tree T.Text +wrongStump = + Branch + (F.col @Double "x" .> F.lit (10.5 :: Double)) + (Leaf "pos") + (Leaf "neg") + +taoNoDegradation :: Test +taoNoDegradation = TestCase $ do + let indices = V.enumFromN 0 20 + initialLoss = computeTreeLoss @T.Text "label" sepDF indices wrongStump + optimized = + taoOptimize @T.Text testCfg "label" sepConds sepDF indices wrongStump + finalLoss = computeTreeLoss @T.Text "label" sepDF indices optimized + assertBool + "taoOptimize must not increase loss" + (finalLoss <= initialLoss + 1e-9) + +taoMonotone :: Test +taoMonotone = TestCase $ do + let indices = V.enumFromN 0 20 + initLoss = computeTreeLoss @T.Text "label" sepDF indices wrongStump + stepTree = taoIteration @T.Text testCfg "label" sepConds sepDF indices + step (tree, _) = + let tree' = stepTree tree + in (tree', computeTreeLoss @T.Text "label" sepDF indices tree') + snapshots = take 6 $ iterate step (wrongStump, initLoss) + losses = map snd snapshots + pairs = zip losses (drop 1 losses) + assertBool + "loss must be non-increasing across taoIteration steps" + (all (\(a, b) -> b <= a + 1e-9) pairs) + +taoConvergesPureLabels :: Test +taoConvergesPureLabels = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (replicate 10 ("A" :: T.Text))) + , ("x", DI.fromList ([1.0 .. 10.0] :: [Double])) + ] + indices = V.enumFromN 0 10 + initTree = Leaf "A" :: Tree T.Text + initLoss = computeTreeLoss @T.Text "label" df indices initTree + result = + taoOptimize @T.Text testCfg "label" sepConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual "pure-label initial loss must be zero" 0.0 initLoss + assertEqual "pure-label final loss must still be zero" 0.0 finalLoss + +taoDeadBranchNoCrash :: Test +taoDeadBranchNoCrash = TestCase $ do + let badCond = F.col @Double "x" .<= F.lit (0.5 :: Double) + indices = V.enumFromN 0 20 + initTree = Branch badCond (Leaf "pos") (Leaf "neg") :: Tree T.Text + result = + taoOptimize @T.Text testCfg "label" [badCond] sepDF indices initTree + finalLoss = computeTreeLoss @T.Text "label" sepDF indices result + assertBool + "dead-branch tree must produce a valid loss in [0,1]" + (finalLoss >= 0.0 && finalLoss <= 1.0) + +------------------------------------------------------------------------ +-- Shared fixtures: 4x4 grid +------------------------------------------------------------------------ + +gridPairs :: [(Double, Double)] +gridPairs = [(x, y) | y <- [1 .. 4], x <- [1 .. 4]] + +gridBaseDF :: D.DataFrame +gridBaseDF = + D.fromNamedColumns + [ ("x", DI.fromList (map fst gridPairs)) + , ("y", DI.fromList (map snd gridPairs)) + ] + +------------------------------------------------------------------------ +-- Oblique recovery tests +------------------------------------------------------------------------ + +taoRecoversSingleObliqueDerived :: Test +taoRecoversSingleObliqueDerived = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("pos" :: T.Text)) + (F.lit ("neg" :: T.Text)) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + initTree = + Branch + (F.col @Double "x" .<= F.lit (2.5 :: Double)) + (Leaf "pos") + (Leaf "neg") :: + Tree T.Text + conds = + [ (F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double) + , (F.col @Double "x" + F.col @Double "y") .> F.lit (4.5 :: Double) + ] + cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + result = taoOptimize @T.Text cfg "label" conds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "TAO recovers single oblique (x+y) split with zero loss" + 0.0 + finalLoss + +taoRecoversNestedObliqueDerived :: Test +taoRecoversNestedObliqueDerived = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("low" :: T.Text)) + ( F.ifThenElse + ((F.col @Double "x" - F.col @Double "y") .<= F.lit (0.5 :: Double)) + (F.lit "mid") + (F.lit "high") + ) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + initTree = + Branch + (F.col @Double "x" .<= F.lit (1.5 :: Double)) + (Leaf "low") + ( Branch + (F.col @Double "y" .<= F.lit (3.5 :: Double)) + (Leaf "mid") + (Leaf "high") + ) :: + Tree T.Text + conds = + [ (F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double) + , (F.col @Double "x" + F.col @Double "y") .> F.lit (4.5 :: Double) + , (F.col @Double "x" - F.col @Double "y") .<= F.lit (0.5 :: Double) + , (F.col @Double "x" - F.col @Double "y") .> F.lit (0.5 :: Double) + ] + cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + result = taoOptimize @T.Text cfg "label" conds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "TAO recovers nested oblique (x+y)/(x-y) tree with zero loss" + 0.0 + finalLoss + +-- Shared setup for C2 (a) and (b): axis-aligned pool only, oblique label. +obliqueAxisAlignedFixture :: + (D.DataFrame, V.Vector Int, [Expr Bool], Tree T.Text) +obliqueAxisAlignedFixture = + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("pos" :: T.Text)) + (F.lit ("neg" :: T.Text)) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + axisConds = + [F.col @Double "x" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + ++ [F.col @Double "y" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + initTree = + Branch + (F.col @Double "x" .<= F.lit (2.5 :: Double)) + (Leaf "pos") + (Leaf "neg") :: + Tree T.Text + in (df, indices, axisConds, initTree) + +-- C2 (a): with the linear solver OFF, axis-aligned pool cannot recover the +-- oblique decision boundary. Preserves the original guarantee of the test. +taoAxisAlignedInsufficientForObliqueDiscreteOnly :: Test +taoAxisAlignedInsufficientForObliqueDiscreteOnly = TestCase $ do + let (df, indices, axisConds, initTree) = obliqueAxisAlignedFixture + cfg = + defaultTreeConfig + { taoIterations = 10 + , expressionPairs = 6 + , minLeafSize = 1 + , useLinearSolver = False + } + result = taoOptimize @T.Text cfg "label" axisConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertBool + "axis-aligned stump cannot recover oblique label without linear solver (loss > 0.1)" + (finalLoss > 0.1) + +-- C2 (b): with the linear solver ON, the L1-LR fit discovers the oblique +-- (x + y) hyperplane even though only axis-aligned conditions are in the pool. +taoLinearRecoversObliqueFromAxisAlignedPool :: Test +taoLinearRecoversObliqueFromAxisAlignedPool = TestCase $ do + let (df, indices, axisConds, initTree) = obliqueAxisAlignedFixture + cfg = + defaultTreeConfig + { taoIterations = 10 + , expressionPairs = 6 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 2 + } + result = taoOptimize @T.Text cfg "label" axisConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "linear solver recovers oblique split from axis-aligned-only pool" + 0.0 + finalLoss + +------------------------------------------------------------------------ +-- Nullable numeric feature tests +------------------------------------------------------------------------ + +-- Cleanly separable nullable column (no actual nulls): Just 1..6 -> "pos", +-- Just 7..12 -> "neg". Exercises the nullable numeric path. +nullableSepDF :: D.DataFrame +nullableSepDF = + D.fromNamedColumns + [ ("label", DI.fromList (replicate 6 "pos" ++ replicate 6 "neg" :: [T.Text])) + , + ( "x" + , DI.fromVector + ( V.fromList $ + map (Just . fromIntegral) ([1 .. 6] :: [Int]) + ++ map (Just . fromIntegral) ([7 .. 12] :: [Int]) :: + V.Vector (Maybe Double) + ) + ) + ] + +-- DF with genuine nulls interspersed. +nullsMixedDF :: D.DataFrame +nullsMixedDF = + D.fromNamedColumns + [ ("label", DI.fromList (["pos", "pos", "pos", "neg", "neg", "neg"] :: [T.Text])) + , + ( "x" + , DI.fromVector + ( V.fromList + [Just 1.0, Nothing, Just 3.0, Just 7.0, Nothing, Just 9.0] :: + V.Vector (Maybe Double) + ) + ) + ] + +-- numericCols picks up DI.fromVector (Maybe Double) as NMaybeDouble. +numericColsNullableDoubleTest :: Test +numericColsNullableDoubleTest = TestCase $ do + let exprs = numericCols nullableSepDF + hasMD = any (\case NMaybeDouble _ -> True; _ -> False) exprs + assertBool + "numericCols finds NMaybeDouble for DI.fromVector (Maybe Double)" + hasMD + +-- numericCols picks up DI.fromVector (Maybe Int) as NMaybeDouble (via whenPresent). +numericColsNullableIntTest :: Test +numericColsNullableIntTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["pos", "neg"] :: [T.Text])) + , + ( "n" + , DI.fromVector (V.fromList [Just (1 :: Int), Just 2] :: V.Vector (Maybe Int)) + ) + ] + hasMD = any (\case NMaybeDouble _ -> True; _ -> False) (numericCols df) + assertBool "numericCols finds NMaybeDouble for DI.fromVector (Maybe Int)" hasMD + +-- generateNumericConds is non-empty for a DF with an DI.fromVector (Maybe Double). +numericCondsNullableNonEmptyTest :: Test +numericCondsNullableNonEmptyTest = + TestCase $ + assertBool + "generateNumericConds non-empty for DI.fromVector (Maybe Double)" + (not (null (generateNumericConds defaultTreeConfig nullableSepDF))) + +-- Null values evaluate to False for threshold conditions (null rows route right). +nullValueRoutesFalseTest :: Test +nullValueRoutesFalseTest = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "B"] :: [T.Text])) + , + ( "x" + , DI.fromVector + (V.fromList [Nothing, Just (5.0 :: Double)] :: V.Vector (Maybe Double)) + ) + ] + cond = F.fromMaybe False (F.col @(Maybe Double) "x" .<= F.lit (6.0 :: Double)) + (lft, rgt) = partitionIndices cond df (V.fromList [0, 1]) + assertBool "null row (idx 0) routes to right (false) partition" (0 `V.elem` rgt) + assertBool "Just 5.0 <= 6.0 routes to left (true) partition" (1 `V.elem` lft) + +-- Nullable feature (no actual nulls) achieves zero loss on cleanly separable data. +nullableFitZeroLossTest :: Test +nullableFitZeroLossTest = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + featureDf = D.exclude ["label"] nullableSepDF + conds = generateNumericConds cfg featureDf + initTree = buildCartTree @T.Text cfg "label" nullableSepDF + indices = V.enumFromN 0 12 + result = taoOptimize @T.Text cfg "label" conds nullableSepDF indices initTree + loss = computeTreeLoss @T.Text "label" nullableSepDF indices result + assertEqual "zero loss on cleanly separable OptionalColumn data" 0.0 loss + +-- fitDecisionTree with genuine nulls: loss is a valid probability and no crash. +nullableFitWithNullsNoCrashTest :: Test +nullableFitWithNullsNoCrashTest = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 3, expressionPairs = 4, minLeafSize = 1} + featureDf = D.exclude ["label"] nullsMixedDF + conds = generateNumericConds cfg featureDf + initTree = buildCartTree @T.Text cfg "label" nullsMixedDF + indices = V.enumFromN 0 6 + result = taoOptimize @T.Text cfg "label" conds nullsMixedDF indices initTree + loss = computeTreeLoss @T.Text "label" nullsMixedDF indices result + assertBool + "loss is in [0,1] with null values present" + (loss >= 0.0 && loss <= 1.0) + +-- numericExprsWithTerms produces cross-column combinations when one col is +-- DI.fromVector (Maybe Double) and another is a plain UnboxedColumn Double. +numericExprsWithTermsMixedTest :: Test +numericExprsWithTermsMixedTest = TestCase $ do + let df = + D.fromNamedColumns + [ + ( "x" + , DI.fromVector + (V.fromList [Just 1.0, Just 2.0, Just 3.0] :: V.Vector (Maybe Double)) + ) + , ("y", DI.fromList ([4.0, 5.0, 6.0] :: [Double])) + ] + cfg = defaultSynthConfig{maxExprDepth = 2, enableArithOps = True} + exprs = numericExprsWithTerms cfg df + assertBool + "more than 2 expressions: base cols + combinations" + (length exprs > 2) + assertBool + "combined exprs include NMaybeDouble (nullable arithmetic)" + (any (\case NMaybeDouble _ -> True; _ -> False) exprs) + +------------------------------------------------------------------------ +-- Probability tree tests +------------------------------------------------------------------------ + +-- probsFromIndices: counts correct on a 3-row slice +probsFromIndicesBasic :: Test +probsFromIndicesBasic = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0] :: [Double])) + ] + probs = probsFromIndices @T.Text "label" df (V.fromList [0, 1, 2]) + assertBool "A prob ≈ 2/3" (abs (probs M.! "A" - 2 / 3) < 1e-9) + assertBool "B prob ≈ 1/3" (abs (probs M.! "B" - 1 / 3) < 1e-9) + +-- probsFromIndices: only a subset of rows counted +probsFromIndicesSubset :: Test +probsFromIndicesSubset = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "B", "B"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0, 4.0] :: [Double])) + ] + probs = probsFromIndices @T.Text "label" df (V.fromList [0, 1]) + assertEqual "only rows 0,1 → A:1.0" (M.fromList [("A", 1.0)]) probs + +-- probsFromIndices: single class → probability 1.0 +probsFromIndicesSingleClass :: Test +probsFromIndicesSingleClass = TestCase $ do + let probs = probsFromIndices @T.Text "label" fixtureDF (V.fromList [0, 2]) + assertEqual "rows 0,2 both A → A:1.0" (M.fromList [("A", 1.0)]) probs + +-- buildProbTree: Leaf preserves distribution +buildProbTreeLeaf :: Test +buildProbTreeLeaf = TestCase $ do + let df = + D.fromNamedColumns + [ ("label", DI.fromList (["A", "A", "A"] :: [T.Text])) + , ("x", DI.fromList ([1.0, 2.0, 3.0] :: [Double])) + ] + pt = buildProbTree @T.Text (Leaf "A") "label" df (V.fromList [0, 1, 2]) + case pt of + Leaf m -> assertEqual "pure-A leaf → {A:1.0}" (M.fromList [("A", 1.0)]) m + _ -> assertFailure "expected Leaf" + +-- buildProbTree: Branch distributes rows to left/right leaves correctly +buildProbTreeBranch :: Test +buildProbTreeBranch = TestCase $ do + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + pt = buildProbTree @T.Text stump "label" fixtureDF allIndices + case pt of + Branch _ (Leaf lm) (Leaf rm) -> do + assertBool "left leaf has A:0.5" (abs (M.findWithDefault 0 "A" lm - 0.5) < 1e-9) + assertBool "left leaf has B:0.5" (abs (M.findWithDefault 0 "B" lm - 0.5) < 1e-9) + assertBool + "right leaf has A:0.5" + (abs (M.findWithDefault 0 "A" rm - 0.5) < 1e-9) + assertBool + "right leaf has C:0.5" + (abs (M.findWithDefault 0 "C" rm - 0.5) < 1e-9) + _ -> assertFailure "expected Branch with two Leaves" + +-- probExprs: leaf tree produces Lit values +probExprsLeaf :: Test +probExprsLeaf = TestCase $ do + let pt = Leaf (M.fromList [("A", 0.75), ("B", 0.25)]) :: ProbTree T.Text + pe = probExprs pt + assertBool "A expr is Lit 0.75" (eqExpr (Lit 0.75) (pe M.! "A")) + assertBool "B expr is Lit 0.25" (eqExpr (Lit 0.25) (pe M.! "B")) + +-- probExprs: class absent from one leaf gets Lit 0.0 on that side +probExprsMissingClass :: Test +probExprsMissingClass = TestCase $ do + let pt = + Branch + splitCond + (Leaf (M.fromList [("A", 1.0)])) + (Leaf (M.fromList [("B", 1.0)])) :: + ProbTree T.Text + pe = probExprs pt + assertBool + "A expr: If cond (Lit 1.0) (Lit 0.0)" + (eqExpr (F.ifThenElse splitCond (Lit 1.0) (Lit 0.0)) (pe M.! "A")) + assertBool + "B expr: If cond (Lit 0.0) (Lit 1.0)" + (eqExpr (F.ifThenElse splitCond (Lit 0.0) (Lit 1.0)) (pe M.! "B")) + +-- probExprs: keys equal all classes that appear across any leaf +probExprsAllClasses :: Test +probExprsAllClasses = TestCase $ do + let pt = + Branch + splitCond + (Leaf (M.fromList [("A", 1.0)])) + (Leaf (M.fromList [("B", 0.6), ("C", 0.4)])) :: + ProbTree T.Text + pe = probExprs pt + assertEqual "three classes in result" (sort ["A", "B", "C"]) (sort (M.keys pe)) + +-- Probabilities sum to 1.0 at every row after applying probExprs +probsSumToOne :: Test +probsSumToOne = TestCase $ do + let stump = Branch splitCond (Leaf "A") (Leaf "B") :: Tree T.Text + pt = buildProbTree @T.Text stump "label" fixtureDF allIndices + pe = probExprs pt + sumExpr = foldl1 (.+) (M.elems pe) + case interpret @Double fixtureDF sumExpr of + Left e -> assertFailure (show e) + Right (DI.TColumn sumCol) -> + case DI.toVector @Double sumCol of + Left e2 -> assertFailure (show e2) + Right vals -> + mapM_ + (\v -> assertBool ("sum ≈ 1.0, got " ++ show v) (abs (v - 1.0) < 1e-9)) + (V.toList vals) + +-- argmax of probExprs agrees with fitDecisionTree on sepDF +probArgmaxMatchesClassifier :: Test +probArgmaxMatchesClassifier = TestCase $ do + let cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + hardExpr = fitDecisionTree @T.Text cfg (Col "label") sepDF + pe = fitProbTree @T.Text cfg (Col "label") sepDF + indices = [0 .. D.nRows sepDF - 1] + case interpret @T.Text sepDF hardExpr of + Left e -> assertFailure (show e) + Right (DI.TColumn hardCol) -> + case DI.toVector @T.Text hardCol of + Left e2 -> assertFailure (show e2) + Right hardVals -> do + probCols <- + mapM + ( \(cls, expr) -> case interpret @Double sepDF expr of + Left e3 -> assertFailure (show e3) >> return (cls, V.empty) + Right (DI.TColumn col2) -> case DI.toVector @Double col2 of + Left e4 -> assertFailure (show e4) >> return (cls, V.empty) + Right v -> return (cls, v) + ) + (M.toList pe) + mapM_ + ( \i -> + let argmax = fst $ maximumBy (compare `on` (V.! i) . snd) probCols + hard = hardVals V.! i + in assertEqual ("row " ++ show i) hard argmax + ) + indices + +------------------------------------------------------------------------ +-- C4-C9 / D-series: linear solver integration tests +------------------------------------------------------------------------ + +-- C4: Nested oblique recovery without oblique hints; label set by two oblique +-- boundaries but only axis-aligned thresholds in the pool. The linear solver +-- should learn both splits and reach zero loss. +taoRecoversNestedObliqueWithoutHint :: Test +taoRecoversNestedObliqueWithoutHint = TestCase $ do + let labelExpr = + F.ifThenElse + ((F.col @Double "x" + F.col @Double "y") .<= F.lit (4.5 :: Double)) + (F.lit ("low" :: T.Text)) + ( F.ifThenElse + ((F.col @Double "x" - F.col @Double "y") .<= F.lit (0.5 :: Double)) + (F.lit "mid") + (F.lit "high") + ) + df = D.derive @T.Text "label" labelExpr gridBaseDF + indices = V.enumFromN 0 16 + initTree = + Branch + (F.col @Double "x" .<= F.lit (1.5 :: Double)) + (Leaf "low") + ( Branch + (F.col @Double "y" .<= F.lit (3.5 :: Double)) + (Leaf "mid") + (Leaf "high") + ) :: + Tree T.Text + axisOnlyConds = + [F.col @Double "x" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + ++ [F.col @Double "y" .<= F.lit (t :: Double) | t <- [1.5, 2.5, 3.5]] + cfg = + defaultTreeConfig + { taoIterations = 20 + , expressionPairs = 6 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 2 + } + result = taoOptimize @T.Text cfg "label" axisOnlyConds df indices initTree + finalLoss = computeTreeLoss @T.Text "label" df indices result + assertEqual + "linear solver recovers nested oblique tree from axis-aligned-only pool" + 0.0 + finalLoss + +-- C5: Monotone loss across iterations with the linear solver enabled. +-- Resolves Issue 1 from the prior plan (currentCond included in the +-- competition pool). +taoMonotoneWithLinear :: Test +taoMonotoneWithLinear = TestCase $ do + let indices = V.enumFromN 0 20 + cfg = defaultTreeConfig{taoIterations = 5, expressionPairs = 4, minLeafSize = 1} + initLoss = computeTreeLoss @T.Text "label" sepDF indices wrongStump + stepTree = taoIteration @T.Text cfg "label" sepConds sepDF indices + step (tree, _) = + let tree' = stepTree tree + in (tree', computeTreeLoss @T.Text "label" sepDF indices tree') + snapshots = take 6 $ iterate step (wrongStump, initLoss) + losses = map snd snapshots + pairs = zip losses (drop 1 losses) + assertBool + ("loss must be non-increasing across iterations (got " ++ show losses ++ ")") + (all (\(a, b) -> b <= a + 1e-9) pairs) + +-- C6: When the discrete pool contains an exact-zero-error split (axis-aligned +-- works perfectly), the competition picks the simpler discrete candidate +-- rather than a similarly-good but more complex linear one. +taoLinearVsDiscreteCompetition :: Test +taoLinearVsDiscreteCompetition = TestCase $ do + let indices = V.enumFromN 0 20 + cfg = + defaultTreeConfig + { taoIterations = 5 + , expressionPairs = 4 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 2 + } + result = taoOptimize @T.Text cfg "label" sepConds sepDF indices wrongStump + finalLoss = computeTreeLoss @T.Text "label" sepDF indices result + assertEqual + "axis-aligned separable data should fit to zero loss" + 0.0 + finalLoss + +-- C8: Linear solver respects the L1 penalty and produces sparse hyperplanes +-- on data where only some features are informative. +taoLinearProducesSparsity :: Test +taoLinearProducesSparsity = TestCase $ do + let n = 50 :: Int + xs = [fromIntegral i / 10 - 2.5 :: Double | i <- [0 .. n - 1]] + avals = xs + bs = map (* 0.7) xs + cs = [fromIntegral ((i * 7) `mod` 11) / 5 - 1 :: Double | i <- [0 .. n - 1]] + ds = [fromIntegral ((i * 13) `mod` 7) / 3 - 1 :: Double | i <- [0 .. n - 1]] + labels = + [ if (avals !! i) + (bs !! i) > 0 then "pos" else "neg" :: T.Text + | i <- [0 .. n - 1] + ] + df = + D.fromNamedColumns + [ ("label", DI.fromList labels) + , ("a", DI.fromList avals) + , ("b", DI.fromList bs) + , ("c", DI.fromList cs) + , ("d", DI.fromList ds) + ] + cfg = + defaultTreeConfig + { maxTreeDepth = 1 + , taoIterations = 10 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 2 + , linearSolverConfig = + (linearSolverConfig defaultTreeConfig) + { DataFrame.LinearSolver.scL1Lambda = 0.05 + } + } + result = fitDecisionTree @T.Text cfg (Col "label") df + rootCols = getColumns result + assertBool + ( "informative columns 'a' or 'b' must appear in the fitted Expr (got " + ++ show rootCols + ++ ")" + ) + ("a" `elem` rootCols || "b" `elem` rootCols) + +-- C9: Determinism — same training data produces an equal (eqExpr) tree. +taoLinearDeterministic :: Test +taoLinearDeterministic = TestCase $ do + let cfg = + defaultTreeConfig + { taoIterations = 5 + , expressionPairs = 4 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 2 + } + r1 = fitDecisionTree @T.Text cfg (Col "label") sepDF + r2 = fitDecisionTree @T.Text cfg (Col "label") sepDF + assertBool "fitDecisionTree is deterministic on the same input" (eqExpr r1 r2) + +-- D1: One care point — solver must not crash; integration should fall back +-- gracefully (via minCarePointsForLinear) and rely on the discrete path. +taoLinearTinyCareSet :: Test +taoLinearTinyCareSet = TestCase $ do + let cfg = + defaultTreeConfig + { taoIterations = 5 + , expressionPairs = 4 + , minLeafSize = 1 + , useLinearSolver = True + , minCarePointsForLinear = 100 + } + result = fitDecisionTree @T.Text cfg (Col "label") sepDF + cfgOff = cfg{useLinearSolver = False} + resultOff = fitDecisionTree @T.Text cfgOff (Col "label") sepDF + assertBool + "skipping linear solver yields same expression as linear-off baseline" + (eqExpr result resultOff) + +------------------------------------------------------------------------ +-- Categorical-condition generator tests (Phase 1-2 of the plan) +------------------------------------------------------------------------ + +-- A binary-target DataFrame with a 5-level Text column whose levels have +-- monotonically-increasing positive rates. Breiman's algorithm should +-- enumerate the 4 contiguous-prefix splits in that exact rate order. +breimanBinaryDF :: D.DataFrame +breimanBinaryDF = + let n = 100 :: Int + mkLabel "a" = "neg" + mkLabel "b" = "neg" + mkLabel "c" = "pos" + mkLabel "d" = "pos" + mkLabel "e" = "pos" + mkLabel _ = "neg" + levels = cycle ["a", "b", "c", "d", "e"] + feats = take n levels + labs = map mkLabel feats + in D.fromUnnamedColumns + [ DI.fromList (map T.pack feats :: [T.Text]) + , DI.fromList (map T.pack labs :: [T.Text]) + ] + |> D.rename "0" "feat" + |> D.rename "1" "label" + +testCategoricalBreimanBinary :: Test +testCategoricalBreimanBinary = TestCase $ do + let ti = requireTargetInfo "label" breimanBinaryDF + conds = + discreteConditions @T.Text + ti + defaultTreeConfig + (D.exclude ["label"] breimanBinaryDF) + feat = "feat" + feats = filter (\c -> feat `elem` getColumns c) conds + assertEqual "Breiman emits k-1 prefixes" 4 (length feats) + +testCategoricalSubsetsMulticlassLowCard :: Test +testCategoricalSubsetsMulticlassLowCard = TestCase $ do + let n = 30 :: Int + feats = take n (cycle ["x", "y", "z"]) + labs = take n (cycle ["A", "B", "C"]) + df = + D.fromUnnamedColumns + [ DI.fromList (map T.pack feats :: [T.Text]) + , DI.fromList (map T.pack labs :: [T.Text]) + ] + |> D.rename "0" "feat" + |> D.rename "1" "label" + ti = requireTargetInfo "label" df + conds = discreteConditions @T.Text ti defaultTreeConfig (D.exclude ["label"] df) + feat = "feat" + feats' = filter (\c -> feat `elem` getColumns c) conds + assertEqual "subsets at low cardinality" 6 (length feats') + +testCategoricalSingletonsMulticlassHighCard :: Test +testCategoricalSingletonsMulticlassHighCard = TestCase $ do + let n = 60 :: Int + feats = take n (cycle ["a", "b", "c", "d", "e", "f"]) + labs = take n (cycle ["A", "B", "C"]) + df = + D.fromUnnamedColumns + [ DI.fromList (map T.pack feats :: [T.Text]) + , DI.fromList (map T.pack labs :: [T.Text]) + ] + |> D.rename "0" "feat" + |> D.rename "1" "label" + ti = requireTargetInfo "label" df + conds = discreteConditions @T.Text ti defaultTreeConfig (D.exclude ["label"] df) + feat = "feat" + feats' = filter (\c -> feat `elem` getColumns c) conds + assertEqual "singletons at high cardinality" 6 (length feats') + +testCategoricalCardZero :: Test +testCategoricalCardZero = TestCase $ do + let df = + D.fromUnnamedColumns + [ DI.fromList ([] :: [T.Text]) + , DI.fromList ([] :: [T.Text]) + ] + |> D.rename "0" "feat" + |> D.rename "1" "label" + ti = requireTargetInfo "label" df + conds = discreteConditions @T.Text ti defaultTreeConfig (D.exclude ["label"] df) + feat = "feat" + feats' = filter (\c -> feat `elem` getColumns c) conds + assertEqual "no candidates on empty column" 0 (length feats') + +testCategoricalNullableBinary :: Test +testCategoricalNullableBinary = TestCase $ do + let feats = + [ Just "a" + , Just "b" + , Just "c" + , Nothing + , Just "a" + , Just "b" + , Just "c" + , Nothing + , Just "a" + , Just "b" + , Just "c" + , Just "a" + , Just "b" + , Just "c" + , Just "a" + , Just "b" + ] + labs = + [ "neg" + , "neg" + , "pos" + , "neg" + , "neg" + , "neg" + , "pos" + , "neg" + , "neg" + , "neg" + , "pos" + , "neg" + , "neg" + , "pos" + , "neg" + , "pos" + ] + df = + D.fromUnnamedColumns + [ DI.fromList (feats :: [Maybe T.Text]) + , DI.fromList (map T.pack labs :: [T.Text]) + ] + |> D.rename "0" "feat" + |> D.rename "1" "label" + ti = requireTargetInfo "label" df + conds = discreteConditions @T.Text ti defaultTreeConfig (D.exclude ["label"] df) + feat = "feat" :: T.Text + feats' = filter (\c -> feat `elem` getColumns c) conds + assertEqual "Breiman prefixes on nullable column ignore nulls" 2 (length feats') + +------------------------------------------------------------------------ +-- PR 2 extended: threshold-consolidation rewrite in combineAndVec / +-- combineOrVec. Positive cases, negative cases, semantic-preservation check. +------------------------------------------------------------------------ + +-- A small synthetic DataFrame to materialize CondVecs against. +threshFixtureDF :: D.DataFrame +threshFixtureDF = + D.fromNamedColumns + [ ("x", DI.fromList ([0.0, 1.0, 2.0, 3.0, 4.0, 5.0] :: [Double])) + , ("y", DI.fromList ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0] :: [Double])) + ] + +materializeOrFail :: Expr Bool -> CondVec +materializeOrFail e = case materializeCondVec threshFixtureDF e of + Just cv -> cv + Nothing -> error "materializeOrFail: condition could not be materialized" + +-- | Helper: assert that two `Expr Bool`s agree by 'eqExpr'. +assertEqExpr :: String -> Expr Bool -> Expr Bool -> Assertion +assertEqExpr msg expected actual = + assertBool + (msg ++ "\n expected: " ++ show expected ++ "\n actual: " ++ show actual) + (eqExpr expected actual) + +-- Eight positive cases. + +threshAndLeq :: Test +threshAndLeq = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .<=. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .<=. F.lit (1.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "AND of x≤3 and x≤1 collapses to x≤1" + (F.col @Double "x" .<=. F.lit (1.0 :: Double)) + (cvExpr r) + +threshOrLeq :: Test +threshOrLeq = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .<=. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .<=. F.lit (1.0 :: Double)) + r = combineOrVec a b + assertEqExpr + "OR of x≤3 and x≤1 collapses to x≤3" + (F.col @Double "x" .<=. F.lit (3.0 :: Double)) + (cvExpr r) + +threshAndLt :: Test +threshAndLt = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .<. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .<. F.lit (1.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "AND of x<3 and x<1 collapses to x<1" + (F.col @Double "x" .<. F.lit (1.0 :: Double)) + (cvExpr r) + +threshOrLt :: Test +threshOrLt = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .<. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .<. F.lit (1.0 :: Double)) + r = combineOrVec a b + assertEqExpr + "OR of x<3 and x<1 collapses to x<3" + (F.col @Double "x" .<. F.lit (3.0 :: Double)) + (cvExpr r) + +threshAndGeq :: Test +threshAndGeq = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>=. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>=. F.lit (3.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "AND of x≥1 and x≥3 collapses to x≥3" + (F.col @Double "x" .>=. F.lit (3.0 :: Double)) + (cvExpr r) + +threshOrGeq :: Test +threshOrGeq = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>=. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>=. F.lit (3.0 :: Double)) + r = combineOrVec a b + assertEqExpr + "OR of x≥1 and x≥3 collapses to x≥1" + (F.col @Double "x" .>=. F.lit (1.0 :: Double)) + (cvExpr r) + +threshAndGt :: Test +threshAndGt = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>. F.lit (3.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "AND of x>1 and x>3 collapses to x>3" + (F.col @Double "x" .>. F.lit (3.0 :: Double)) + (cvExpr r) + +threshOrGt :: Test +threshOrGt = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>. F.lit (3.0 :: Double)) + r = combineOrVec a b + assertEqExpr + "OR of x>1 and x>3 collapses to x>1" + (F.col @Double "x" .>. F.lit (1.0 :: Double)) + (cvExpr r) + +-- Six negative cases: rewrite must NOT fire. + +threshNegMixedDirection :: Test +threshNegMixedDirection = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .<. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>=. F.lit (1.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "mixed-direction AND keeps generic F.and form" + (F.and (cvExpr a) (cvExpr b)) + (cvExpr r) + +threshNegCrossColumn :: Test +threshNegCrossColumn = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "y" .>. F.lit (3.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "cross-column AND keeps generic F.and form" + (F.and (cvExpr a) (cvExpr b)) + (cvExpr r) + +threshNegMixedOpFamily :: Test +threshNegMixedOpFamily = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .<. F.lit (4.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "different-op-family AND keeps generic F.and form" + (F.and (cvExpr a) (cvExpr b)) + (cvExpr r) + +threshNegEqualityOp :: Test +threshNegEqualityOp = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .==. F.lit (3.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .==. F.lit (1.0 :: Double)) + r = combineOrVec a b + assertEqExpr + "equality OR keeps generic F.or form" + (F.or (cvExpr a) (cvExpr b)) + (cvExpr r) + +threshNegLitOnLeft :: Test +threshNegLitOnLeft = TestCase $ do + let a = materializeOrFail (F.lit (1.0 :: Double) .<. F.col @Double "x") + b = materializeOrFail (F.lit (3.0 :: Double) .<. F.col @Double "x") + r = combineAndVec a b + assertEqExpr + "Lit-on-left AND keeps generic F.and form" + (F.and (cvExpr a) (cvExpr b)) + (cvExpr r) + +threshNegNonLiteralRhs :: Test +threshNegNonLiteralRhs = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.col @Double "y") + b = materializeOrFail (F.col @Double "x" .>. F.lit (3.0 :: Double)) + r = combineAndVec a b + assertEqExpr + "non-literal RHS AND keeps generic F.and form" + (F.and (cvExpr a) (cvExpr b)) + (cvExpr r) + +-- Semantic-preservation spot check: the consolidated cvVec matches the +-- elementwise AND/OR of the inputs at every row of a synthetic DataFrame. +threshSemanticPreservation :: Test +threshSemanticPreservation = TestCase $ do + let a = materializeOrFail (F.col @Double "x" .>. F.lit (1.0 :: Double)) + b = materializeOrFail (F.col @Double "x" .>. F.lit (3.0 :: Double)) + rAnd = combineAndVec a b + rOr = combineOrVec a b + expectedAnd = VU.zipWith (&&) (cvVec a) (cvVec b) + expectedOr = VU.zipWith (||) (cvVec a) (cvVec b) + assertEqual + "consolidated AND vec matches elementwise &&" + expectedAnd + (cvVec rAnd) + assertEqual + "consolidated OR vec matches elementwise ||" + expectedOr + (cvVec rOr) + +------------------------------------------------------------------------ +-- Test list +------------------------------------------------------------------------ + +tests :: [Test] +tests = + [ TestLabel "carePointsBothWrong" carePointsBothWrong + , TestLabel "carePointsLeftCorrect" carePointsLeftCorrect + , TestLabel "carePointsRightCorrect" carePointsRightCorrect + , TestLabel "carePointsMixed" carePointsMixed + , TestLabel "carePointsBothCorrect" carePointsBothCorrect + , TestLabel "majorityVoteTest" majorityVoteTest + , TestLabel "majorityVoteSubset" majorityVoteSubset + , TestLabel "computeLossZero" computeLossZero + , TestLabel "computeLossHalf" computeLossHalf + , TestLabel "partitionDisjoint" partitionDisjoint + , TestLabel "partitionUnion" partitionUnion + , TestLabel "countErrorsAllCorrect" countErrorsAllCorrect + , TestLabel "countErrorsAllWrong" countErrorsAllWrong + , TestLabel "predictLeaf" predictLeaf + , TestLabel "predictBranch" predictBranch + , TestLabel "taoNoDegradation" taoNoDegradation + , TestLabel "taoMonotone" taoMonotone + , TestLabel "taoConvergesPureLabels" taoConvergesPureLabels + , TestLabel "taoDeadBranchNoCrash" taoDeadBranchNoCrash + , TestLabel "taoRecoversSingleObliqueDerived" taoRecoversSingleObliqueDerived + , TestLabel "taoRecoversNestedObliqueDerived" taoRecoversNestedObliqueDerived + , TestLabel + "C2a taoAxisAlignedInsufficientForObliqueDiscreteOnly" + taoAxisAlignedInsufficientForObliqueDiscreteOnly + , TestLabel + "C2b taoLinearRecoversObliqueFromAxisAlignedPool" + taoLinearRecoversObliqueFromAxisAlignedPool + , TestLabel "numericColsNullableDouble" numericColsNullableDoubleTest + , TestLabel "numericColsNullableInt" numericColsNullableIntTest + , TestLabel "numericCondsNullableNonEmpty" numericCondsNullableNonEmptyTest + , TestLabel "nullValueRoutesFalse" nullValueRoutesFalseTest + , TestLabel "nullableFitZeroLoss" nullableFitZeroLossTest + , TestLabel "nullableFitWithNullsNoCrash" nullableFitWithNullsNoCrashTest + , TestLabel "numericExprsWithTermsMixed" numericExprsWithTermsMixedTest + , TestLabel "probsFromIndicesBasic" probsFromIndicesBasic + , TestLabel "probsFromIndicesSubset" probsFromIndicesSubset + , TestLabel "probsFromIndicesSingleClass" probsFromIndicesSingleClass + , TestLabel "buildProbTreeLeaf" buildProbTreeLeaf + , TestLabel "buildProbTreeBranch" buildProbTreeBranch + , TestLabel "probExprsLeaf" probExprsLeaf + , TestLabel "probExprsMissingClass" probExprsMissingClass + , TestLabel "probExprsAllClasses" probExprsAllClasses + , TestLabel "probsSumToOne" probsSumToOne + , TestLabel "probArgmaxMatchesClassifier" probArgmaxMatchesClassifier + , TestLabel + "C4 taoRecoversNestedObliqueWithoutHint" + taoRecoversNestedObliqueWithoutHint + , TestLabel "C5 taoMonotoneWithLinear" taoMonotoneWithLinear + , TestLabel "C6 taoLinearVsDiscreteCompetition" taoLinearVsDiscreteCompetition + , TestLabel "C8 taoLinearProducesSparsity" taoLinearProducesSparsity + , TestLabel "C9 taoLinearDeterministic" taoLinearDeterministic + , TestLabel "D1 taoLinearTinyCareSet" taoLinearTinyCareSet + , TestLabel "E1 categoricalBreimanBinary" testCategoricalBreimanBinary + , TestLabel + "E2 categoricalSubsetsMulticlassLowCard" + testCategoricalSubsetsMulticlassLowCard + , TestLabel + "E3 categoricalSingletonsMulticlassHighCard" + testCategoricalSingletonsMulticlassHighCard + , TestLabel "E4 categoricalCardZero" testCategoricalCardZero + , TestLabel "E5 categoricalNullableBinary" testCategoricalNullableBinary + , -- PR 2 extended: threshold-consolidation rewrite (positive cases). + TestLabel "F1 threshAndLeq" threshAndLeq + , TestLabel "F2 threshOrLeq" threshOrLeq + , TestLabel "F3 threshAndLt" threshAndLt + , TestLabel "F4 threshOrLt" threshOrLt + , TestLabel "F5 threshAndGeq" threshAndGeq + , TestLabel "F6 threshOrGeq" threshOrGeq + , TestLabel "F7 threshAndGt" threshAndGt + , TestLabel "F8 threshOrGt" threshOrGt + , -- PR 2 extended: negative cases (rewrite must NOT fire). + TestLabel "F9 threshNegMixedDirection" threshNegMixedDirection + , TestLabel "F10 threshNegCrossColumn" threshNegCrossColumn + , TestLabel "F11 threshNegMixedOpFamily" threshNegMixedOpFamily + , TestLabel "F12 threshNegEqualityOp" threshNegEqualityOp + , TestLabel "F13 threshNegLitOnLeft" threshNegLitOnLeft + , TestLabel "F14 threshNegNonLiteralRhs" threshNegNonLiteralRhs + , TestLabel "F15 threshSemanticPreservation" threshSemanticPreservation + ] diff --git a/tests-internal/Learn/EdgeCases.hs b/tests-internal/Learn/EdgeCases.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Learn/EdgeCases.hs @@ -0,0 +1,441 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Edge-case / degeneracy tests (category 7) and numerical-stability tests +(category 8) for @dataframe-learn@. Each asserts the mathematically correct +result or a specific documented degenerate behaviour — never "it ran". +-} +module Learn.EdgeCases (tests) where + +import qualified Data.Map.Strict as M +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr) +import DataFrame.Internal.Interpreter (interpret) +import qualified DataFrameApi as D + +import DataFrame.GMM +import DataFrame.KMeans +import DataFrame.LinearAlgebra (logSumExp) +import DataFrame.LinearAlgebra.Eigen (jacobiEigenSym) +import DataFrame.LinearAlgebra.Solve (choleskySolve) +import DataFrame.LinearModel +import DataFrame.LinearSolver (sigmoid) +import DataFrame.PCA + +import DataFrame.Internal.Statistics (correlation', variance') + +import Test.HUnit + +-- Helpers (mirrors Learn.Models) -------------------------------------------- + +interpD :: D.DataFrame -> Expr Double -> [Double] +interpD df e = case interpret @Double df e of + Right (TColumn c) -> either (const []) VU.toList (toVector @Double @VU.Vector c) + Left err -> error (show err) + +interpI :: D.DataFrame -> Expr Int -> [Int] +interpI df e = case interpret @Int df e of + Right (TColumn c) -> either (const []) VU.toList (toVector @Int @VU.Vector c) + Left err -> error (show err) + +mat :: [[Double]] -> V.Vector (VU.Vector Double) +mat = V.fromList . map VU.fromList + +close :: Double -> Double -> Double -> Bool +close tol a b = abs (a - b) <= tol + +finite :: Double -> Bool +finite x = not (isNaN x) && not (isInfinite x) + +-- =========================================================================== +-- Category 8: numerical stability +-- =========================================================================== + +{- logSumExp on a large-positive cluster: log(e^1000+e^1001+e^1002) = + 1002 + log(e^-2+e^-1+1) ≈ 1002.40760596. A naive log.sum.map exp + overflows to +Inf here. -} +testLogSumExpLargePositive :: Test +testLogSumExpLargePositive = TestCase $ do + let lse = logSumExp (VU.fromList [1000, 1001, 1002]) + expected = 1002 + log (exp (-2) + exp (-1) + 1) + assertBool "lse large-positive is finite" (finite lse) + assertBool + "lse [1000,1001,1002] == 1002 + log(e^-2+e^-1+1)" + (close 1e-9 lse expected) + assertBool "lse literal ~ 1002.40760596" (close 1e-7 lse 1002.4076059644443) + +{- logSumExp on a large-negative cluster: log(e^-1000+e^-1001+e^-1002) = + -1000 + log(1+e^-1+e^-2) ≈ -999.59239403. A naive impl underflows every + term to 0 -> log 0 = -Inf. -} +testLogSumExpLargeNegative :: Test +testLogSumExpLargeNegative = TestCase $ do + let lse = logSumExp (VU.fromList [-1000, -1001, -1002]) + expected = -1000 + log (1 + exp (-1) + exp (-2)) + assertBool "lse large-negative is finite" (finite lse) + assertBool + "lse [-1000,-1001,-1002] == -1000 + log(1+e^-1+e^-2)" + (close 1e-9 lse expected) + assertBool "lse literal ~ -999.59239403" (close 1e-7 lse (-999.5923940355557)) + +{- logSumExp of a single element is that element exactly (m + log 1). -} +testLogSumExpSingleton :: Test +testLogSumExpSingleton = TestCase $ do + assertBool "lse [42] == 42" (close 0 (logSumExp (VU.fromList [42])) 42) + +{- sigmoid at extreme arguments must stay in [0,1] and not NaN/Inf: + sigmoid(1000) ≈ 1, sigmoid(-1000) ≈ 0. The stable branch avoids the + overflow a naive 1/(1+exp(-z)) hits. -} +testSigmoidExtreme :: Test +testSigmoidExtreme = TestCase $ do + let sp = sigmoid 1000 + sn = sigmoid (-1000) + assertBool "sigmoid(1000) finite" (finite sp) + assertBool "sigmoid(-1000) finite" (finite sn) + assertBool "sigmoid(1000) in [0,1]" (sp >= 0 && sp <= 1) + assertBool "sigmoid(-1000) in [0,1]" (sn >= 0 && sn <= 1) + assertBool "sigmoid(1000) == 1" (close 1e-12 sp 1) + assertBool "sigmoid(-1000) == 0" (close 1e-12 sn 0) + assertBool "sigmoid(0) == 0.5" (close 1e-15 (sigmoid 0) 0.5) + assertBool + "sigmoid antisymmetric at 3" + (close 1e-12 (sigmoid (-3)) (1 - sigmoid 3)) + +{- Variance of large-but-low-variance data [1e8+1,+2,+3]: true sample variance + is 1.0. A naive sum-of-squares formula loses all precision (catastrophic + cancellation); Welford keeps it ~1. -} +testVarianceCatastrophicCancellation :: Test +testVarianceCatastrophicCancellation = TestCase $ do + let xs = VU.fromList [1e8 + 1, 1e8 + 2, 1e8 + 3] :: VU.Vector Double + v = variance' xs + assertBool "variance finite" (finite v) + assertBool "variance is positive" (v > 0) + assertBool "variance of shifted {1,2,3} == 1.0" (close 1e-6 v 1.0) + +{- Variance of an identical column is exactly 0 (not a tiny negative from + cancellation). -} +testVarianceConstant :: Test +testVarianceConstant = TestCase $ do + let v = variance' (VU.replicate 100 (7.0 :: Double)) + assertEqual "variance of constant column is 0" 0 v + +{- Variance of fewer than two samples is defined to be 0 (computeVariance guard), + not NaN from a /0. -} +testVarianceSingleton :: Test +testVarianceSingleton = TestCase $ do + assertEqual + "variance of one sample is 0" + 0 + (variance' (VU.fromList [3.5 :: Double])) + +{- Correlation of a perfectly linear pair is exactly +1 (and -1 reversed), + computed stably. y = 2x+1 over a spread of x. -} +testCorrelationPerfect :: Test +testCorrelationPerfect = TestCase $ do + let xs = VU.fromList [1, 2, 3, 4, 5] :: VU.Vector Double + ys = VU.map (\x -> 2 * x + 1) xs + yneg = VU.map (\x -> -(2 * x) + 1) xs + case correlation' xs ys of + Just r -> assertBool "corr(x, 2x+1) == 1" (close 1e-9 r 1) + Nothing -> assertFailure "expected a correlation" + case correlation' xs yneg of + Just r -> assertBool "corr(x, -2x+1) == -1" (close 1e-9 r (-1)) + Nothing -> assertFailure "expected a correlation" + +{- Degeneracy: correlation against a constant column. The denominator is 0, so + the library returns Just NaN — this pins the actual (non-throwing) behaviour; + a future Nothing/0 would flag the contract change. -} +testCorrelationConstantColumnIsNaN :: Test +testCorrelationConstantColumnIsNaN = TestCase $ do + let xs = VU.fromList [1, 2, 3, 4, 5] :: VU.Vector Double + ys = VU.replicate 5 (9.0 :: Double) + case correlation' xs ys of + Just r -> + assertBool + "corr with a zero-variance column is NaN (degenerate denom)" + (isNaN r) + Nothing -> assertFailure "correlation' returned Nothing (contract changed)" + +{- correlation' on n<2 returns Nothing (guarded), not a NaN. -} +testCorrelationTooFew :: Test +testCorrelationTooFew = TestCase $ do + assertEqual + "correlation of one point is Nothing" + Nothing + (correlation' (VU.fromList [1]) (VU.fromList [2])) + +-- =========================================================================== +-- Category 8: stability inside the model expr layer +-- =========================================================================== + +{- Logistic probability expressions at extreme feature values (|x| ~ 1e6) must + remain finite and in [0,1]; a non-stable 1/(1+exp(-margin)) would overflow. + Also checks both one-vs-rest class probabilities are present. -} +testLogisticProbsExtremeFeatures :: Test +testLogisticProbsExtremeFeatures = TestCase $ do + let trainDf = + D.fromNamedColumns + [ ("x", DI.fromList ([-3, -2, -1, -0.5, 0.5, 1, 2, 3] :: [Double])) + , ("label", DI.fromList ([0, 0, 0, 0, 1, 1, 1, 1] :: [Int])) + ] + m = fit defaultLogisticConfig (F.col @Int "label") trainDf + probs = logisticProbExprs m + extremeDf = + D.fromNamedColumns + [("x", DI.fromList ([-1e6, -1, 0, 1, 1e6] :: [Double]))] + assertEqual "two one-vs-rest probability exprs" 2 (M.size probs) + let allProbVals = + concat [interpD extremeDf e | e <- M.elems probs] + assertBool "all logistic probabilities finite" (all finite allProbVals) + assertBool + "all logistic probabilities in [0,1]" + (all (\p -> p >= 0 && p <= 1) allProbVals) + +-- =========================================================================== +-- Category 7: edge cases / degeneracy on models +-- =========================================================================== + +{- One-row OLS (x=2, y=7): under-determined, so olsSolve falls back to + ridge(1e-8). The fit must interpolate the single point (predict 7 at x=2) + and never be NaN/Inf. -} +testOLSOneRow :: Test +testOLSOneRow = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([2] :: [Double])) + , ("y", DI.fromList ([7] :: [Double])) + ] + m = fit defaultLinearConfig (F.col @Double "y") df + preds = interpD df (predict m) + assertBool "single OLS prediction finite" (all finite preds) + assertBool + "OLS fits the single training point" + (case preds of [p] -> close 1e-3 p 7; _ -> False) + +{- All-identical labels in logistic regression: every row is class 1. There is + exactly one class, so predict must return that class for every row (no + division-by-zero, no crash, no spurious second class). -} +testLogisticSingleClass :: Test +testLogisticSingleClass = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([-2, -1, 0, 1, 2] :: [Double])) + , ("label", DI.fromList ([1, 1, 1, 1, 1] :: [Int])) + ] + m = fit defaultLogisticConfig (F.col @Int "label") df + preds = interpI df (predict m) + assertEqual + "single-class logistic predicts that class everywhere" + [1, 1, 1, 1, 1] + preds + +{- Perfectly separable logistic data plus a constant (zero-variance) feature: + the constant column is dropped, the informative one still separates, and + prediction recovers the labels exactly. -} +testLogisticConstantFeatureIgnored :: Test +testLogisticConstantFeatureIgnored = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([-3, -2, -1, -0.5, 0.5, 1, 2, 3] :: [Double])) + , ("const", DI.fromList (replicate 8 (5.0 :: Double))) + , ("label", DI.fromList ([0, 0, 0, 0, 1, 1, 1, 1] :: [Int])) + ] + m = fit defaultLogisticConfig (F.col @Int "label") df + preds = interpI df (predict m) + assertEqual + "logistic ignores constant column and separates" + [0, 0, 0, 0, 1, 1, 1, 1] + preds + +{- Linear regression with an irrelevant constant feature: the constant column is + near-constant (variance 0) and must get a zero weight, while the informative + coefficient is recovered. y = 3x + 4 exactly. -} +testLinearConstantFeatureZeroWeight :: Test +testLinearConstantFeatureZeroWeight = TestCase $ do + let xs = [1, 2, 3, 4, 5, 6] :: [Double] + df = + D.fromNamedColumns + [ ("x", DI.fromList xs) + , ("const", DI.fromList (replicate 6 (5.0 :: Double))) + , ("y", DI.fromList [3 * x + 4 | x <- xs]) + ] + m = fit defaultLinearConfig (F.col @Double "y") df + coefs = VU.toList (regCoef m) + names = V.toList (regFeatureNames m) + byName = zip names coefs + case lookup "const" byName of + Just w -> assertBool "constant feature gets ~zero weight" (close 1e-6 w 0) + Nothing -> assertFailure "const column missing from feature names" + let preds = interpD df (predict m) + truth = [3 * x + 4 | x <- xs] + assertBool + "regression with constant feature still fits y=3x+4" + (and (zipWith (close 1e-4) preds truth)) + +{- k-means with k > n: the library clamps k = min k (max 1 n), so 3 rows with + kmK=10 must yield exactly 3 centres, in-range labels, and finite centres — + not a crash or degenerate centres. -} +testKMeansKGreaterThanN :: Test +testKMeansKGreaterThanN = TestCase $ do + let df = + D.fromNamedColumns + [ ("a", DI.fromList ([0, 5, 10] :: [Double])) + , ("b", DI.fromList ([0, 5, 10] :: [Double])) + ] + cfg = defaultKMeansConfig{kmK = 10, kmNInit = 3, kmSeed = 1} + m = fit cfg [F.col @Double "a", F.col @Double "b"] df + assertEqual "k clamped to n=3 centres" 3 (V.length (kmCenters m)) + let labels = VU.toList (kmLabels m) + assertBool "all labels in [0,2]" (all (\l -> l >= 0 && l < 3) labels) + assertBool + "all centre coordinates finite" + (all (VU.all finite) (V.toList (kmCenters m))) + assertBool "k=n clustering has ~zero inertia" (close 1e-9 (kmInertia m) 0) + +{- k-means on an all-identical feature set: any clustering has inertia 0 and + centres equal that point, staying finite (no NaN from an empty cluster or a + zero distance sum in k-means++ sampling). -} +testKMeansAllIdentical :: Test +testKMeansAllIdentical = TestCase $ do + let df = + D.fromNamedColumns + [ ("a", DI.fromList (replicate 6 (4.0 :: Double))) + , ("b", DI.fromList (replicate 6 (4.0 :: Double))) + ] + cfg = defaultKMeansConfig{kmK = 2, kmNInit = 3, kmSeed = 1} + m = fit cfg [F.col @Double "a", F.col @Double "b"] df + assertBool + "centres finite on degenerate identical data" + (all (VU.all finite) (V.toList (kmCenters m))) + assertBool "inertia is 0 for identical points" (close 1e-12 (kmInertia m) 0) + let assigns = interpI df (predict m) + assertBool "assignments in [0,1]" (all (\l -> l >= 0 && l < 2) assigns) + +{- GMM with k > n: clamped like k-means (k = min k (max 1 n)). Two rows, gmmK=5 + must yield 2 components, finite weights summing to ~1, and a finite log + likelihood -- not a crash or NaN. -} +testGMMKGreaterThanN :: Test +testGMMKGreaterThanN = TestCase $ do + let df = + D.fromNamedColumns + [ ("a", DI.fromList ([0, 10] :: [Double])) + , ("b", DI.fromList ([0, 10] :: [Double])) + ] + cfg = defaultGMMConfig{gmmK = 5, gmmSeed = 1} + m = fit cfg [F.col @Double "a", F.col @Double "b"] df + assertEqual "GMM k clamped to n=2" 2 (VU.length (gmmWeights m)) + assertBool "GMM weights finite" (VU.all finite (gmmWeights m)) + assertBool "GMM weights sum to ~1" (close 1e-6 (VU.sum (gmmWeights m)) 1) + assertBool "GMM log-likelihood finite" (finite (gmmLogLikelihood m)) + +{- PCA on one informative axis plus a constant column: the constant's explained + ratio must be ~0 while ratios sum to ~1 and stay finite. Catches a + divide-by-zero in ratio normalisation when an eigenvalue is 0. -} +testPCAConstantColumn :: Test +testPCAConstantColumn = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([-2, -1, 0, 1, 2] :: [Double])) + , ("const", DI.fromList (replicate 5 (3.0 :: Double))) + ] + m = + fit + (PCAConfig (NComp 2) False) + [F.col @Double "x", F.col @Double "const"] + df + ratio = VU.toList (pcaExplainedVarianceRatio m) + assertBool "explained ratios finite" (all finite ratio) + assertBool "explained ratios sum to ~1" (close 1e-9 (sum ratio) 1) + r0 <- case ratio of + (r : _) -> pure r + [] -> assertFailure "expected explained-variance ratios" + assertBool "first PC explains ~all variance" (close 1e-9 r0 1) + let es = map snd (pcaExprs m) + assertBool + "pca exprs finite on constant column" + (all (all finite . interpD df) es) + +{- PCA on extreme-scale features (1e8-shifted), no standardisation: components + must stay unit-length and finite — catches cancellation in the covariance / + eigendecomposition at large magnitudes. -} +testPCAExtremeScale :: Test +testPCAExtremeScale = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([1e8 + 1, 1e8 + 2, 1e8 + 3, 1e8 + 4] :: [Double])) + , ("y", DI.fromList ([1e8 + 1, 1e8 + 2, 1e8 + 3, 1e8 + 4] :: [Double])) + ] + m = + fit + (PCAConfig (NComp 1) False) + [F.col @Double "x", F.col @Double "y"] + df + comp0 = pcaComponents m V.! 0 + nrm = sqrt (VU.sum (VU.map (^ (2 :: Int)) comp0)) + assertBool "component finite at 1e8 scale" (VU.all finite comp0) + assertBool "component unit length at 1e8 scale" (close 1e-6 nrm 1) + assertBool + "explained ratio finite at 1e8 scale" + (VU.all finite (pcaExplainedVarianceRatio m)) + +-- =========================================================================== +-- Category 7/8: linear-algebra degeneracy +-- =========================================================================== + +{- choleskySolve on a non-positive-definite (zero) matrix returns Nothing rather + than crashing or producing NaNs. A 2x2 all-zero matrix has a zero pivot. -} +testCholeskyNonPD :: Test +testCholeskyNonPD = TestCase $ do + assertEqual + "cholesky of singular zero matrix is Nothing" + Nothing + (choleskySolve (mat [[0, 0], [0, 0]]) (VU.fromList [1, 1])) + +{- Jacobi eigendecomposition of a 1x1 matrix (one-column degenerate case): + eigenvalue is the single entry, eigenvector is [1] (sign-canonicalised). -} +testJacobi1x1 :: Test +testJacobi1x1 = TestCase $ do + let (ev, vecs) = jacobiEigenSym (mat [[5]]) + assertBool "1x1 eigenvalue is the entry" (close 1e-12 (ev VU.! 0) 5) + assertBool "1x1 eigenvector is [1]" (close 1e-12 ((vecs V.! 0) VU.! 0) 1) + +{- Jacobi on an identity matrix: both eigenvalues are exactly 1 and the result is + finite (the rotation angle code must not divide by zero when off-diagonals are + already 0). -} +testJacobiIdentity :: Test +testJacobiIdentity = TestCase $ do + let (ev, vecs) = jacobiEigenSym (mat [[1, 0], [0, 1]]) + assertBool "identity eigenvalues both 1" (VU.all (close 1e-12 1) ev) + assertBool "identity eigenvectors finite" (all (VU.all finite) (V.toList vecs)) + +tests :: [Test] +tests = + [ testLogSumExpLargePositive + , testLogSumExpLargeNegative + , testLogSumExpSingleton + , testSigmoidExtreme + , testVarianceCatastrophicCancellation + , testVarianceConstant + , testVarianceSingleton + , testCorrelationPerfect + , testCorrelationConstantColumnIsNaN + , testCorrelationTooFew + , testLogisticProbsExtremeFeatures + , testOLSOneRow + , testLogisticSingleClass + , testLogisticConstantFeatureIgnored + , testLinearConstantFeatureZeroWeight + , testKMeansKGreaterThanN + , testKMeansAllIdentical + , testGMMKGreaterThanN + , testPCAConstantColumn + , testPCAExtremeScale + , testCholeskyNonPD + , testJacobi1x1 + , testJacobiIdentity + ] diff --git a/tests-internal/Learn/NumericalRigor.hs b/tests-internal/Learn/NumericalRigor.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Learn/NumericalRigor.hs @@ -0,0 +1,408 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Numerical-rigor suite: gradient checks (cat 4), reproducibility of the +stochastic models (cat 16), and statistical properties of the RNG and splitters +(cat 5). Every test is constructed to FAIL on a real bug. +-} +module Learn.NumericalRigor (tests) where + +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import qualified DataFrameApi as D + +import DataFrame.GMM +import DataFrame.KMeans +import DataFrame.LinearSolver.Loss ( + SmoothLoss (..), + logisticLoss, + sigmoid, + sqHingeLoss, + squaredLoss, + ) +import DataFrame.Random +import DataFrame.SVM.RFF + +import qualified Data.Vector.Unboxed as VU +import System.Random (mkStdGen) +import Test.HUnit + +-- --------------------------------------------------------------------------- +-- Independent reference loss VALUE functions. The library exposes only the +-- gradient 'slGradZ'; these reconstruct the scalar value so the finite +-- difference is computed independently of the analytic gradient. +-- --------------------------------------------------------------------------- + +-- | @½ (z - y)²@ (squaredLoss). +squaredValue :: Double -> Double -> Double +squaredValue y z = 0.5 * (z - y) * (z - y) + +-- | @log (1 + exp (-y z))@ (logisticLoss). softplus form, numerically stable. +logisticValue :: Double -> Double -> Double +logisticValue y z = + let m = negate (y * z) + in if m >= 0 then m + log1pExp (negate m) else log1pExp m + where + log1pExp x = log (1 + exp x) + +-- | @(max 0 (1 - y z))²@ (sqHingeLoss). +sqHingeValue :: Double -> Double -> Double +sqHingeValue y z = let m = 1 - y * z in if m > 0 then m * m else 0 + +-- | Central finite difference of @f@ in @z@ at @(y, z)@. +centralDiff :: + (Double -> Double -> Double) -> Double -> Double -> Double -> Double +centralDiff f eps y z = (f y (z + eps) - f y (z - eps)) / (2 * eps) + +{- | Gradient-check one 'SmoothLoss' against its reference value fn over a grid +of @(y, z)@ points. Fails if the analytic gradient differs from the finite +difference anywhere beyond @tol@. +-} +gradCheck :: + String -> + SmoothLoss -> + (Double -> Double -> Double) -> + Double -> + Test +gradCheck nm loss valueFn tol = + TestCase $ + mapM_ check [(y, z) | y <- [-1, 1], z <- zs] + where + zs = [-3.7, -2.1, -1.25, -0.6, -0.15, 0.22, 0.55, 1.4, 1.9, 2.8, 4.3] + check (y, z) = do + let analytic = slGradZ loss y z + fd = centralDiff valueFn 1e-6 y z + ok = abs (analytic - fd) <= tol * (1 + abs fd) + assertBool + ( nm + ++ " grad mismatch at (y=" + ++ show y + ++ ", z=" + ++ show z + ++ "): analytic=" + ++ show analytic + ++ " finite-diff=" + ++ show fd + ) + ok + +-- | The squared-loss gradient @z - y@ checked against ½(z-y)². +testSquaredGradient :: Test +testSquaredGradient = gradCheck "squared" squaredLoss squaredValue 1e-5 + +-- | The logistic gradient @-y·σ(-y z)@ checked against log(1+exp(-yz)). +testLogisticGradient :: Test +testLogisticGradient = gradCheck "logistic" logisticLoss logisticValue 1e-5 + +-- | The squared-hinge gradient @-2y·max(0,1-yz)@ checked against (max 0 (1-yz))². +testSqHingeGradient :: Test +testSqHingeGradient = gradCheck "squared_hinge" sqHingeLoss sqHingeValue 1e-5 + +{- | Sign sanity that a self-referential check can't pass: where the loss rises +in @z@ the gradient must be positive (and vice versa). Catches a flipped-sign +gradient even with the right magnitude. +-} +testGradientSigns :: Test +testGradientSigns = TestCase $ do + assertBool + "squared: grad>0 when z>y (loss rising)" + (slGradZ squaredLoss 1.0 3.0 > 0) + assertBool + "squared: grad<0 when z0" + (slGradZ logisticLoss (-1.0) 0.5 > 0) + assertBool + "sqHinge active margin y=+1: grad<0" + (slGradZ sqHingeLoss 1.0 0.0 < 0) + assertBool + "sqHinge satisfied margin: grad==0" + (slGradZ sqHingeLoss 1.0 5.0 == 0) + +-- --------------------------------------------------------------------------- +-- Statistical / distributional tests (cat 5). Tolerances are CI-derived from +-- the sample size (N = 100k → 6σ band ~ 5e-3). A broken sampler — constant or +-- biased — falls well outside. +-- --------------------------------------------------------------------------- + +-- | Draw @n@ uniforms threading the generator; returns the sample list. +drawUniforms :: Int -> Gen -> [Double] +drawUniforms n g0 = go n g0 [] + where + go 0 _ acc = acc + go k g acc = let (x, g') = nextDouble g in go (k - 1) g' (x : acc) + +mean :: [Double] -> Double +mean xs = sum xs / fromIntegral (length xs) + +variance :: [Double] -> Double +variance xs = + let m = mean xs + in sum [(x - m) * (x - m) | x <- xs] / fromIntegral (length xs) + +{- | @nextDouble@ is ~uniform on [0,1): mean ≈ 0.5 and variance ≈ 1/12, each +within a 6σ CI for 100k samples, AND every draw is in [0,1). A constant or +out-of-range sampler fails; a biased one (mean drifts) fails. +-} +testUniformDistribution :: Test +testUniformDistribution = TestCase $ do + let n = 100000 :: Int + xs = drawUniforms n (mkGen 20240613) + m = mean xs + v = variance xs + seMean = (1 / sqrt 12) / sqrt (fromIntegral n) + assertBool "all uniforms in [0,1)" (all (\x -> x >= 0 && x < 1) xs) + assertBool + ("uniform mean ~0.5, got " ++ show m) + (abs (m - 0.5) <= 6 * seMean) + assertBool + ("uniform variance ~1/12, got " ++ show v) + (abs (v - 1 / 12) <= 0.01) + assertBool "uniform actually varies" (maximum xs - minimum xs > 0.9) + +{- | Box-Muller @gaussianVector@ produces standard normals: sample mean ≈ 0 and +variance ≈ 1 over 100k draws (6σ band ~ 2e-2). Catches a Box-Muller bug that +shifts the mean or scales the variance. +-} +testGaussianMoments :: Test +testGaussianMoments = TestCase $ do + let n = 100000 :: Int + (vec, _) = gaussianVector n (mkGen 777) + xs = VU.toList vec + m = mean xs + v = variance xs + seMean = 1 / sqrt (fromIntegral n) + assertBool "gaussianVector length" (VU.length vec == n) + assertBool + "gaussian samples finite" + (all (\x -> not (isNaN x) && not (isInfinite x)) xs) + assertBool + ("gaussian mean ~0, got " ++ show m) + (abs m <= 6 * seMean) + assertBool + ("gaussian variance ~1, got " ++ show v) + (abs (v - 1) <= 0.03) + let tail2 = fromIntegral (length (filter (\x -> abs x > 2) xs)) / fromIntegral n + assertBool + ("gaussian |x|>2 frequency ~0.0455, got " ++ show tail2) + (abs (tail2 - 0.0455) <= 0.01) + +{- | @randomSplit frac@ over many seeds: the realized train fraction matches +@frac@ within a binomial CI, and train+test always sums to the input row count +(cat 2 invariant). A split that loses/dupes rows, or ignores @frac@, fails. +-} +testSplitProportions :: Test +testSplitProportions = TestCase $ do + let nRows = 4000 + frac = 0.7 :: Double + df = D.fromNamedColumns [("x", DI.fromList [1 .. nRows :: Int])] + seeds = [1 .. 25] :: [Int] + seFrac = sqrt (frac * (1 - frac) / fromIntegral nRows) + mapM_ + ( \s -> do + let (tr, te) = D.randomSplit (mkStdGen s) frac df + nTr = fst (D.dimensions tr) + nTe = fst (D.dimensions te) + realized = fromIntegral nTr / fromIntegral nRows + assertEqual + ("split preserves rows (seed " ++ show s ++ ")") + nRows + (nTr + nTe) + assertBool + ( "split fraction ~" + ++ show frac + ++ " (seed " + ++ show s + ++ "), got " + ++ show realized + ) + (abs (realized - frac) <= 5 * seFrac) + ) + seeds + +{- | k-means inertia on a clean, well-separated blob is stable across seeds and +never beats the true within-cluster sum of squares of the generating partition. +A broken inertia would report values below the optimum or wildly seed-dependent. +-} +testKMeansInertiaStable :: Test +testKMeansInertiaStable = TestCase $ do + let + as = [0, 0.1, -0.1, 0.05, -0.05, 10, 10.1, 9.9, 10.05, 9.95] :: [Double] + bs = [0, -0.1, 0.1, 0.05, -0.05, 10, 9.9, 10.1, 9.95, 10.05] :: [Double] + df = D.fromNamedColumns [("a", DI.fromList as), ("b", DI.fromList bs)] + fitSeed s = + kmInertia $ + fit + defaultKMeansConfig{kmK = 2, kmNInit = 1, kmSeed = s} + [F.col @Double "a", F.col @Double "b"] + df + inertias = map fitSeed [0 .. 19] + blob1 = take 5 (zip as bs) + blob2 = zip (drop 5 as) (drop 5 bs) + ssOf pts = + let mx = sum (map fst pts) / 5 + my = sum (map snd pts) / 5 + in sum [(x - mx) ^ (2 :: Int) + (y - my) ^ (2 :: Int) | (x, y) <- pts] + optimum = ssOf blob1 + ssOf blob2 + best = minimum inertias + worst = maximum inertias + assertBool + "k-means inertia finite" + (all (\i -> not (isNaN i) && not (isInfinite i)) inertias) + assertBool + ( "k-means inertia never below optimum (" + ++ show optimum + ++ "), best=" + ++ show best + ) + (best >= optimum - 1e-9) + assertBool + ( "k-means inertia stable across seeds, best=" + ++ show best + ++ " worst=" + ++ show worst + ) + (worst - best <= 1e-6 + 1e-3 * optimum) + +-- --------------------------------------------------------------------------- +-- Reproducibility tests (cat 16). Each compares two fits with the same seed +-- (determinism) and asserts a different seed CAN change the model, so a +-- constant-returning stub wouldn't pass. +-- --------------------------------------------------------------------------- + +blobsDF :: D.DataFrame +blobsDF = + D.fromNamedColumns + [ + ( "a" + , DI.fromList ([0, 0.2, -0.1, 0.1, 8, 8.1, 7.9, 8.2, 0.05, 8.05] :: [Double]) + ) + , + ( "b" + , DI.fromList ([0, -0.1, 0.2, 0.0, 5, 5.2, 4.9, 5.1, 0.1, 5.05] :: [Double]) + ) + ] + +-- | Many-cluster frame so k-means++ seeding genuinely depends on the seed. +spreadDF :: D.DataFrame +spreadDF = + D.fromNamedColumns + [ ("a", DI.fromList ([0, 1, 2, 10, 11, 12, 20, 21, 22, 30, 31, 32] :: [Double])) + , ("b", DI.fromList ([0, 1, 0, 10, 11, 10, 0, 1, 0, 10, 11, 10] :: [Double])) + ] + +{- | k-means: same seed → identical KMeansModel (Eq derived); a different seed +on a multi-blob frame CAN yield different centers (the seed actually drives +k-means++ init). Single-init so the seed is not washed out by nInit restarts. +-} +testKMeansReproducible :: Test +testKMeansReproducible = TestCase $ do + let cfg s = defaultKMeansConfig{kmK = 4, kmNInit = 1, kmMaxIter = 1, kmSeed = s} + run s = fit (cfg s) [F.col @Double "a", F.col @Double "b"] + a = run 1 spreadDF + b = run 1 spreadDF + assertEqual "k-means same seed identical model" a b + let centersFor s = kmCenters (run s spreadDF) + base = centersFor 1 + anyDiffer = any (\s -> centersFor s /= base) [2, 3, 5, 7, 11, 13] + assertBool "k-means: a different seed changes the model" anyDiffer + +{- | GMM: same seed → identical GMMModel (Eq derived); a different seed CAN move +the fitted means (seed drives the responsibility init via sampleIndices). +-} +testGMMReproducible :: Test +testGMMReproducible = TestCase $ do + let cfg s = defaultGMMConfig{gmmK = 2, gmmMaxIter = 1, gmmSeed = s} + run s = fit (cfg s) [F.col @Double "a", F.col @Double "b"] blobsDF + a = run 1 + b = run 1 + assertEqual "GMM same seed identical model" a b + let meansFor s = gmmMeans (run s) + base = meansFor 1 + anyDiffer = any (\s -> meansFor s /= base) [2, 3, 4, 5, 6, 7, 8] + assertBool "GMM: a different seed changes the means" anyDiffer + +{- | RFF-SVM: same seed → identical random projection AND identical fitted SVC +coefficients; a different seed changes the random Fourier features. The model +type only derives Show, so compare representative numeric fields directly. +-} +testRFFReproducible :: Test +testRFFReproducible = TestCase $ do + let clsDF = + D.fromNamedColumns + [ ("x", DI.fromList ([-3, -2, -1, -0.5, 0.5, 1, 2, 3] :: [Double])) + , ("label", DI.fromList ([0, 0, 0, 0, 1, 1, 1, 1] :: [Int])) + ] + cfg s = defaultRFFConfig{rffD = 40, rffGamma = 0.2, rffSeed = s} + run s = fit (cfg s) (F.col @Int "label") clsDF + a = run 5 + b = run 5 + assertEqual "RFF same seed: same projection B" (rffB a) (rffB b) + assertEqual "RFF same seed: same coefficients" (rffCoef a) (rffCoef b) + assertEqual "RFF same seed: same intercept" (rffIntercept a) (rffIntercept b) + let projFor s = rffB (run s) + base = projFor 5 + anyDiffer = any (\s -> projFor s /= base) [1, 2, 3, 6, 7, 8] + assertBool "RFF: a different seed changes the projection" anyDiffer + +-- --------------------------------------------------------------------------- +-- randomSplit determinism + row-count invariant (cat 16 + 2). +-- --------------------------------------------------------------------------- + +{- | @randomSplit@ with the same seed is bit-identical (compared via the +prettyPrinted frames), preserves total rows, and a different seed CAN change the +partition. +-} +testSplitReproducible :: Test +testSplitReproducible = TestCase $ do + let df = D.fromNamedColumns [("x", DI.fromList [1 .. 200 :: Int])] + (tr1, te1) = D.randomSplit (mkStdGen 42) 0.6 df + (tr2, te2) = D.randomSplit (mkStdGen 42) 0.6 df + assertBool "randomSplit same seed: same train" (tr1 == tr2) + assertBool "randomSplit same seed: same test" (te1 == te2) + assertEqual + "randomSplit preserves row count" + 200 + (fst (D.dimensions tr1) + fst (D.dimensions te1)) + let trainFor s = fst (D.randomSplit (mkStdGen s) 0.6 df) + base = trainFor 42 + anyDiffer = any (\s -> trainFor s /= base) [1, 2, 3, 7, 99] + assertBool "randomSplit: a different seed changes the partition" anyDiffer + +-- --------------------------------------------------------------------------- +-- A self-consistency sanity for the helpers above so a broken reference value +-- fn cannot make the gradient checks vacuous: the reference squared value must +-- actually be ½(z-y)² at a known point. +-- --------------------------------------------------------------------------- +testReferenceValueSanity :: Test +testReferenceValueSanity = TestCase $ do + assertBool + "squaredValue at (y=2,z=5) == 4.5" + (abs (squaredValue 2 5 - 4.5) < 1e-12) + assertBool + "logisticValue at (y=1,z=0) == log 2" + (abs (logisticValue 1 0 - log 2) < 1e-12) + assertBool "sqHingeValue at (y=1,z=0) == 1" (abs (sqHingeValue 1 0 - 1) < 1e-12) + assertBool "sigmoid 0 == 0.5" (abs (sigmoid 0 - 0.5) < 1e-12) + +tests :: [Test] +tests = + [ testReferenceValueSanity + , testSquaredGradient + , testLogisticGradient + , testSqHingeGradient + , testGradientSigns + , testUniformDistribution + , testGaussianMoments + , testSplitProportions + , testKMeansInertiaStable + , testKMeansReproducible + , testGMMReproducible + , testRFFReproducible + , testSplitReproducible + ] diff --git a/tests-internal/Learn/Numerics.hs b/tests-internal/Learn/Numerics.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Learn/Numerics.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module Learn.Numerics (tests) where + +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import DataFrame.Model (fit, predict) +import Test.HUnit + +import DataFrame.LinearAlgebra +import DataFrame.LinearAlgebra.Eigen +import DataFrame.LinearAlgebra.Solve +import DataFrame.Random + +mat :: [[Double]] -> Matrix +mat = V.fromList . map VU.fromList + +approx :: Double -> Double -> Double -> Bool +approx tol a b = abs (a - b) <= tol + +vApprox :: Double -> VU.Vector Double -> VU.Vector Double -> Bool +vApprox tol a b = + VU.length a == VU.length b && VU.and (VU.zipWith (approx tol) a b) + +testQR :: Test +testQR = TestCase $ do + let a = mat [[1, 1], [1, 2], [1, 3], [1, 4]] + b = VU.fromList [3, 5, 7, 9] + case qrLeastSquares a b of + Right x -> + assertBool "QR recovers [1,2]" (vApprox 1e-9 x (VU.fromList [1, 2])) + Left cols -> assertFailure ("unexpected rank deficiency: " ++ show cols) + +testQRRankDeficient :: Test +testQRRankDeficient = TestCase $ do + let a = mat [[1, 2], [2, 4], [3, 6]] + case qrLeastSquares a (VU.fromList [1, 2, 3]) of + Left _ -> pure () + Right _ -> assertFailure "expected rank deficiency on collinear columns" + +testCholesky :: Test +testCholesky = TestCase $ do + let a = mat [[4, 2], [2, 3]] + case choleskySolve a (VU.fromList [2, 1]) of + Just x -> + assertBool "cholesky solves" (vApprox 1e-9 x (VU.fromList [0.5, 0])) + Nothing -> assertFailure "expected PD" + assertEqual "non-PD rejected" Nothing (cholesky (mat [[1, 2], [2, 1]])) + +testJacobi :: Test +testJacobi = TestCase $ do + let (ev, vecs) = jacobiEigenSym (mat [[2, 1], [1, 2]]) + assertBool "eigenvalues [3,1]" (vApprox 1e-9 ev (VU.fromList [3, 1])) + let v0 = vecs V.! 0 + recon = matVec (mat [[2, 1], [1, 2]]) v0 + assertBool + "A v = lambda v" + (vApprox 1e-9 recon (scaleV (ev VU.! 0) v0)) + +testPowerIter :: Test +testPowerIter = TestCase $ do + let (lam, _) = powerIterTop 200 (mat [[2, 1], [1, 2]]) + assertBool "dominant eigenvalue ~3" (approx 1e-6 lam 3) + +testLogSumExp :: Test +testLogSumExp = TestCase $ do + let xs = VU.fromList [1000, 1001, 1002] + lse = logSumExp xs + assertBool "logSumExp stable, finite" (not (isNaN lse) && not (isInfinite lse)) + assertBool "logSumExp >= max" (lse >= 1002) + +testRngDeterminism :: Test +testRngDeterminism = TestCase $ do + let (a, _) = shuffleInts 50 (mkGen 11) + (b, _) = shuffleInts 50 (mkGen 11) + assertEqual "same seed same shuffle" a b + assertBool + "is a permutation" + (VU.toList (VU.modify (const (pure ())) a) /= [] && all (`VU.elem` a) [0 .. 49]) + +testRngRange :: Test +testRngRange = TestCase $ do + let go 0 _ acc = acc + go k g acc = + let (x, g') = nextIntR (3, 7) g in go (k - 1 :: Int) g' (x : acc) + xs = go 1000 (mkGen 5) [] + assertBool "ints within range" (all (\x -> x >= 3 && x <= 7) xs) + +tests :: [Test] +tests = + [ testQR + , testQRRankDeficient + , testCholesky + , testJacobi + , testPowerIter + , testLogSumExp + , testRngDeterminism + , testRngRange + ] diff --git a/tests-internal/Learn/Symbolic.hs b/tests-internal/Learn/Symbolic.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Learn/Symbolic.hs @@ -0,0 +1,138 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Learn.Symbolic (tests) where + +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (TypedColumn (..), toVector) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr) +import DataFrame.Internal.Interpreter (interpret) +import qualified DataFrameApi as D + +import DataFrame.PCA.Kernel +import DataFrame.SVM.RFF +import DataFrame.SymbolicRegression +import DataFrame.SymbolicRegression.Expr +import DataFrame.SymbolicRegression.Optimize (meanSquaredError) +import DataFrame.SymbolicRegression.Simplify (simplify) + +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Test.HUnit + +interpD :: D.DataFrame -> Expr Double -> [Double] +interpD df e = case interpret @Double df e of + Right (TColumn c) -> either (const []) VU.toList (toVector @Double @VU.Vector c) + Left err -> error (show err) + +interpI :: D.DataFrame -> Expr Int -> [Int] +interpI df e = case interpret @Int df e of + Right (TColumn c) -> either (const []) VU.toList (toVector @Int @VU.Vector c) + Left err -> error (show err) + +testSRRecovers :: Test +testSRRecovers = TestCase $ do + let xs = [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6] :: [Double] + df = + D.fromNamedColumns + [ ("x", DI.fromList xs) + , ("y", DI.fromList [x * x + x | x <- xs]) + ] + m = + fit + defaultSRConfig + { srSeed = 3 + , srGenerations = 60 + , srPopSize = 300 + , srUnaryOps = [] + } + (F.col @Double "y") + df + preds = interpD df (srBest m) + truth = interpD df (F.col @Double "y") + err = sum (zipWith (\p t -> (p - t) ^ (2 :: Int)) preds truth) / 10 + assertBool "SR recovers x*x+x to low error" (err < 1e-6) + assertBool "Pareto front non-empty" (not (null (srPareto m))) + +testSRDeterminism :: Test +testSRDeterminism = TestCase $ do + let xs = [1 .. 8] :: [Double] + df = + D.fromNamedColumns + [ ("x", DI.fromList xs) + , ("y", DI.fromList (map (\x -> 2 * x + 1) xs)) + ] + run = + fit + defaultSRConfig{srSeed = 9, srGenerations = 20} + (F.col @Double "y") + df + assertEqual "SR deterministic best MSE" (srBestMSE run) (srBestMSE (rerun df)) + where + rerun = + fit + defaultSRConfig{srSeed = 9, srGenerations = 20} + (F.col @Double "y") + +testSimplifyPreservesEval :: Test +testSimplifyPreservesEval = TestCase $ do + let feats = V.singleton (VU.fromList [1, 2, 3, 4 :: Double]) + n = 4 + e = SBin SAdd (SBin SMul (SVar 0) (SConst 1)) (SConst 0) + target = VU.fromList [1, 2, 3, 4] + assertBool + "simplify preserves evaluation" + ( abs + (meanSquaredError feats n target e - meanSquaredError feats n target (simplify e)) + < 1e-12 + ) + assertBool "simplify is size non-increasing" (srSize (simplify e) <= srSize e) + assertEqual "simplify idempotent" (simplify e) (simplify (simplify e)) + +testKernelPCA :: Test +testKernelPCA = TestCase $ do + let df = + D.fromNamedColumns + [ ("a", DI.fromList ([0, 0.2, -0.1, 0.1, 8, 8.1, 7.9, 8.2] :: [Double])) + , ("b", DI.fromList ([0, -0.1, 0.2, 0.0, 5, 5.2, 4.9, 5.1] :: [Double])) + ] + m = + fit + defaultKernelPCAConfig{kpcaNComponents = 2, kpcaNLandmarks = 8, kpcaSeed = 1} + [F.col @Double "a", F.col @Double "b"] + df + pc1 = case kernelPCAExprs m of + ((_, e) : _) -> interpD df e + [] -> error "testKernelPCA: no kPCA components" + assertBool "kPCA finite" (not (any isNaN pc1)) + pc1First <- case pc1 of + (p : _) -> pure p + [] -> assertFailure "kPCA produced no projection" + assertBool + "kPCA first component separates blobs" + (signum pc1First /= signum (last pc1)) + +testRFFSVM :: Test +testRFFSVM = TestCase $ do + let df = + D.fromNamedColumns + [ ("x", DI.fromList ([-3, -2, -1, -0.5, 0.5, 1, 2, 3] :: [Double])) + , ("label", DI.fromList ([0, 0, 0, 0, 1, 1, 1, 1] :: [Int])) + ] + m = + fit + defaultRFFConfig{rffD = 80, rffGamma = 0.2, rffSeed = 2} + (F.col @Int "label") + df + preds = interpI df (predict m) + assertEqual "RFF SVM separates" [0, 0, 0, 0, 1, 1, 1, 1] preds + +tests :: [Test] +tests = + [ testSRRecovers + , testSRDeterminism + , testSimplifyPreservesEval + , testKernelPCA + , testRFFSVM + ] diff --git a/tests-internal/LinearSolver.hs b/tests-internal/LinearSolver.hs new file mode 100644 --- /dev/null +++ b/tests-internal/LinearSolver.hs @@ -0,0 +1,782 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module LinearSolver where + +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (getColumns) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.LinearSolver +import qualified DataFrameApi as D + +import Data.List (sort) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import System.Random (mkStdGen, randomR) +import Test.HUnit + +------------------------------------------------------------------------ +-- Test fixtures and helpers +------------------------------------------------------------------------ + +-- Generate n points with d features, each value uniform in [-1, 1], from a seed. +syntheticPoints :: Int -> Int -> Int -> V.Vector (VU.Vector Double) +syntheticPoints seed n d = + let (rows, _) = foldr step ([], mkStdGen seed) [1 .. n] + in V.fromList (take n rows) + where + step _ (acc, g) = + let (row, g') = genRow d g + in (row : acc, g') + genRow k g0 = go k g0 [] + where + go 0 g xs = (VU.fromList (reverse xs), g) + go i g xs = + let (v, g') = randomR (-1.0 :: Double, 1.0) g + in go (i - 1) g' (v : xs) + +-- Label each row by sign(w . x + b); +1 if score > 0, else -1. +labelsForHyperplane :: + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + Double -> + VU.Vector Double +labelsForHyperplane rows w b = + VU.generate + (V.length rows) + ( \i -> + let score = dotProduct w (rows V.! i) + b + in if score > 0 then 1 else -1 + ) + +-- Cosine similarity between two non-zero vectors. +cosineSim :: VU.Vector Double -> VU.Vector Double -> Double +cosineSim u v = + let nu = sqrt (dotProduct u u) + nv = sqrt (dotProduct v v) + in if nu == 0 || nv == 0 then 0 else dotProduct u v / (nu * nv) + +-- Predict +1 or -1 from a fitted LinearModel. +predict :: LinearModel -> VU.Vector Double -> Double +predict m x = + let score = dotProduct (lmWeights m) x + lmIntercept m + in if score > 0 then 1 else -1 + +-- Predict directly on standardized features (skipping de-standardization). +predictStandardized :: VU.Vector Double -> Double -> VU.Vector Double -> Double +predictStandardized w b x = + if dotProduct w x + b > 0 then 1 else -1 + +-- Average binary logistic loss at (w, b); the branch on @margin@ keeps +-- @log(1 + exp(-margin))@ numerically stable. +logisticLoss :: + V.Vector (VU.Vector Double) -> + VU.Vector Double -> + VU.Vector Double -> + Double -> + Double +logisticLoss features labels w b = + let n = V.length features + loss i = + let yi = labels VU.! i + row = features V.! i + margin = yi * (dotProduct w row + b) + in if margin >= 0 + then log (1 + exp (-margin)) + else (-margin) + log (1 + exp margin) + in sum [loss i | i <- [0 .. n - 1]] / fromIntegral n + +------------------------------------------------------------------------ +-- A1: Recover known hyperplane with no L1 +------------------------------------------------------------------------ + +testA1RecoverHyperplane :: Test +testA1RecoverHyperplane = TestCase $ do + let groundTruth = VU.fromList [0.7, -0.5] + groundBias = 0.3 + rows = syntheticPoints 1 200 2 + labels = labelsForHyperplane rows groundTruth groundBias + cfg = + defaultSolverConfig + { scL1Lambda = 0 + , scL2Lambda = 0 + , scMaxIter = 500 + , scTol = 1e-6 + } + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + cosSim = cosineSim (lmWeights model) groundTruth + sameSignAll = + all + (\i -> predict model (rows V.! i) == labels VU.! i) + [0 .. V.length rows - 1] + assertBool + ("recovered weights should align with ground truth (cos = " ++ show cosSim ++ ")") + (cosSim > 0.99) + assertBool "all training points predicted correctly" sameSignAll + +------------------------------------------------------------------------ +-- A2: L1 produces sparse weights +------------------------------------------------------------------------ + +testA2L1Sparsity :: Test +testA2L1Sparsity = TestCase $ do + let groundTruth = VU.fromList [0, 1.2, 0, 0, -1.5, 0, 0, 0, 0, 0] + groundBias = 0 + rows = syntheticPoints 7 500 10 + labels = labelsForHyperplane rows groundTruth groundBias + cfg = + defaultSolverConfig + { scL1Lambda = 0.1 + , scL2Lambda = 0 + , scMaxIter = 500 + , scTol = 1e-6 + } + names = V.fromList [T.pack ("f" ++ show i) | i <- [0 .. 9 :: Int]] + model = fitL1Logistic cfg rows labels names + ws = VU.toList (lmWeights model) + nonZeroIdxs = [i | (i, w) <- zip [0 :: Int ..] ws, w /= 0] + zeroIdxs = [i | (i, w) <- zip [0 :: Int ..] ws, w == 0] + assertBool + ( "informative feature 1 should have non-zero weight (got " + ++ show (ws !! 1) + ++ ")" + ) + (ws !! 1 /= 0) + assertBool + ( "informative feature 4 should have non-zero weight (got " + ++ show (ws !! 4) + ++ ")" + ) + (ws !! 4 /= 0) + let noiseFeatures = [0, 2, 3, 5, 6, 7, 8, 9] :: [Int] + noiseZero = length [i | i <- noiseFeatures, i `elem` zeroIdxs] + assertBool + ( "at least 6 noise features zeroed (got " + ++ show noiseZero + ++ "; non-zero idxs = " + ++ show nonZeroIdxs + ++ ")" + ) + (noiseZero >= 6) + +------------------------------------------------------------------------ +-- A3: Convergence on well-conditioned input +------------------------------------------------------------------------ + +testA3Convergence :: Test +testA3Convergence = TestCase $ do + let groundTruth = VU.fromList [1.0, -0.5, 0.7] + rows = syntheticPoints 2 300 3 + labels = labelsForHyperplane rows groundTruth 0 + cfg = + defaultSolverConfig + { scL1Lambda = 0.01 + , scL2Lambda = 0 + , scMaxIter = 1000 + , scTol = 1e-5 + } + model = fitL1Logistic cfg rows labels (V.fromList ["a", "b", "c"]) + (rowsStd, _, _, _) = standardize rows + ws = lmWeights model + b = lmIntercept model + loss0 = logisticLoss rowsStd labels (VU.replicate 3 0) 0 + lossFit = logisticLoss rows labels ws b + assertBool + ( "loss decreased from initial (initial=" + ++ show loss0 + ++ ", final=" + ++ show lossFit + ++ ")" + ) + (lossFit < loss0) + +------------------------------------------------------------------------ +-- A4: Final loss <= initial loss (monotone or near-monotone in FISTA) +------------------------------------------------------------------------ + +testA4LossNotIncreasing :: Test +testA4LossNotIncreasing = TestCase $ do + let groundTruth = VU.fromList [0.8, 0.4] + rows = syntheticPoints 3 100 2 + labels = labelsForHyperplane rows groundTruth 0 + cfg = defaultSolverConfig{scL1Lambda = 0.05, scL2Lambda = 0, scMaxIter = 100} + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + loss0 = logisticLoss rows labels (VU.replicate 2 0) 0 + lossFit = logisticLoss rows labels (lmWeights model) (lmIntercept model) + assertBool + ( "final loss must be <= initial loss (l0=" + ++ show loss0 + ++ ", lf=" + ++ show lossFit + ++ ")" + ) + (lossFit <= loss0 + 1e-9) + +------------------------------------------------------------------------ +-- A5: Degenerate input — all labels +1 +------------------------------------------------------------------------ + +testA5AllSameDirection :: Test +testA5AllSameDirection = TestCase $ do + let rows = syntheticPoints 4 50 3 + labels = VU.replicate 50 1.0 + cfg = defaultSolverConfig{scL1Lambda = 0.01, scL2Lambda = 0, scMaxIter = 100} + model = fitL1Logistic cfg rows labels (V.fromList ["a", "b", "c"]) + ws = VU.toList (lmWeights model) + b = lmIntercept model + anyNaN = any isNaN ws || isNaN b + anyInf = any isInfinite ws || isInfinite b + allPositive = all (\i -> predict model (rows V.! i) == 1) [0 .. V.length rows - 1] + assertBool "no NaN in weights/intercept" (not anyNaN) + assertBool "no Inf in weights/intercept" (not anyInf) + assertBool + "all-same labels should produce a positive-predicting model" + allPositive + +------------------------------------------------------------------------ +-- A6: Degenerate — empty input +------------------------------------------------------------------------ + +testA6Empty :: Test +testA6Empty = TestCase $ do + let cfg = defaultSolverConfig + emptyRows = V.empty :: V.Vector (VU.Vector Double) + emptyLabels = VU.empty :: VU.Vector Double + names = V.fromList ["a", "b"] + model = fitL1Logistic cfg emptyRows emptyLabels names + assertEqual + "empty input -> 2 zero weights" + (VU.fromList [0, 0]) + (lmWeights model) + assertEqual "empty input -> zero intercept" 0 (lmIntercept model) + +------------------------------------------------------------------------ +-- A7: Degenerate — constant feature +------------------------------------------------------------------------ + +testA7ConstantFeature :: Test +testA7ConstantFeature = TestCase $ do + let baseRows = syntheticPoints 5 100 1 + rows = + V.map + (\row -> VU.fromList (0.5 : VU.toList row)) + baseRows + groundTruth = VU.fromList [0.0, 1.0] + labels = labelsForHyperplane rows groundTruth 0 + cfg = + defaultSolverConfig + { scL1Lambda = 0.01 + , scL2Lambda = 0 + , scMaxIter = 300 + , scTol = 1e-6 + } + model = fitL1Logistic cfg rows labels (V.fromList ["constant", "signal"]) + ws = VU.toList (lmWeights model) + anyBad = any (\x -> isNaN x || isInfinite x) ws + w0 <- case ws of + (w : _) -> pure w + [] -> assertFailure "expected at least one weight" + assertBool + ("constant feature weight ~ 0 (got " ++ show w0 ++ ")") + (abs w0 < 1e-6) + assertBool + ("signal feature non-zero (got " ++ show (ws !! 1) ++ ")") + (ws !! 1 /= 0) + assertBool "no NaN/Inf" (not anyBad) + +------------------------------------------------------------------------ +-- A8: Numerical stability with large feature values +------------------------------------------------------------------------ + +testA8LargeValues :: Test +testA8LargeValues = TestCase $ do + let scale = 1000.0 :: Double + baseRows = syntheticPoints 6 100 2 + rows = V.map (VU.map (* scale)) baseRows + groundTruth = VU.fromList [0.5, -0.7] + labels = labelsForHyperplane rows groundTruth 0 + cfg = defaultSolverConfig{scL1Lambda = 0.01, scL2Lambda = 0, scMaxIter = 300} + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + ws = VU.toList (lmWeights model) + b = lmIntercept model + anyBad = any (\x -> isNaN x || isInfinite x) (b : ws) + sameSigns = + length + [ () | i <- [0 .. V.length rows - 1], predict model (rows V.! i) == labels VU.! i + ] + assertBool "no NaN/Inf with scaled features" (not anyBad) + assertBool + ( "should correctly classify the vast majority of rows (" + ++ show sameSigns + ++ "/100)" + ) + (sameSigns >= 90) + +------------------------------------------------------------------------ +-- A9: Standardization round-trip — recovered weights point in the true +-- direction even when raw-feature scales differ by orders of magnitude. +------------------------------------------------------------------------ + +testA9StandardizationRoundTrip :: Test +testA9StandardizationRoundTrip = TestCase $ do + let nRows = 80 :: Int + col0 = [fromIntegral i * 5 :: Double | i <- [0 .. nRows - 1]] + col1 = [fromIntegral i * 0.0025 :: Double | i <- [0 .. nRows - 1]] + rows = V.fromList [VU.fromList [c0, c1] | (c0, c1) <- zip col0 col1] + labels = + VU.fromList + [ if (c0 - 200) + 1000 * (c1 - 0.1) > 0 then 1.0 else -1.0 + | (c0, c1) <- zip col0 col1 + ] + cfg = + defaultSolverConfig + { scL1Lambda = 1.0e-4 + , scL2Lambda = 0 + , scMaxIter = 2000 + , scTol = 1.0e-7 + } + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + truthDir = VU.fromList [1.0, 1000.0] + cs = cosineSim (lmWeights model) truthDir + trainPreds = + [predict model (rows V.! i) | i <- [0 .. nRows - 1]] + trainLabs = + [labels VU.! i | i <- [0 .. nRows - 1]] + correct = + length + [() | (p, l) <- zip trainPreds trainLabs, p == l] + assertEqual "all training points correctly classified" nRows correct + assertBool + ( "recovered raw weights align with ground-truth direction across " + ++ "vastly different feature scales (cos = " + ++ show cs + ++ ")" + ) + (cs > 0.95) + +------------------------------------------------------------------------ +-- A10: Determinism — same input -> same output +------------------------------------------------------------------------ + +testA10Determinism :: Test +testA10Determinism = TestCase $ do + let groundTruth = VU.fromList [0.6, 0.4] + rows = syntheticPoints 9 60 2 + labels = labelsForHyperplane rows groundTruth 0 + cfg = defaultSolverConfig{scL1Lambda = 0.05, scL2Lambda = 0, scMaxIter = 200} + m1 = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + m2 = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + assertEqual "same input -> same weights" (lmWeights m1) (lmWeights m2) + assertEqual "same input -> same intercept" (lmIntercept m1) (lmIntercept m2) + +------------------------------------------------------------------------ +-- A11: Two-feature ground truth recovery (w_2/w_1 ratio) +------------------------------------------------------------------------ + +testA11GroundTruthRatio :: Test +testA11GroundTruthRatio = TestCase $ do + let groundTruth = VU.fromList [1.0, 2.0] + groundBias = -3.0 + n = 500 + baseRows = syntheticPoints 10 n 2 + rows = V.map (VU.map (* 3)) baseRows + labels = labelsForHyperplane rows groundTruth groundBias + cfg = + defaultSolverConfig + { scL1Lambda = 0.001 + , scL2Lambda = 0 + , scMaxIter = 1000 + , scTol = 1e-7 + } + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + ws = lmWeights model + b = lmIntercept model + ratio = (ws VU.! 1) / (ws VU.! 0) + biasRatio = b / (ws VU.! 0) + assertBool + ("w2/w1 should approximate 2.0 (got " ++ show ratio ++ ")") + (ratio > 1.7 && ratio < 2.3) + assertBool + ("b/w1 should approximate -3.0 (got " ++ show biasRatio ++ ")") + (biasRatio > -3.4 && biasRatio < -2.6) + +------------------------------------------------------------------------ +-- B1: modelToExpr produces a well-typed Expr Bool +------------------------------------------------------------------------ + +testB1ExprWellTyped :: Test +testB1ExprWellTyped = TestCase $ do + let model = + LinearModel + { lmWeights = VU.fromList [1.0, -2.0] + , lmIntercept = 0.5 + , lmFeatureNames = V.fromList ["x", "y"] + } + expr = modelToExpr model + df = + D.fromNamedColumns + [ ("x", DI.fromList ([0.0, 1.0, 2.0] :: [Double])) + , ("y", DI.fromList ([0.0, 0.0, 5.0] :: [Double])) + ] + manual = + [ (1.0 * 0.0 - 2.0 * 0.0 + 0.5) > 0 + , (1.0 * 1.0 - 2.0 * 0.0 + 0.5) > 0 + , (1.0 * 2.0 - 2.0 * 5.0 + 0.5) > 0 + ] + case interpret @Bool df expr of + Left e -> assertFailure ("interpret failed: " ++ show e) + Right (DI.TColumn col) -> case DI.toVector @Bool col of + Left e -> assertFailure ("toVector failed: " ++ show e) + Right vals -> + assertEqual "Expr matches manual evaluation" manual (V.toList vals) + +------------------------------------------------------------------------ +-- B2: Zero weights are dropped from the resulting Expr +------------------------------------------------------------------------ + +testB2ZeroWeightsPruned :: Test +testB2ZeroWeightsPruned = TestCase $ do + let model = + LinearModel + { lmWeights = VU.fromList [0.0, 1.5, 0.0] + , lmIntercept = 0.0 + , lmFeatureNames = V.fromList ["a", "b", "c"] + } + expr = modelToExpr model + cols = sort (getColumns expr) + assertEqual "only column b appears in the Expr" ["b"] cols + +------------------------------------------------------------------------ +-- A14: Constant feature at large raw value — weight must be exactly 0 +-- and no NaN/Inf leaks into the rest of the fit. +------------------------------------------------------------------------ + +testA14ConstantHugeValue :: Test +testA14ConstantHugeValue = TestCase $ do + let baseRows = syntheticPoints 14 100 1 + rows = + V.map + (\row -> VU.fromList (1.0e8 : VU.toList row)) + baseRows + labels = labelsForHyperplane rows (VU.fromList [0.0, 1.0]) 0 + cfg = defaultSolverConfig{scL1Lambda = 0.01, scL2Lambda = 0, scMaxIter = 300} + model = fitL1Logistic cfg rows labels (V.fromList ["constant", "signal"]) + ws = VU.toList (lmWeights model) + b = lmIntercept model + anyBad = any (\v -> isNaN v || isInfinite v) (b : ws) + assertBool "no NaN/Inf with constant-at-1e8 feature" (not anyBad) + w0 <- case ws of + (w : _) -> pure w + [] -> assertFailure "expected at least one weight" + assertEqual + "constant feature is dropped — weight is exactly zero" + 0 + w0 + assertBool + ("signal feature has non-zero weight (got " ++ show (ws !! 1) ++ ")") + (ws !! 1 /= 0) + +------------------------------------------------------------------------ +-- A15: Variance exactly zero (all rows identical for that column). +------------------------------------------------------------------------ + +testA15AllZeroFeature :: Test +testA15AllZeroFeature = TestCase $ do + let baseRows = syntheticPoints 15 80 1 + rows = + V.map + (\row -> VU.fromList (0.0 : VU.toList row)) + baseRows + labels = labelsForHyperplane rows (VU.fromList [0.0, 1.0]) 0 + cfg = defaultSolverConfig{scL1Lambda = 0.01, scL2Lambda = 0, scMaxIter = 300} + model = fitL1Logistic cfg rows labels (V.fromList ["zero", "signal"]) + ws = VU.toList (lmWeights model) + w0 <- case ws of + (w : _) -> pure w + [] -> assertFailure "expected at least one weight" + assertEqual "zero-variance column has weight zero" 0 w0 + assertBool ("signal weight non-zero (" ++ show (ws !! 1) ++ ")") (ws !! 1 /= 0) + +------------------------------------------------------------------------ +-- A16: Severely imbalanced labels (99:1) — should not collapse to a +-- constant predictor on the majority class without some learning. +------------------------------------------------------------------------ + +testA16ImbalancedLabels :: Test +testA16ImbalancedLabels = TestCase $ do + let nPos = 99 + nNeg = 1 + n = nPos + nNeg + rows = syntheticPoints 16 n 2 + labels = + VU.fromList + (replicate nPos 1.0 ++ replicate nNeg (-1.0)) + cfg = defaultSolverConfig{scL1Lambda = 0.01, scL2Lambda = 0, scMaxIter = 500} + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + ws = VU.toList (lmWeights model) + b = lmIntercept model + anyBad = any (\v -> isNaN v || isInfinite v) (b : ws) + assertBool "no NaN/Inf with 99:1 imbalance" (not anyBad) + assertBool ("intercept favors majority class (got b=" ++ show b ++ ")") (b > 0) + +------------------------------------------------------------------------ +-- A17: Mixed per-feature raw scales — should not diverge. +------------------------------------------------------------------------ + +testA17ImbalancedRawScales :: Test +testA17ImbalancedRawScales = TestCase $ do + let baseRows = syntheticPoints 17 100 3 + rows = + V.map + ( \row -> + let v0 = row VU.! 0 + v1 = row VU.! 1 + v2 = row VU.! 2 + in VU.fromList [1.0e-6 * v0, v1, 1.0e6 * v2] + ) + baseRows + labels = labelsForHyperplane baseRows (VU.fromList [1.0, -0.5, 0.7]) 0 + cfg = defaultSolverConfig{scL1Lambda = 1.0e-4, scL2Lambda = 0, scMaxIter = 500} + model = fitL1Logistic cfg rows labels (V.fromList ["tiny", "unit", "huge"]) + ws = VU.toList (lmWeights model) + b = lmIntercept model + anyBad = any (\v -> isNaN v || isInfinite v) (b : ws) + assertBool ("no NaN/Inf with mixed scales (ws=" ++ show ws ++ ")") (not anyBad) + let preds = [predict model (rows V.! i) | i <- [0 .. V.length rows - 1]] + lbls = [labels VU.! i | i <- [0 .. VU.length labels - 1]] + correct = length [() | (p, l) <- zip preds lbls, p == l] + assertBool + ("non-divergent under wild scales (got " ++ show correct ++ "/100)") + (correct >= 65) + +------------------------------------------------------------------------ +-- A12: maxIter = 0 returns the initial point unchanged +------------------------------------------------------------------------ + +testA12MaxIterZero :: Test +testA12MaxIterZero = TestCase $ do + let rows = syntheticPoints 20 50 2 + labels = labelsForHyperplane rows (VU.fromList [1.0, -0.5]) 0 + cfg = defaultSolverConfig{scMaxIter = 0} + model = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + assertEqual + "maxIter=0 returns zero weights" + (VU.fromList [0, 0]) + (lmWeights model) + assertEqual "maxIter=0 returns zero intercept" 0 (lmIntercept model) + +------------------------------------------------------------------------ +-- A13: maxIter = 1 takes exactly one prox step (results differ from +-- the initial zero point but may not be near the optimum). +------------------------------------------------------------------------ + +testA13MaxIterOne :: Test +testA13MaxIterOne = TestCase $ do + let rows = syntheticPoints 21 80 2 + labels = labelsForHyperplane rows (VU.fromList [1.0, -0.5]) 0 + cfg = defaultSolverConfig{scMaxIter = 1, scL1Lambda = 0.001, scL2Lambda = 0} + cfg0 = cfg{scMaxIter = 0} + m1 = fitL1Logistic cfg rows labels (V.fromList ["x", "y"]) + m0 = fitL1Logistic cfg0 rows labels (V.fromList ["x", "y"]) + anyNonZero v = not (VU.all (== 0) v) + assertEqual "baseline m0 weights are zero" (VU.fromList [0, 0]) (lmWeights m0) + assertBool + ("maxIter=1 must change at least one weight (got " ++ show (lmWeights m1) ++ ")") + (anyNonZero (lmWeights m1) || lmIntercept m1 /= 0) + let badW = VU.any (\x -> isNaN x || isInfinite x) (lmWeights m1) + badB = isNaN (lmIntercept m1) || isInfinite (lmIntercept m1) + assertBool "no NaN/Inf after one iteration" (not (badW || badB)) + +------------------------------------------------------------------------ +-- PR 3: Elastic Net recovery on correlated-feature pairs. Pure L1 picks one +-- of two correlated informative features; Elastic Net keeps both non-zero +-- (Zou & Hastie 2005 grouping effect). Cases: ρ ≈ 0.97 and ρ ≈ 0.7. +------------------------------------------------------------------------ + +-- Generate two correlated features f0, f1 with correlation ρ, plus +-- noise features f2..f7. Truth is sign(f0 + f1). +correlatedPairData :: + Int -> Double -> (V.Vector (VU.Vector Double), VU.Vector Double) +correlatedPairData seed rho = + let n = 400 :: Int + d = 8 :: Int + g0 = mkStdGen seed + drawUnit = randomR (-1.0 :: Double, 1.0) + drawRow !gIn = + let (z0, g1) = drawUnit gIn + (epsRaw, g2) = drawUnit g1 + eps = epsRaw * sqrt (max 0 (1 - rho * rho)) + f0 = z0 + f1 = rho * z0 + eps + drawNoise k g + | k >= d - 2 = ([], g) + | otherwise = + let (x, g') = drawUnit g + (xs, g'') = drawNoise (k + 1) g' + in (x : xs, g'') + (noise, g3) = drawNoise 0 g2 + row = f0 : f1 : noise + in (VU.fromList row, g3) + go 0 _ acc = reverse acc + go k g acc = + let (r, g') = drawRow g + in go (k - 1) g' (r : acc) + rows = V.fromList (go n g0 []) + labels = + VU.generate n $ \i -> + let r = rows V.! i + s = VU.unsafeIndex r 0 + VU.unsafeIndex r 1 + in if s > 0 then 1.0 else -1.0 + in (rows, labels) + +testA19ElasticNetRecoveryHigh :: Test +testA19ElasticNetRecoveryHigh = TestCase $ do + let (rows, labels) = correlatedPairData 31 0.97 + names = V.fromList ["f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7"] + cfgEN = defaultSolverConfig{scL1Lambda = 0.05, scL2Lambda = 0.05, scMaxIter = 1000} + men = fitL1Logistic cfgEN rows labels names + wEN = VU.toList (lmWeights men) + nzCount xs = length (filter (/= 0) xs) + (aEN, bEN) = case wEN of + (a : b : _) -> (a, b) + _ -> error "elastic-net test: expected at least two weights" + assertBool + ("ρ=0.97 EN keeps f0 non-zero; wEN[:2] = " ++ show (take 2 wEN)) + (aEN /= 0) + assertBool + ("ρ=0.97 EN keeps f1 non-zero; wEN[:2] = " ++ show (take 2 wEN)) + (bEN /= 0) + let ratio = abs aEN / max (abs bEN) 1e-9 + assertBool + ("ρ=0.97 EN grouping: |w0/w1| ∈ [0.33, 3.0]; got ratio=" ++ show ratio) + (ratio >= 0.33 && ratio <= 3.0) + assertBool + ("ρ=0.97 EN sparsity: total non-zero ≤ 5; got " ++ show (nzCount wEN)) + (nzCount wEN <= 5) + +testA19ElasticNetRecoveryMid :: Test +testA19ElasticNetRecoveryMid = TestCase $ do + let (rows, labels) = correlatedPairData 37 0.7 + names = V.fromList ["f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7"] + cfgEN = defaultSolverConfig{scL1Lambda = 0.05, scL2Lambda = 0.05, scMaxIter = 1000} + men = fitL1Logistic cfgEN rows labels names + wEN = VU.toList (lmWeights men) + (aEN, bEN) = case wEN of + (a : b : _) -> (a, b) + _ -> error "elastic-net test: expected at least two weights" + assertBool + ("ρ=0.7 EN keeps f0 non-zero; wEN[:2] = " ++ show (take 2 wEN)) + (aEN /= 0) + assertBool + ("ρ=0.7 EN keeps f1 non-zero; wEN[:2] = " ++ show (take 2 wEN)) + (bEN /= 0) + let ratio = abs aEN / max (abs bEN) 1e-9 + assertBool + ("ρ=0.7 EN grouping: |w0/w1| ∈ [0.33, 3.0]; got ratio=" ++ show ratio) + (ratio >= 0.33 && ratio <= 3.0) + +------------------------------------------------------------------------ +-- PR 3: A20 — class-balanced fit on 95/5 imbalance. Unweighted, the intercept +-- polarises toward logit(0.95) ≈ 2.94; with class-balanced weights it sits +-- near 0 and predictions become roughly balanced on a symmetric test set. +------------------------------------------------------------------------ + +testA20ClassBalancedFit :: Test +testA20ClassBalancedFit = TestCase $ do + let n = 200 :: Int + nPos = 190 :: Int + g0 = mkStdGen 41 + drawN = randomR (-1.0 :: Double, 1.0) + drawRowAt mu g = + let (z, g') = drawN g + x = mu + 0.6 * z + in (VU.singleton x, g') + rowsAndLabels = + let go _ 0 _ acc = reverse acc + go !pCnt k g acc = + let !mu = if pCnt > 0 then 0.15 else -0.15 + (row, g') = drawRowAt mu g + !y = if pCnt > 0 then 1.0 else -1.0 + in go (pCnt - 1) (k - 1) g' ((row, y) : acc) + in go nPos n g0 [] + rows = V.fromList (map fst rowsAndLabels) + labels = VU.fromList (map snd rowsAndLabels) + names = V.fromList ["x"] + cfgUnbal = + defaultSolverConfig + { scL1Lambda = 0.001 + , scL2Lambda = 0 + , scMaxIter = 2000 + , scTol = 1e-7 + , scSampleWeights = Nothing + } + nNeg = n - nPos + balanced = + VU.generate n $ \i -> + let !y = VU.unsafeIndex labels i + in if y > 0 + then fromIntegral n / (2 * fromIntegral nPos) + else fromIntegral n / (2 * fromIntegral nNeg) + cfgBal = cfgUnbal{scSampleWeights = Just balanced} + mUnbal = fitL1Logistic cfgUnbal rows labels names + mBal = fitL1Logistic cfgBal rows labels names + bUnbal = lmIntercept mUnbal + bBal = lmIntercept mBal + testRows = + V.fromList + ( replicate 100 (VU.singleton 0.15) + ++ replicate 100 (VU.singleton (-0.15)) + ) + predFracPos m = + let preds = V.map (predict m) testRows + ps = V.length (V.filter (> 0) preds) + in fromIntegral ps / fromIntegral (V.length testRows) :: Double + fracUnbal = predFracPos mUnbal + fracBal = predFracPos mBal + assertBool + ("unbalanced |b| > 2.0; got " ++ show bUnbal) + (abs bUnbal > 2.0) + assertBool + ("balanced |b| < 0.3; got " ++ show bBal) + (abs bBal < 0.3) + assertBool + ("unbalanced fraction-positive on balanced test ≥ 0.90; got " ++ show fracUnbal) + (fracUnbal >= 0.90) + assertBool + ( "balanced fraction-positive on balanced test ∈ [0.40, 0.60]; got " + ++ show fracBal + ) + (fracBal >= 0.40 && fracBal <= 0.60) + +------------------------------------------------------------------------ +-- Test list +------------------------------------------------------------------------ + +tests :: [Test] +tests = + [ TestLabel "A1 recover known hyperplane" testA1RecoverHyperplane + , TestLabel "A2 L1 sparsity" testA2L1Sparsity + , TestLabel "A3 convergence" testA3Convergence + , TestLabel "A4 loss not increasing" testA4LossNotIncreasing + , TestLabel "A5 all same direction" testA5AllSameDirection + , TestLabel "A6 empty input" testA6Empty + , TestLabel "A7 constant feature" testA7ConstantFeature + , TestLabel "A8 large feature values" testA8LargeValues + , TestLabel "A9 standardization round-trip" testA9StandardizationRoundTrip + , TestLabel "A10 determinism" testA10Determinism + , TestLabel "A11 ground truth ratio" testA11GroundTruthRatio + , TestLabel "A12 maxIter zero" testA12MaxIterZero + , TestLabel "A13 maxIter one" testA13MaxIterOne + , TestLabel "A14 constant huge value" testA14ConstantHugeValue + , TestLabel "A15 all-zero feature" testA15AllZeroFeature + , TestLabel "A16 imbalanced 99:1 labels" testA16ImbalancedLabels + , TestLabel "A17 imbalanced raw scales" testA17ImbalancedRawScales + , TestLabel "B1 Expr well-typed" testB1ExprWellTyped + , TestLabel "B2 zero weights pruned" testB2ZeroWeightsPruned + , -- PR 3: Elastic Net + class-balanced weights. + TestLabel "A19 Elastic Net grouping ρ=0.97" testA19ElasticNetRecoveryHigh + , TestLabel "A19 Elastic Net grouping ρ=0.7" testA19ElasticNetRecoveryMid + , TestLabel "A20 class-balanced fit on 95/5" testA20ClassBalancedFit + ] diff --git a/tests-internal/Main.hs b/tests-internal/Main.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Main.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module Main where + +import qualified System.Exit as Exit + +import Test.HUnit +import Test.QuickCheck + +import qualified Cart +import qualified DecisionTree +import qualified Learn.EdgeCases +import qualified Learn.NumericalRigor +import qualified Learn.Numerics +import qualified Learn.Symbolic +import qualified LinearSolver +import qualified Properties.Simplify +import qualified TreePruning +import qualified Worklist + +tests :: Test +tests = + TestList $ + Cart.tests + ++ DecisionTree.tests + ++ LinearSolver.tests + ++ TreePruning.tests + ++ Worklist.tests + ++ Learn.Numerics.tests + ++ Learn.Symbolic.tests + ++ Learn.EdgeCases.tests + ++ Learn.NumericalRigor.tests + +isSuccessful :: Result -> Bool +isSuccessful (Success{}) = True +isSuccessful _ = False + +main :: IO () +main = do + result <- runTestTT tests + if failures result > 0 || errors result > 0 + then Exit.exitFailure + else do + simpRes <- mapM (quickCheckWithResult stdArgs) Properties.Simplify.tests + wlRes <- mapM (quickCheckWithResult stdArgs) Worklist.props + if not (all isSuccessful simpRes) || not (all isSuccessful wlRes) + then Exit.exitFailure + else Exit.exitSuccess diff --git a/tests-internal/Properties/Simplify.hs b/tests-internal/Properties/Simplify.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Properties/Simplify.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | Property tests for the simplifier and tree-pruning pass: @simplify@ +preserves denotation (over Bool and Maybe Bool, incl. NaN\/null\/boundary rows) +and is idempotent; @pruneDead@ preserves the function the tree computes. +-} +module Properties.Simplify (tests) where + +import DataFrame.DecisionTree.Predict (predictWithTree) +import DataFrame.DecisionTree.Prune (pruneDead) +import DataFrame.DecisionTree.Types (Tree (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Column (TypedColumn (TColumn), toVector) +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr, eqExpr) +import DataFrame.Internal.Interpreter (interpret) +import DataFrame.Internal.Simplify (simplify) +import DataFrame.Operators +import qualified DataFrameApi as D + +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import Test.QuickCheck + +-- A fixture spanning the interesting rows: exact thresholds, gaps, NaN, null. +fixtureDF :: D.DataFrame +fixtureDF = + D.fromNamedColumns + [ + ( "x" + , DI.fromList ([10, 20, 25, 30, 35, 40, 50, 0 / 0, -(1 / 0), 1 / 0] :: [Double]) + ) + , ("n", DI.fromList ([10, 20, 25, 30, 35, 40, 50, 0, 100, -5] :: [Int])) + , + ( "m" + , DI.fromList + ( [ Just 10 + , Nothing + , Just 30 + , Just 35 + , Nothing + , Just 50 + , Just 0 + , Just 30 + , Nothing + , Just 40 + ] :: + [Maybe Double] + ) + ) + ] + +thresholds :: [Double] +thresholds = [20, 25, 30, 35, 40] + +-- ---- generators ---- + +-- strict-Bool comparison atoms over the Double column "x" and Int column "n" +genAtomBool :: Gen (Expr Bool) +genAtomBool = do + t <- elements thresholds + oneof + [ elements + [ F.col @Double "x" .< F.lit t + , F.col @Double "x" .<= F.lit t + , F.col @Double "x" .> F.lit t + , F.col @Double "x" .>= F.lit t + , F.col @Double "x" .== F.lit t + , F.col @Double "x" ./= F.lit t + ] + , elements + [ F.toDouble (F.col @Int "n") .< F.lit t + , F.toDouble (F.col @Int "n") .<= F.lit t + , F.toDouble (F.col @Int "n") .> F.lit t + , F.toDouble (F.col @Int "n") .>= F.lit t + ] + ] + +genBoolExpr :: Int -> Gen (Expr Bool) +genBoolExpr d + | d <= 0 = genAtomBool + | otherwise = + oneof + [ genAtomBool + , F.and <$> genBoolExpr (d - 1) <*> genBoolExpr (d - 1) + , F.or <$> genBoolExpr (d - 1) <*> genBoolExpr (d - 1) + , F.not <$> genBoolExpr (d - 1) + ] + +-- nullable comparison atoms over the Maybe Double column "m" +genAtomMaybe :: Gen (Expr (Maybe Bool)) +genAtomMaybe = do + t <- elements thresholds + elements + [ F.col @(Maybe Double) "m" .< F.lit t + , F.col @(Maybe Double) "m" .<= F.lit t + , F.col @(Maybe Double) "m" .> F.lit t + , F.col @(Maybe Double) "m" .>= F.lit t + , F.col @(Maybe Double) "m" .== F.lit t + , F.col @(Maybe Double) "m" ./= F.lit t + ] + +genMaybeExpr :: Int -> Gen (Expr (Maybe Bool)) +genMaybeExpr d + | d <= 0 = genAtomMaybe + | otherwise = + oneof + [ genAtomMaybe + , (.&&) <$> genMaybeExpr (d - 1) <*> genMaybeExpr (d - 1) + , (.||) <$> genMaybeExpr (d - 1) <*> genMaybeExpr (d - 1) + ] + +genTree :: Int -> Gen (Tree T.Text) +genTree d + | d <= 0 = Leaf <$> elements ["A", "B", "C"] + | otherwise = + oneof + [ Leaf <$> elements ["A", "B", "C"] + , do + cond <- genAtomBool + Branch cond <$> genTree (d - 1) <*> genTree (d - 1) + ] + +-- ---- evaluation helpers ---- + +evalBool :: D.DataFrame -> Expr Bool -> Maybe (VU.Vector Bool) +evalBool df e = case interpret @Bool df e of + Right (TColumn tcol) -> either (const Nothing) Just (toVector @Bool @VU.Vector tcol) + Left _ -> Nothing + +evalMaybe :: D.DataFrame -> Expr (Maybe Bool) -> Maybe (V.Vector (Maybe Bool)) +evalMaybe df e = case interpret @(Maybe Bool) df e of + Right (TColumn tcol) -> either (const Nothing) Just (toVector @(Maybe Bool) @V.Vector tcol) + Left _ -> Nothing + +-- ---- properties ---- + +prop_simplifyPreservesBool :: Property +prop_simplifyPreservesBool = + forAll (genBoolExpr 4) $ \e -> + evalBool fixtureDF e === evalBool fixtureDF (simplify e) + +prop_simplifyPreservesMaybe :: Property +prop_simplifyPreservesMaybe = + forAll (genMaybeExpr 3) $ \e -> + evalMaybe fixtureDF e === evalMaybe fixtureDF (simplify e) + +prop_simplifyIdempotent :: Property +prop_simplifyIdempotent = + forAll (genBoolExpr 4) $ \e -> + let s = simplify e in property (eqExpr (simplify s) s) + +prop_pruneDeadPreserves :: Property +prop_pruneDeadPreserves = + forAll (genTree 4) $ \t -> + let n = D.nRows fixtureDF + predAll tr = [predictWithTree @T.Text "x" fixtureDF i tr | i <- [0 .. n - 1]] + in predAll (pruneDead t) === predAll t + +tests :: [Property] +tests = + [ prop_simplifyPreservesBool + , prop_simplifyPreservesMaybe + , prop_simplifyIdempotent + , prop_pruneDeadPreserves + ] diff --git a/tests-internal/TreePruning.hs b/tests-internal/TreePruning.hs new file mode 100644 --- /dev/null +++ b/tests-internal/TreePruning.hs @@ -0,0 +1,115 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | Specification for the fitted-tree pruning pass ('pruneDead'): path-condition +entailment, the false-edge NaN gate, and same-branch collapse. +-} +module TreePruning (tests) where + +import DataFrame.DecisionTree.Prune (pruneDead) +import DataFrame.DecisionTree.Types (Tree (..)) +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression (eqExpr) +import DataFrame.Operators + +import qualified Data.Text as T +import Test.HUnit + +treeEq :: (Eq a) => Tree a -> Tree a -> Bool +treeEq (Leaf x) (Leaf y) = x == y +treeEq (Branch c1 l1 r1) (Branch c2 l2 r2) = eqExpr c1 c2 && treeEq l1 l2 && treeEq r1 r2 +treeEq _ _ = False + +prunesTo :: String -> Tree T.Text -> Tree T.Text -> Test +prunesTo label input want = + TestLabel label . TestCase $ + assertBool + (label ++ ": got " ++ show (pruneDead input) ++ " want " ++ show want) + (treeEq (pruneDead input) want) + +preserved :: String -> Tree T.Text -> Test +preserved label t = prunesTo label t t + +pathEntailment :: [Test] +pathEntailment = + [ prunesTo + "ancestor entails child keeps true subtree" + ( Branch + (F.col @Double "age" .> F.lit (50 :: Double)) + (Branch (F.col @Double "age" .> F.lit (30 :: Double)) (Leaf "a") (Leaf "b")) + (Leaf "c") + ) + (Branch (F.col @Double "age" .> F.lit (50 :: Double)) (Leaf "a") (Leaf "c")) + , prunesTo + "ancestor refutes child keeps false subtree" + ( Branch + (F.col @Double "age" .> F.lit (50 :: Double)) + (Branch (F.col @Double "age" .< F.lit (40 :: Double)) (Leaf "a") (Leaf "b")) + (Leaf "c") + ) + (Branch (F.col @Double "age" .> F.lit (50 :: Double)) (Leaf "b") (Leaf "c")) + ] + +falseEdgeGate :: [Test] +falseEdgeGate = + [ prunesTo + "integral false edge entails child" + ( Branch + (F.toDouble (F.col @Int "ai") .> F.lit (50 :: Double)) + (Leaf "c") + ( Branch + (F.toDouble (F.col @Int "ai") .< F.lit (60 :: Double)) + (Leaf "a") + (Leaf "b") + ) + ) + ( Branch + (F.toDouble (F.col @Int "ai") .> F.lit (50 :: Double)) + (Leaf "c") + (Leaf "a") + ) + ] + +sameBranchCollapse :: [Test] +sameBranchCollapse = + [ prunesTo + "equal leaves collapse the branch" + (Branch (F.col @Double "age" .> F.lit (50 :: Double)) (Leaf "a") (Leaf "a")) + (Leaf "a") + , prunesTo + "collapse cascades upward" + ( Branch + (F.col @Double "age" .> F.lit (50 :: Double)) + (Branch (F.col @Double "hours" .> F.lit (40 :: Double)) (Leaf "a") (Leaf "a")) + (Leaf "a") + ) + (Leaf "a") + ] + +preservedTrees :: [Test] +preservedTrees = + [ preserved + "child not tight enough is kept" + ( Branch + (F.col @Double "age" .> F.lit (50 :: Double)) + (Branch (F.col @Double "age" .> F.lit (60 :: Double)) (Leaf "a") (Leaf "b")) + (Leaf "c") + ) + , preserved + "double false edge is kept (NaN)" + ( Branch + (F.col @Double "weight" .> F.lit (50 :: Double)) + (Leaf "c") + (Branch (F.col @Double "weight" .< F.lit (60 :: Double)) (Leaf "a") (Leaf "b")) + ) + , preserved + "cross-column descendant is kept" + ( Branch + (F.col @Double "age" .> F.lit (50 :: Double)) + (Branch (F.col @Double "income" .> F.lit (30000 :: Double)) (Leaf "a") (Leaf "b")) + (Leaf "c") + ) + ] + +tests :: [Test] +tests = concat [pathEntailment, falseEdgeGate, sameBranchCollapse, preservedTrees] diff --git a/tests-internal/Worklist.hs b/tests-internal/Worklist.hs new file mode 100644 --- /dev/null +++ b/tests-internal/Worklist.hs @@ -0,0 +1,421 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | TDD spec for the saturation worklist 'saturateCandidates' replacing the +lazy generate-all 'boolExprsVec', which is kept as the behaviour-preservation +oracle. +-} +module Worklist (tests, props) where + +import DataFrame.DecisionTree.CondVec ( + CondVec (..), + combineAndVec, + combineOrVec, + materializeCondVec, + ) +import DataFrame.DecisionTree.Pool ( + DedupMode (Structural, TruthVector), + boolExprsVec, + saturateCandidates, + ) +import qualified DataFrame.Functions as F +import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression ( + Expr, + compareExpr, + eSize, + eqExpr, + normalize, + ) +import DataFrame.Operators +import qualified DataFrameApi as D + +import Data.Function (on) +import Data.List (minimumBy, nubBy) +import qualified Data.Maybe +import qualified Data.Set as Set +import qualified Data.Vector.Unboxed as VU +import Test.HUnit +import Test.QuickCheck + +-- Fixture: x = 0..5, y = 5..0 (anti-correlated), z scrambled (independent of both). +-- Note x>2 and y<3 share the truth vector [F,F,F,T,T,T], so the truth-vector mode +-- must collapse them; z gives a third column for non-consolidating cross-column combos. +fixtureDF :: D.DataFrame +fixtureDF = + D.fromNamedColumns + [ ("x", DI.fromList ([0, 1, 2, 3, 4, 5] :: [Double])) + , ("y", DI.fromList ([5, 4, 3, 2, 1, 0] :: [Double])) + , ("z", DI.fromList ([2, 5, 1, 4, 0, 3] :: [Double])) + ] + +mat :: Expr Bool -> CondVec +mat e = + Data.Maybe.fromMaybe + (error "Worklist.mat: could not materialize") + (materializeCondVec fixtureDF e) + +xGt, xLt, yGt, yLt, zGt, zLt :: Double -> CondVec +xGt n = mat (F.col @Double "x" .>. F.lit n) +xLt n = mat (F.col @Double "x" .<. F.lit n) +yGt n = mat (F.col @Double "y" .>. F.lit n) +yLt n = mat (F.col @Double "y" .<. F.lit n) +zGt n = mat (F.col @Double "z" .>. F.lit n) +zLt n = mat (F.col @Double "z" .<. F.lit n) + +-- Same truth vector as 'xGt 2' ([F,F,F,T,T,T]) but eSize 4 vs 3 — a non-degenerate +-- truth-vector collision for the min-eSize representative rule. +notLe2 :: CondVec +notLe2 = mat (F.not (F.col @Double "x" .<=. F.lit 2)) + +litTrue :: CondVec +litTrue = mat (F.lit True) + +keyOf :: CondVec -> String +keyOf = show . normalize . cvExpr + +keySet :: [CondVec] -> Set.Set String +keySet = Set.fromList . map keyOf + +truthSet :: [CondVec] -> Set.Set [Bool] +truthSet = Set.fromList . map (VU.toList . cvVec) + +-- Mirrors 'evalWithPenaltyVec' (DecisionTree.hs): score = (#care-point errors, eSize), +-- depending only on the cached vector + size, so distinct same-vector same-size atoms tie. +penBy :: [Bool] -> CondVec -> (Int, Int) +penBy lbls cv = + ( length (filter id (zipWith (/=) lbls (VU.toList (cvVec cv)))) + , eSize (cvExpr cv) + ) + +-- The candidate 'bestDiscreteCandidate' would select: the first 'minimumBy penalty' winner. +argminKey :: [Bool] -> [CondVec] -> String +argminKey lbls = keyOf . minimumBy (compare `on` penBy lbls) + +-- Oracle: the current depth-bounded generate-all. +ref :: Int -> [CondVec] -> [CondVec] +ref d base = boolExprsVec base base 0 d + +base3 :: [CondVec] +base3 = [xGt 2, xGt 4, yGt 2] + +-- x>2 and y<3 share the truth vector [F,F,F,T,T,T], so truth-vector mode collapses +-- this 3-atom base to 2 distinct vectors while structural mode keeps all three. +collBase :: [CondVec] +collBase = [xGt 2, yLt 3, yGt 2] + +-- Wider fixture (3 independent-ish columns, 10 rows) yielding many distinct truth +-- vectors — broader coverage for the truth-vector floor / dedup than the 6-row x/y fixture. +wideDF :: D.DataFrame +wideDF = + D.fromNamedColumns + [ ("a", DI.fromList ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9] :: [Double])) + , ("b", DI.fromList ([9, 7, 5, 3, 1, 8, 6, 4, 2, 0] :: [Double])) + , ("c", DI.fromList ([1, 1, 2, 2, 3, 3, 4, 4, 5, 5] :: [Double])) + ] + +matW :: Expr Bool -> CondVec +matW e = + Data.Maybe.fromMaybe + (error "Worklist.matW: could not materialize") + (materializeCondVec wideDF e) + +wideBase :: [CondVec] +wideBase = + [ matW (F.col @Double "a" .>. F.lit 3) + , matW (F.col @Double "b" .<. F.lit 5) + , matW (F.col @Double "c" .>=. F.lit 3) + ] + +------------------------------------------------------------------------ +-- HUnit cases +------------------------------------------------------------------------ + +tests :: [Test] +tests = + [ TestLabel "structural: same distinct set as oracle" . TestCase $ + assertEqual + "keySet" + (keySet (ref 2 base3)) + (keySet (saturateCandidates Structural 2 base3)) + , TestLabel "structural: output deduped (length == distinct keys)" . TestCase $ + let out = saturateCandidates Structural 2 base3 + in assertEqual "no eqExpr duplicates" (Set.size (keySet out)) (length out) + , TestLabel "structural: base atoms all present" . TestCase $ + assertBool "base subset of output" $ + keySet base3 `Set.isSubsetOf` keySet (saturateCandidates Structural 1 base3) + , TestLabel "structural: deterministic" . TestCase $ + assertEqual + "two runs identical" + (map keyOf (saturateCandidates Structural 2 base3)) + (map keyOf (saturateCandidates Structural 2 base3)) + , TestLabel "structural: consolidation flows through the worklist" . TestCase $ + assertEqual + "consolidated set" + (keySet [xGt 2, xGt 4]) + (keySet (saturateCandidates Structural 2 [xGt 2, xGt 4])) + , TestLabel "truth-vector: all output truth vectors distinct" . TestCase $ + let out = saturateCandidates TruthVector 2 [xGt 2, yLt 3] + in assertEqual "distinct cvVecs" (Set.size (truthSet out)) (length out) + , TestLabel "truth-vector: collapses same-truth atoms (x>2, y<3)" . TestCase $ + assertEqual + "collapsed to one" + 1 + (length (saturateCandidates TruthVector 1 [xGt 2, yLt 3])) + , TestLabel "truth-vector: reaches the semantic floor (no split dropped)" + . TestCase + $ assertEqual + "same distinct truth vectors as oracle" + (truthSet (ref 2 base3)) + (truthSet (saturateCandidates TruthVector 2 base3)) + , TestLabel "truth-vector: strictly fewer candidates when a collision exists" + . TestCase + $ assertBool "|truth| < |structural|" + $ length (saturateCandidates TruthVector 2 collBase) + < length (saturateCandidates Structural 2 collBase) + , TestLabel "truth-vector: keeps the minimum-eSize representative" . TestCase $ + let out = saturateCandidates TruthVector 1 [xGt 2, notLe2] + in do + assertEqual "min eSize survivor" [3] (map (eSize . cvExpr) out) + assertEqual "survivor is x>2" [keyOf (xGt 2)] (map keyOf out) + , TestLabel "truth-vector: tie-break independent of input order" . TestCase $ + assertEqual + "order-independent survivor" + (map keyOf (saturateCandidates TruthVector 1 [xGt 2, yLt 3])) + (map keyOf (saturateCandidates TruthVector 1 [yLt 3, xGt 2])) + , TestLabel "edge: empty base" . TestCase $ + assertEqual + "empty in, empty out" + Set.empty + (keySet (saturateCandidates Structural 2 [])) + , TestLabel "edge: singleton base" . TestCase $ + assertEqual + "same as oracle" + (keySet (ref 2 [xGt 2])) + (keySet (saturateCandidates Structural 2 [xGt 2])) + , TestLabel "edge: maxDepth 0 is the base only" . TestCase $ + assertEqual + "no expansion" + (keySet base3) + (keySet (saturateCandidates Structural 0 base3)) + , TestLabel "edge: duplicate-seeded base" . TestCase $ + let base = [xGt 2, xGt 2, yGt 2] + in assertEqual + "dedups seed, same as oracle" + (keySet (ref 2 base)) + (keySet (saturateCandidates Structural 2 base)) + , TestLabel "edge: literal operand" . TestCase $ + let base = [xGt 2, litTrue] + in assertEqual + "same as oracle" + (keySet (ref 2 base)) + (keySet (saturateCandidates Structural 2 base)) + , TestLabel "law: cvVec is a homomorphism over AND" . TestCase $ + assertEqual + "cvVec(a∧b) == cvVec a && cvVec b" + (VU.toList (VU.zipWith (&&) (cvVec (xGt 2)) (cvVec (yGt 2)))) + (VU.toList (cvVec (combineAndVec (xGt 2) (yGt 2)))) + , TestLabel "law: cvVec is a homomorphism over OR" . TestCase $ + assertEqual + "cvVec(a∨b) == cvVec a || cvVec b" + (VU.toList (VU.zipWith (||) (cvVec (xGt 2)) (cvVec (yGt 2)))) + (VU.toList (cvVec (combineOrVec (xGt 2) (yGt 2)))) + , TestLabel "law: consolidated expr re-interprets to its cached vector" . TestCase $ + let c = combineAndVec (xGt 2) (xGt 4) + in assertEqual + "cached == re-interpreted" + (VU.toList (cvVec c)) + (VU.toList (cvVec (mat (cvExpr c)))) + , TestLabel "structural: output order matches the deduped oracle" . TestCase $ + assertEqual + "deduped-oracle order" + (map keyOf (nubBy ((==) `on` keyOf) (ref 2 base3))) + (map keyOf (saturateCandidates Structural 2 base3)) + , TestLabel + "structural: matches oracle set+order at depth 3+4 (non-consolidating)" + . TestCase + $ let b = [xGt 2, yGt 2, zGt 1] + deduped d = map keyOf (nubBy ((==) `on` keyOf) (ref d b)) + out d = map keyOf (saturateCandidates Structural d b) + in do + assertEqual "set d3" (Set.fromList (deduped 3)) (Set.fromList (out 3)) + assertEqual "order d3" (deduped 3) (out 3) + assertEqual "set d4" (Set.fromList (deduped 4)) (Set.fromList (out 4)) + assertEqual "order d4" (deduped 4) (out 4) + , TestLabel "structural: stabilizes at fixpoint (depth cap is a no-op past it)" + . TestCase + $ assertEqual + "depth 2 == depth 5" + (keySet (saturateCandidates Structural 2 [xGt 1, xGt 2, xGt 3, xGt 4])) + (keySet (saturateCandidates Structural 5 [xGt 1, xGt 2, xGt 3, xGt 4])) + , TestLabel "truth-vector: reaches the floor on a wider fixture" . TestCase $ + assertEqual + "same distinct truth vectors as oracle" + (truthSet (ref 2 wideBase)) + (truthSet (saturateCandidates TruthVector 2 wideBase)) + , TestLabel "selection: surfaces the oracle's winning combination" . TestCase $ + let lbls = [False, False, False, True, True, False] + base = [xGt 2, xLt 5, yGt 2] + in assertEqual + "same argmin as oracle" + (argminKey lbls (ref 2 base)) + (argminKey lbls (saturateCandidates Structural 2 base)) + , TestLabel "selection: tie-winner tracks input order, matching the oracle" + . TestCase + $ let lbls = [False, False, False, True, True, True] + in do + assertEqual + "x>2-first order" + (argminKey lbls (ref 2 [xGt 2, yLt 3])) + (argminKey lbls (saturateCandidates Structural 2 [xGt 2, yLt 3])) + assertEqual + "y<3-first order" + (argminKey lbls (ref 2 [yLt 3, xGt 2])) + (argminKey lbls (saturateCandidates Structural 2 [yLt 3, xGt 2])) + , TestLabel + "bounded: output is the distinct closure, below the oracle's materialized count" + . TestCase + $ let base = [xGt 1, xGt 2, xGt 3, xGt 4] + gen = ref 3 base + in do + assertEqual + "output bounded to the distinct closure" + (Set.size (keySet gen)) + (length (saturateCandidates Structural 3 base)) + assertBool + "oracle materializes more than the closure (the explosion the worklist avoids)" + (Set.size (keySet gen) < length gen) + , TestLabel "structural: maxDepth 1 is base-only (no combination round)" + . TestCase + $ assertEqual + "no combination at depth 1" + (keySet (ref 1 [xGt 2, yGt 2])) + (keySet (saturateCandidates Structural 1 [xGt 2, yGt 2])) + , TestLabel "structural: base atoms survive the combination round" . TestCase $ + assertBool "base subset of output at depth 2" $ + keySet base3 `Set.isSubsetOf` keySet (saturateCandidates Structural 2 base3) + , TestLabel + "structural: re-saturating a closed base is stable (fixpoint idempotence)" + . TestCase + $ let b = [xGt 1, xGt 2, xGt 3, xGt 4] + in assertEqual + "saturate ∘ saturate == saturate" + (keySet (saturateCandidates Structural 2 b)) + (keySet (saturateCandidates Structural 2 (saturateCandidates Structural 2 b))) + , TestLabel "law: combiner key is order-independent (congruence basis)" . TestCase $ + do + assertEqual + "AND consolidation commutes at the key" + (keyOf (combineAndVec (xGt 2) (xGt 4))) + (keyOf (combineAndVec (xGt 4) (xGt 2))) + assertEqual + "OR consolidation commutes at the key" + (keyOf (combineOrVec (xGt 2) (xGt 4))) + (keyOf (combineOrVec (xGt 4) (xGt 2))) + assertEqual + "cross-column AND commutes at the key" + (keyOf (combineAndVec (xGt 2) (yGt 2))) + (keyOf (combineAndVec (yGt 2) (xGt 2))) + , TestLabel + "truth-vector: section is the (eSize, compareExpr)-minimum of the fiber" + . TestCase + $ let fiber = [xGt 2, yLt 3] + cmp a b = + compare (eSize (cvExpr a)) (eSize (cvExpr b)) + <> compareExpr (cvExpr a) (cvExpr b) + want = keyOf (minimumBy cmp fiber) + in assertEqual + "min-section survivor" + [want] + (map keyOf (saturateCandidates TruthVector 1 fiber)) + ] + +------------------------------------------------------------------------ +-- QuickCheck properties (over generated base pools and depths) +------------------------------------------------------------------------ + +genAtom :: Gen CondVec +genAtom = + elements + [xGt 1, xGt 2, xGt 3, xLt 2, xLt 4, yGt 1, yGt 3, yLt 3, zGt 1, zGt 3, zLt 4] + +genBase :: Gen [CondVec] +genBase = choose (2, 5) >>= \k -> vectorOf k genAtom + +-- Random label vector of the fixture's length (6 rows), for selection-preservation. +genLabels :: Gen [Bool] +genLabels = vectorOf 6 (elements [False, True]) + +prop_structuralSameSet :: Property +prop_structuralSameSet = + forAllBlind genBase $ \base -> + forAll (choose (1, 3)) $ \d -> + counterexample (show (map keyOf base, d)) $ + keySet (saturateCandidates Structural d base) === keySet (ref d base) + +prop_truthVectorFloor :: Property +prop_truthVectorFloor = + forAllBlind genBase $ \base -> + forAll (choose (1, 3)) $ \d -> + counterexample (show (map keyOf base, d)) $ + truthSet (saturateCandidates TruthVector d base) === truthSet (ref d base) + +-- The candidate set depends only on the base as a set, not its input order. +-- (Output order tracks input order — the byte-identity contract; see selection tests.) +prop_orderInvariant :: Property +prop_orderInvariant = + forAllBlind genBase $ \base -> + forAllBlind (shuffle base) $ \base' -> + forAll (choose (1, 3)) $ \d -> + counterexample (show (map keyOf base, map keyOf base', d)) $ + keySet (saturateCandidates Structural d base) + === keySet (saturateCandidates Structural d base') + +-- The candidate the consumer's 'minimumBy penaltyCV' selects is byte-identical to the oracle's, +-- for any label vector (d >= 2 so combinations exist). This is the model-preservation contract. +prop_selectionPreserved :: Property +prop_selectionPreserved = + forAllBlind genBase $ \base -> + forAllBlind genLabels $ \lbls -> + forAll (choose (2, 3)) $ \d -> + counterexample (show (map keyOf base, lbls, d)) $ + argminKey lbls (saturateCandidates Structural d base) + === argminKey lbls (ref d base) + +-- The full output (order included) is byte-identical to the deduped oracle, at every depth. +-- Subsumes selection-preservation for ANY (cvVec,eSize)-penalty, and stresses the +-- frontier:=admitted optimisation past the depth where it could first diverge. +prop_orderMatchesOracle :: Property +prop_orderMatchesOracle = + forAllBlind genBase $ \base -> + forAll (choose (2, 3)) $ \d -> + counterexample (show (map keyOf base, d)) $ + map keyOf (saturateCandidates Structural d base) + === map keyOf (nubBy ((==) `on` keyOf) (ref d base)) + +-- The structural key faithfully represents the 'eqExpr' quotient on the candidate domain +-- (atoms and their AND/OR products): show.normalize merges exactly what eqExpr merges. +genCand :: Gen CondVec +genCand = + oneof + [ genAtom + , combineAndVec <$> genAtom <*> genAtom + , combineOrVec <$> genAtom <*> genAtom + ] + +prop_keyFaithful :: Property +prop_keyFaithful = + forAllBlind genCand $ \a -> + forAllBlind genCand $ \b -> + counterexample (keyOf a ++ " vs " ++ keyOf b) $ + (keyOf a == keyOf b) === eqExpr (cvExpr a) (cvExpr b) + +props :: [Property] +props = + [ prop_structuralSameSet + , prop_truthVectorFloor + , prop_orderInvariant + , prop_selectionPreserved + , prop_orderMatchesOracle + , prop_keyFaithful + ]