diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for symbolic-regression + +## 0.1.0.0 + +* Basic integration with srtree diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Michael Chavinda + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,104 @@ +# symbolic-regression + +A Haskell library (based on [eggp](https://github.com/folivetti/eggp) which is in turn based on [srtree](https://github.com/folivetti/srtree)) for symbolic regression on DataFrames. Automatically discover mathematical expressions that best fit your data using genetic programming with e-graph optimization. + +## Overview + +symbolic-regression integrates symbolic regression capabilities into a DataFrame workflow. Given a target column and a dataset, it evolves mathematical expressions that predict the target variable, returning a Pareto front of expressions trading off complexity and accuracy. + +## Quick Start + +```haskell +ghci> import qualified DataFrame as D +ghci> import DataFrame.Functions ((.=)) +ghci> import Symbolic.Regression + +-- Load your data +ghci> df <- D.readParquet "./data/mtcars.parquet" + +-- Run symbolic regression to predict 'mpg' +-- NOTE: ALL COLUMNS MUST BE CONVERTED TO DOUBLE FIRST +-- e.g df' = D.derive "some_column" (F.toDouble (F.col @Int "some_column")) df +-- Symbolic regression will by default only use the double columns +-- otherwise. +ghci> exprs <- fit defaultRegressionConfig mpg df + +-- View discovered expressions (Pareto front from simplest to most complex) +ghci> map D.prettyPrint exprs +-- [ qsec, +-- , 57.33 / wt +-- , 10.75 + (1557.67 / disp)] + +-- Create named expressions that we'll use in a dataframe. +ghci> levels = zipWith (.=) ["level_1", "level_2", "level_3"] exprs + +-- Show the various predictions in our dataframe. +ghci> D.deriveMany levels df + +-- Or pick the best one +ghci> D.derive "prediction" (last exprs) df +``` + +## Configuration + +Customize the search with `RegressionConfig`: + +```haskell +data RegressionConfig = RegressionConfig + { generations :: Int -- Number of evolutionary generations (default: 100) + , maxExpressionSize :: Int -- Maximum tree depth/complexity (default: 5) + , numFolds :: Int -- Cross-validation folds (default: 3) + , showTrace :: Bool -- Print progress during evolution (default: True) + , lossFunction :: Distribution -- MSE, Gaussian, Poisson, etc. (default: MSE) + , numOptimisationIterations :: Int -- Parameter optimization iterations (default: 30) + , numParameterRetries :: Int -- Retries for parameter fitting (default: 2) + , populationSize :: Int -- Population size (default: 100) + , tournamentSize :: Int -- Tournament selection size (default: 3) + , crossoverProbability :: Double -- Crossover rate (default: 0.95) + , mutationProbability :: Double -- Mutation rate (default: 0.3) + , unaryFunctions :: [...] -- Unary operations to include + , binaryFunctions :: [...] -- Binary operations to include + , numParams :: Int -- Number of parameters (-1 for auto) + , generational :: Bool -- Use generational replacement (default: False) + , simplifyExpressions :: Bool -- Simplify output expressions (default: True) + , maxTime :: Int -- Time limit in seconds (-1 for none) + , dumpTo :: String -- Save e-graph state to file + , loadFrom :: String -- Load e-graph state from file + } +``` + +### Example: Custom Configuration + +```haskell +myConfig :: RegressionConfig +myConfig = defaultRegressionConfig + { generations = 200 + , maxExpressionSize = 7 + , populationSize = 200 + } + +exprs <- fit myConfig targetColumn df +``` + +## Output + +`fit` returns a list of expressions representing the Pareto front, ordered by complexity (simplest first). Each expression: + +- Is a valid `Expr Double` that can be used with DataFrame operations +- Represents a different trade-off between simplicity and accuracy +- Has optimized numerical constants + +## How It Works + +1. **Genetic Programming**: Evolves a population of expression trees through selection, crossover, and mutation +2. **E-graph Optimization**: Uses equality saturation to discover equivalent expressions and simplify +3. **Parameter Optimization**: Fits numerical constants using nonlinear optimization +4. **Pareto Selection**: Returns expressions across the complexity-accuracy frontier + +## Dependencies + +### System dependencies +To install symbolic-regression you'll need: +* libz: `sudo apt install libz-dev` +* libnlopt: `sudo apt install libnlopt-dev` +* libgmp: `sudo apt install libgmp-dev` diff --git a/src/Symbolic/Regression.hs b/src/Symbolic/Regression.hs new file mode 100644 --- /dev/null +++ b/src/Symbolic/Regression.hs @@ -0,0 +1,704 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : Symbolic.Regression +Description : Symbolic regression for DataFrames using genetic programming with e-graph optimization + +This module provides symbolic regression capabilities for DataFrame workflows. +Given a target column and a dataset, it evolves mathematical expressions that +predict the target variable, returning a Pareto front of expressions trading +off complexity and accuracy. + += Quick Start + +@ +import qualified DataFrame as D +import DataFrame.Functions ((.=)) +import Symbolic.Regression + +-- Load your data +df <- D.readParquet "./data/mtcars.parquet" + +-- Run symbolic regression to predict 'mpg' +exprs <- fit defaultRegressionConfig mpg df + +-- Use the best expression +D.derive "prediction" (last exprs) df +@ + += Important Notes + +All columns used in regression must be converted to 'Double' first. +Symbolic regression will by default only use the double columns. + += How It Works + +1. __Genetic Programming__: Evolves a population of expression trees through + selection, crossover, and mutation +2. __E-graph Optimization__: Uses equality saturation to discover equivalent + expressions and simplify +3. __Parameter Optimization__: Fits numerical constants using nonlinear optimization +4. __Pareto Selection__: Returns expressions across the complexity-accuracy frontier +-} +module Symbolic.Regression ( + -- * Main API + fit, + + -- * Configuration + RegressionConfig (..), + defaultRegressionConfig, +) where + +import Control.Exception (throw) +import Control.Monad.State.Strict +import Data.Massiv.Array as MA hiding (forM, forM_) +import qualified Data.Text as T +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified DataFrame as D +import qualified DataFrame.Functions as F +import DataFrame.Internal.Expression +import System.Random + +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Simplify hiding (myCost) +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.ModelSelection (fractionalBayesFactor) +import Control.Lens (over) +import Control.Monad ( + filterM, + forM, + forM_, + replicateM, + unless, + when, + (>=>), + ) +import Data.Binary (decode, encode) +import qualified Data.ByteString.Lazy as BS +import Data.Function (on) +import Data.Functor +import qualified Data.HashSet as Set +import qualified Data.IntMap.Strict as IM +import Data.List ( + intercalate, + maximumBy, + nub, + zip4, + ) +import Data.List.Split (splitOn) +import qualified Data.Map.Strict as Map +import Data.Maybe (fromJust, fromMaybe) +import Data.SRTree +import Data.SRTree.Datasets +import qualified Data.SRTree.Internal as SI +import Data.SRTree.Print +import Data.SRTree.Random + +import Algorithm.EqSat (runEqSat) +import Algorithm.EqSat.SearchSR +import Data.Time.Clock.POSIX +import Text.ParseSR + +{- | Configuration for the symbolic regression algorithm. + +Use 'defaultRegressionConfig' as a starting point and modify fields as needed: + +@ +myConfig :: RegressionConfig +myConfig = defaultRegressionConfig + { generations = 200 + , maxExpressionSize = 7 + , populationSize = 200 + } +@ +-} +data RegressionConfig = RegressionConfig + { generations :: Int + -- ^ Number of evolutionary generations to run (default: 100) + , maxExpressionSize :: Int + -- ^ Maximum tree depth\/complexity for generated expressions (default: 5) + , numFolds :: Int + -- ^ Number of cross-validation folds (default: 3) + , showTrace :: Bool + -- ^ Whether to print progress during evolution (default: 'True') + , lossFunction :: Distribution + -- ^ Loss function to optimize: 'MSE', 'Gaussian', 'Poisson', etc. (default: 'MSE') + , numOptimisationIterations :: Int + -- ^ Number of iterations for parameter optimization (default: 30) + , numParameterRetries :: Int + -- ^ Number of retries for parameter fitting (default: 2) + , populationSize :: Int + -- ^ Size of the expression population (default: 100) + , tournamentSize :: Int + -- ^ Number of individuals in tournament selection (default: 3) + , crossoverProbability :: Double + -- ^ Probability of crossover between expressions (default: 0.95) + , mutationProbability :: Double + -- ^ Probability of mutation (default: 0.3) + , unaryFunctions :: [D.Expr Double -> D.Expr Double] + -- ^ Unary operations to include in the search space (default: @[]@) + , binaryFunctions :: [D.Expr Double -> D.Expr Double -> D.Expr Double] + {- ^ Binary operations to include in the search space + (default: @[(+), (-), (*), (\/)]@) + -} + , numParams :: Int + -- ^ Number of parameters to use. Set to @-1@ for automatic detection (default: -1) + , generational :: Bool + -- ^ Whether to use generational replacement strategy (default: 'False') + , simplifyExpressions :: Bool + -- ^ Whether to simplify output expressions using e-graph optimization (default: 'True') + , maxTime :: Int + -- ^ Time limit in seconds. Set to @-1@ for no limit (default: -1) + , dumpTo :: String + -- ^ File path to save e-graph state for later resumption (default: @\"\"@) + , loadFrom :: String + -- ^ File path to load e-graph state from a previous run (default: @\"\"@) + } + +{- | Default configuration for symbolic regression. + +Provides sensible defaults for most use cases: + +* 100 generations with population size 100 +* Maximum expression size of 5 +* 3-fold cross-validation +* MSE loss function +* Basic arithmetic operations: @+@, @-@, @*@, @\/@ + +Modify specific fields to customize the search behavior. +-} +defaultRegressionConfig :: RegressionConfig +defaultRegressionConfig = + RegressionConfig + { generations = 100 + , maxExpressionSize = 5 + , numFolds = 3 + , showTrace = True + , lossFunction = MSE + , numOptimisationIterations = 30 + , numParameterRetries = 2 + , populationSize = 100 + , tournamentSize = 3 + , crossoverProbability = 0.95 + , mutationProbability = 0.3 + , unaryFunctions = [] + , binaryFunctions = [(+), (-), (*), (/)] + , numParams = -1 + , generational = False + , simplifyExpressions = True + , maxTime = -1 + , dumpTo = "" + , loadFrom = "" + } + +{- | Run symbolic regression to discover mathematical expressions that fit the data. + +Returns a list of expressions representing the Pareto front, ordered by +complexity (simplest first). Each expression: + +* Is a valid @'D.Expr' 'Double'@ that can be used with DataFrame operations +* Represents a different trade-off between simplicity and accuracy +* Has optimized numerical constants + += Example + +@ +exprs <- fit defaultRegressionConfig targetColumn df + +-- View discovered expressions +map D.prettyPrint exprs +-- [\"qsec\", \"57.33 \/ wt\", \"10.75 + (1557.67 \/ disp)\"] + +-- Use expressions in DataFrame operations +D.derive \"prediction\" (last exprs) df +@ + += Important + +All columns must be converted to 'Double' before running regression. +The algorithm will only use double-typed columns as features. +-} +fit :: + -- | Configuration controlling the search algorithm + RegressionConfig -> + -- | Target column expression to predict + D.Expr Double -> + -- | Input DataFrame containing features and target + D.DataFrame -> + -- | Pareto front of expressions, ordered simplest to most complex + IO [D.Expr Double] +fit cfg targetColumn df = do + g <- getStdGen + let + df' = + D.exclude + [F.name targetColumn] + (D.selectBy [D.byProperty (D.hasElemType @Double)] df) + matrix = either throw id (D.toDoubleMatrix df') + features = fromLists' Seq (V.toList (V.map VU.toList matrix)) :: Array S Ix2 Double + target' = fromLists' Seq (D.columnAsList targetColumn df) :: Array S Ix1 Double + nonterminals = + intercalate + "," + ( Prelude.map + (toNonTerminal . (\f -> f (F.col "fake1") (F.col "fake2"))) + (binaryFunctions cfg) + ) + varnames = + intercalate + "," + ( Prelude.map + T.unpack + (Prelude.filter (/= F.name targetColumn) (D.columnNames df)) + ) + alg = + evalStateT + ( egraphGP + cfg + nonterminals + varnames + [((features, target', Nothing), (features, target', Nothing))] + [(features, target', Nothing)] + ) + emptyGraph + fmap (Prelude.map (toExpr df')) (evalStateT alg g) + +toExpr :: D.DataFrame -> Fix SRTree -> Expr Double +toExpr _ (Fix (Const value)) = Lit value +toExpr df (Fix (Var ix)) = Col (D.columnNames df !! ix) +toExpr df (Fix (Bin op left right)) = case op of + SI.Add -> toExpr df left + toExpr df right + SI.Sub -> toExpr df left - toExpr df right + SI.Mul -> toExpr df left * toExpr df right + SI.Div -> toExpr df left / toExpr df right + treeOp -> error ("UNIMPLEMENTED OPERATION: " ++ show treeOp) +toExpr _ _ = error "UNIMPLEMENTED" + +toNonTerminal :: D.Expr Double -> String +toNonTerminal (BinaryOp "add" _ _ _) = "add" +toNonTerminal (BinaryOp "sub" _ _ _) = "sub" +toNonTerminal (BinaryOp "mult" _ _ _) = "mul" +toNonTerminal (BinaryOp "divide" _ _ _) = "div" +toNonTerminal e = error ("Unsupported operation: " ++ show e) + +egraphGP :: + RegressionConfig -> + String -> -- nonterminals + String -> -- varnames + [(DataSet, DataSet)] -> + [DataSet] -> + StateT EGraph (StateT StdGen IO) [Fix SRTree] +egraphGP cfg nonterminals varnames dataTrainVals dataTests = do + unless (null (loadFrom cfg)) $ + io (BS.readFile (loadFrom cfg)) >>= \eg -> put (decode eg) + + _ <- insertTerms + evaluateUnevaluated fitFun + + t0 <- io getPOSIXTime + + pop <- replicateM (populationSize cfg) $ do + ec <- insertRndExpr (maxExpressionSize cfg) rndTerm rndNonTerm >>= canonical + _ <- updateIfNothing fitFun ec + pure ec + pop' <- Prelude.mapM canonical pop + + output <- + if showTrace cfg + then forM (Prelude.zip [0 ..] pop') $ uncurry printExpr' + else pure [] + + let mTime = + if maxTime cfg < 0 then Nothing else Just (fromIntegral $ maxTime cfg - 5) + (_, _, _) <- iterateFor (generations cfg) t0 mTime (pop', output, populationSize cfg) $ \_ (ps', out, curIx) -> do + newPop' <- replicateM (populationSize cfg) (evolve ps') + + out' <- + if showTrace cfg + then forM (Prelude.zip [curIx ..] newPop') $ uncurry printExpr' + else pure [] + + totSz <- gets (Map.size . _eNodeToEClass) + let full = totSz > max maxMem (populationSize cfg) + when full (cleanEGraph >> cleanDB) + + newPop <- + if generational cfg + then Prelude.mapM canonical newPop' + else do + pareto <- + concat <$> forM [1 .. maxExpressionSize cfg] (`getTopFitEClassWithSize` 2) + let remainder = populationSize cfg - length pareto + lft <- + if full + then getTopFitEClassThat remainder (const True) + else pure $ Prelude.take remainder newPop' + Prelude.mapM canonical (pareto <> lft) + pure (newPop, out <> out', curIx + populationSize cfg) + + unless (null (dumpTo cfg)) $ + get >>= (io . BS.writeFile (dumpTo cfg) . encode) + paretoFront' fitFun (maxExpressionSize cfg) + where + maxMem = 2000000 + fitFun = + fitnessMV + shouldReparam + (numParameterRetries cfg) + (numOptimisationIterations cfg) + (lossFunction cfg) + dataTrainVals + nonTerms = parseNonTerms nonterminals + (Sz2 _ nFeats) = case dataTrainVals of + [] -> Sz2 0 0 + (h : _) -> MA.size (getX . fst $ h) + params = + if numParams cfg == -1 + then [param 0] + else Prelude.map param [0 .. numParams cfg - 1] + shouldReparam = numParams cfg == -1 + relabel = if shouldReparam then relabelParams else relabelParamsOrder + terms = + if lossFunction cfg == ROXY + then var 0 : params + else [var ix | ix <- [0 .. nFeats - 1]] + uniNonTerms = [t | t <- nonTerms, isUni t] + binNonTerms = [t | t <- nonTerms, isBin t] + + isUni (Uni _ _) = True + isUni _ = False + + isBin (Bin{}) = True + isBin _ = False + + cleanEGraph = do + let nParetos = 10 + io . putStrLn $ "cleaning" + pareto <- + forM [1 .. maxExpressionSize cfg] (`getTopFitEClassWithSize` nParetos) + >>= Prelude.mapM canonical . concat + infos <- forM pareto (\c -> gets (fmap _info . (IM.!? c) . _eClass)) + exprs <- forM pareto getBestExpr + put emptyGraph + newIds <- fromTrees myCost $ Prelude.map relabel exprs + forM_ (Prelude.zip newIds (Prelude.reverse infos)) $ \(eId, info') -> + case info' of + Nothing -> pure () + Just i'' -> insertFitness eId (fromJust $ _fitness i'') (_theta i'') + + rndTerm = do + coin <- toss + if coin || numParams cfg == 0 then randomFrom terms else randomFrom params + + rndNonTerm = randomFrom nonTerms + + refitChanged = do + ids <- + (gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList) + Data.Functor.<&> nub + modify' $ over (eDB . refits) (const Set.empty) + forM_ ids $ \ec -> do + t <- getBestExpr ec + (f, p) <- fitFun t + insertFitness ec f p + + iterateFor 0 _ _ xs _ = pure xs + iterateFor n t0' maxT xs f = do + xs' <- f n xs + t1 <- io getPOSIXTime + let delta = t1 - t0' + maxT' = subtract delta <$> maxT + case maxT' of + Nothing -> iterateFor (n - 1) t1 maxT' xs' f + Just mt -> + if mt <= 0 + then pure xs + else iterateFor (n - 1) t1 maxT' xs' f + + evolve xs' = do + xs <- Prelude.mapM canonical xs' + parents' <- tournament xs + offspring <- combine parents' + if numParams cfg == 0 + then runEqSat myCost rewritesWithConstant 1 >> cleanDB >> refitChanged + else runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged + canonical offspring >>= updateIfNothing fitFun >> pure () + canonical offspring + + tournament xs = do + p1 <- applyTournament xs >>= canonical + p2 <- applyTournament xs >>= canonical + pure (p1, p2) + + applyTournament :: [EClassId] -> RndEGraph EClassId + applyTournament xs = do + challengers <- + replicateM (tournamentSize cfg) (rnd $ randomFrom xs) >>= traverse canonical + fits <- Prelude.map fromJust <$> Prelude.mapM getFitness challengers + pure . snd . maximumBy (compare `on` fst) $ Prelude.zip fits challengers + + combine (p1, p2) = crossover p1 p2 >>= mutate >>= canonical + + crossover p1 p2 = do + sz <- getSize p1 + coin <- rnd $ tossBiased (crossoverProbability cfg) + if sz == 1 || not coin + then rnd (randomFrom [p1, p2]) + else do + pos <- rnd $ randomRange (1, sz - 1) + cands <- getAllSubClasses p2 + tree <- getSubtree pos 0 Nothing [] cands p1 + fromTree myCost (relabel tree) >>= canonical + + getSubtree :: + Int -> + Int -> + Maybe (EClassId -> ENode) -> + [Maybe (EClassId -> ENode)] -> + [EClassId] -> + EClassId -> + RndEGraph (Fix SRTree) + getSubtree 0 sz (Just parent) mGrandParents cands p' = do + p <- canonical p' + candidates' <- + filterM (fmap (< maxExpressionSize cfg - sz) . getSize) cands + candidates <- + filterM (doesNotExistGens mGrandParents . parent) candidates' + >>= traverse canonical + if null candidates + then getBestExpr p + else do + subtree <- rnd (randomFrom candidates) + getBestExpr subtree + getSubtree pos sz parent mGrandParents cands p' = do + p <- canonical p' + root <- getBestENode p >>= canonize + case root of + Param ix -> pure . Fix $ Param ix + Const x -> pure . Fix $ Const x + Var ix -> pure . Fix $ Var ix + Uni f t' -> do + t <- canonical t' + Fix . Uni f + <$> getSubtree (pos - 1) (sz + 1) (Just $ Uni f) (parent : mGrandParents) cands t + Bin op l'' r'' -> do + l <- canonical l'' + r <- canonical r'' + szLft <- getSize l + szRgt <- getSize r + if szLft < pos + then do + l' <- getBestExpr l + r' <- + getSubtree + (pos - szLft - 1) + (sz + szLft + 1) + (Just $ Bin op l) + (parent : mGrandParents) + cands + r + pure . Fix $ Bin op l' r' + else do + l' <- + getSubtree + (pos - 1) + (sz + szRgt + 1) + (Just (\t -> Bin op t r)) + (parent : mGrandParents) + cands + l + r' <- getBestExpr r + pure . Fix $ Bin op l' r' + + getAllSubClasses p' = do + p <- canonical p' + en <- getBestENode p + case en of + Bin _ l r -> do + ls <- getAllSubClasses l + rs <- getAllSubClasses r + pure (p : ls <> rs) + Uni _ t -> (p :) <$> getAllSubClasses t + _ -> pure [p] + + mutate p = do + sz <- getSize p + coin <- rnd $ tossBiased (mutationProbability cfg) + if coin + then do + pos <- rnd $ randomRange (0, sz - 1) + tree <- mutAt pos (maxExpressionSize cfg) Nothing p + fromTree myCost (relabel tree) >>= canonical + else pure p + + peel :: Fix SRTree -> SRTree () + peel (Fix (Bin op _ _)) = Bin op () () + peel (Fix (Uni f _)) = Uni f () + peel (Fix (Param ix)) = Param ix + peel (Fix (Var ix)) = Var ix + peel (Fix (Const x)) = Const x + + mutAt :: + Int -> Int -> Maybe (EClassId -> ENode) -> EClassId -> RndEGraph (Fix SRTree) + mutAt 0 sizeLeft Nothing _ = insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical >>= getBestExpr + mutAt 0 1 _ _ = rnd $ randomFrom terms + mutAt 0 sizeLeft (Just parent) _ = do + ec <- insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical + (Fix tree) <- getBestExpr ec + root <- getBestENode ec + exist <- canonize (parent ec) >>= doesExist + if exist + then do + let children = childrenOf root + candidates <- case length children of + 0 -> + filterM + (checkToken parent . replaceChildren children) + (Prelude.map peel terms) + 1 -> filterM (checkToken parent . replaceChildren children) uniNonTerms + 2 -> filterM (checkToken parent . replaceChildren children) binNonTerms + _ -> pure [] + if null candidates + then pure $ Fix tree + else do + newToken <- rnd (randomFrom candidates) + pure . Fix $ replaceChildren (childrenOf tree) newToken + else pure . Fix $ tree + mutAt pos sizeLeft _ p' = do + p <- canonical p' + root <- getBestENode p >>= canonize + case root of + Param ix -> pure . Fix $ Param ix + Const x -> pure . Fix $ Const x + Var ix -> pure . Fix $ Var ix + Uni f t' -> + canonical t' + >>= ( fmap (Fix . Uni f) + . mutAt (pos - 1) (sizeLeft - 1) (Just $ Uni f) + ) + Bin op ln rn -> do + l <- canonical ln + r <- canonical rn + szLft <- getSize l + szRgt <- getSize r + if szLft < pos + then do + l' <- getBestExpr l + r' <- mutAt (pos - szLft - 1) (sizeLeft - szLft - 1) (Just $ Bin op l) r + pure . Fix $ Bin op l' r' + else do + l' <- mutAt (pos - 1) (sizeLeft - szRgt - 1) (Just (\t -> Bin op t r)) l + r' <- getBestExpr r + pure . Fix $ Bin op l' r' + + printExpr' :: Int -> EClassId -> RndEGraph [String] + printExpr' ix ec' = do + ec <- canonical ec' + thetas' <- gets (fmap (_theta . _info) . (IM.!? ec) . _eClass) + bestExpr <- + (if simplifyExpressions cfg then simplifyEqSatDefault else id) + <$> getBestExpr ec + + let best' = + if shouldReparam then relabelParams bestExpr else relabelParamsOrder bestExpr + nParams' = countParamsUniq best' + fromSz (MA.Sz x) = x + nThetas = fmap (Prelude.map (fromSz . MA.size)) thetas' + (_, thetas) <- + if maybe False (Prelude.any (/= nParams')) nThetas + then fitFun best' + else pure (1.0, fromMaybe [] thetas') + + maxLoss <- negate . fromJust <$> getFitness ec + forM (Data.List.zip4 [(0 :: Int) ..] dataTrainVals dataTests thetas) $ \(view, (dataTrain, dataVal), dataTest, theta') -> do + let (x, y, mYErr) = dataTrain + (x_val, y_val, mYErr_val) = dataVal + (x_te, y_te, mYErr_te) = dataTest + distribution = lossFunction cfg + + expr = paramsToConst (MA.toList theta') best' + showNA z = if isNaN z then "" else show z + r2_train = r2 x y best' theta' + r2_val = r2 x_val y_val best' theta' + r2_te = r2 x_te y_te best' theta' + nll_train = nll distribution mYErr x y best' theta' + nll_val = nll distribution mYErr_val x_val y_val best' theta' + nll_te = nll distribution mYErr_te x_te y_te best' theta' + mdl_train = fractionalBayesFactor distribution mYErr x y theta' best' + mdl_val = fractionalBayesFactor distribution mYErr_val x_val y_val theta' best' + mdl_te = fractionalBayesFactor distribution mYErr_te x_te y_te theta' best' + vals = + intercalate "," $ + Prelude.map + showNA + [ nll_train + , nll_val + , nll_te + , maxLoss + , r2_train + , r2_val + , r2_te + , mdl_train + , mdl_val + , mdl_te + ] + thetaStr = intercalate ";" $ Prelude.map show (MA.toList theta') + showExprFun = if null varnames then showExpr else showExprWithVars (splitOn "," varnames) + showLatexFun = if null varnames then showLatex else showLatexWithVars (splitOn "," varnames) + pure $ + show ix + <> "," + <> show view + <> "," + <> showExprFun expr + <> "," + <> "\"" + <> showPython best' + <> "\"," + <> "\"$$" + <> showLatexFun best' + <> "$$\"," + <> thetaStr + <> "," + <> show @Int (countNodes $ convertProtectedOps expr) + <> "," + <> vals + + insertTerms = forM terms (fromTree myCost >=> canonical) + + paretoFront' _ maxSize' = go 1 (-(1.0 / 0.0)) + where + go :: Int -> Double -> RndEGraph [Fix SRTree] + go n f + | n > maxSize' = pure [] + | otherwise = do + ecList <- getBestExprWithSize n + if not (null ecList) + then do + let (ec', mf) = case ecList of + [] -> (0, Nothing) + (e : _) -> e + f' = fromJust mf + improved = f' >= f && not (isNaN f') && not (isInfinite f') + ec <- canonical ec' + if improved + then do + thetas' <- gets (fmap (_theta . _info) . (IM.!? ec) . _eClass) + bestExpr <- + relabelParams . (if simplifyExpressions cfg then simplifyEqSatDefault else id) + <$> getBestExpr ec + let t = case thetas' of + Just (h : _) -> paramsToConst (MA.toList h) bestExpr + _ -> Fix (Const 0) -- Not sure if this makes sense as a default. + ts <- go (n + 1) (max f f') + pure (t : ts) + else go (n + 1) (max f f') + else go (n + 1) f diff --git a/symbolic-regression.cabal b/symbolic-regression.cabal new file mode 100644 --- /dev/null +++ b/symbolic-regression.cabal @@ -0,0 +1,61 @@ +cabal-version: 3.0 +name: symbolic-regression +version: 0.1.0.0 +synopsis: Symbolic Regression in Haskell +license: MIT +license-file: LICENSE +author: DataHaskell +maintainer: mschavinda@gmail.com +category: Data +build-type: Simple +extra-doc-files: CHANGELOG.md README.md + +common warnings + ghc-options: -Wall + +library + import: warnings + exposed-modules: Symbolic.Regression + build-depends: base >= 4 && <5 + , dataframe ^>= 0.4 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , binary >=0.8.9.1 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.1 && <1.1 + , mtl >=2.2 && <2.4 + , optparse-applicative >=0.17 && <0.19 + , random >=1.2 && <1.4 + , scheduler >=2.0.0.1 && <3 + , split >=0.2.5 && <0.3 + , srtree >= 2.0.1.5 + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unliftio >=0.2.10 && <1 + , unliftio-core >=0.2.1 && <1 + , unordered-containers ==0.2.* + , text >= 2.0 && < 3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + , directory + , time + hs-source-dirs: src + default-language: Haskell2010 + +test-suite symbolic-regression-test + import: warnings + default-language: Haskell2010 + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + build-depends: + base >= 4 && <5, + symbolic-regression diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,4 @@ +module Main (main) where + +main :: IO () +main = putStrLn "Test suite not yet implemented."