symbolic-regression (empty) → 0.1.0.0
raw patch · 6 files changed
+898/−0 lines, 6 filesdep +attoparsecdep +attoparsec-exprdep +base
Dependencies added: attoparsec, attoparsec-expr, base, binary, bytestring, containers, dataframe, directory, dlist, exceptions, filepath, hashable, ieee754, lens, list-shuffle, massiv, mtl, optparse-applicative, random, scheduler, split, srtree, statistics, symbolic-regression, text, time, transformers, unliftio, unliftio-core, unordered-containers, vector, zlib
Files
- CHANGELOG.md +5/−0
- LICENSE +20/−0
- README.md +104/−0
- src/Symbolic/Regression.hs +704/−0
- symbolic-regression.cabal +61/−0
- test/Main.hs +4/−0
+ CHANGELOG.md view
@@ -0,0 +1,5 @@+# Revision history for symbolic-regression++## 0.1.0.0++* Basic integration with srtree
+ LICENSE view
@@ -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.
+ README.md view
@@ -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`
+ src/Symbolic/Regression.hs view
@@ -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
+ symbolic-regression.cabal view
@@ -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
+ test/Main.hs view
@@ -0,0 +1,4 @@+module Main (main) where++main :: IO ()+main = putStrLn "Test suite not yet implemented."