packages feed

srtree 2.0.0.4 → 2.0.1.2

raw patch · 33 files changed

+2007/−1961 lines, 33 filesdep ~massivdep ~splitnew-component:exe:eggpnew-component:exe:symregg

Dependency ranges changed: massiv, split

Files

ChangeLog.md view
@@ -1,5 +1,24 @@ # Changelog for srtree +## 2.0.1.2++- Fix bug where the parameters were printed as `t[:,ix]` instead of `t[ix]` in `showPython`++## 2.0.1.1++- MSE loss is now the default+- Renamed `--distribution` argument to `--loss` in eggp and easter.+- Fixed bug with Gaussian distribution and fixed number of parameters.+- Fixed bug in which `--number-params 0` would create parameters.+- Fixed bug in `rEGGression` that pattern matched equivalent expressions.+- Support to `--numpy` flag that prints the output as a numpy expression (experimental, eggp only).+- Support to `--simplify` flag that simplifies the expressions before displaying (experimental, eggp only).++## 2.0.1.0++- Support to Multiview Symbolic Regression in eggp and EASTER.+- Support to `--number-params` argument that limits the maximum number of parameters and allow repated parameters in an expression.+ ## 2.0.0.4  - Cleaned up test cases (they were deprecated), will include new ones later
README.md view
@@ -1,7 +1,18 @@ # srtree: A supporting library for tree-based symbolic regression  -`srtree` is a Haskell library that implements a tree-based structure for expressions and supporting functions to be used in the context of **symbolic regression**.+`srtree` is a Haskell library that implements a tree-based structure for expressions and supporting functions to be used in the context of **symbolic regression** (SR). +This repository is also the home for different algorithm implementations for SR and software tools to support the post-processing of SR models (please refer to their corresponding README files):++- [srsimplify](apps/srsimplify/README.md): a parser and simplification tool supporting the output of many popular SR algorithms.+- [srtools](apps/srtools/README.md): a tool that can be used to evaluate symbolic regression expressions and create nice reports with confidence intervals. +- [tinygp](apps/tinygp/README.md): a simple GP implementation based on tinyGP.+- [rEGGression](apps/rEGGression/README.md): nonlinear regression models exploration and query system with e-graphs (egg).+- [symregg](apps/symregg/README.md): Equality graph Assisted Search Technique for Equation Recovery.+- [eggp](apps/eggp/README.md): E-graph Genetic Programming.++## SRTree+ The expression structure is defined as a fixed-point of a mix of unary and binary tree. This makes it easier to implement supporting functions that requires the traversal of the trees. Also, since it is a parameterized structure, we can creating partial trees to pattern math structures of interest. This structure may contain four types of nodes: @@ -33,11 +44,6 @@ calculating the derivatives, printing, generating random trees, simplifying the expression, calculating overall statistics, optimizing parameters, and model selection metrics.  -Together with this library, we provide example applications (please refer to their corresponding README files):--- [srsimplify](apps/srsimplify/README.md): a parser and simplification tool supporting the output of many popular SR algorithms.-- [srtools](apps/srtools/README.md): a tool that can be used to evaluate symbolic regression expressions and create nice reports with confidence intervals. -- [tinygp](apps/tinygp/README.md): a simple GP implementation based on tinyGP.  ## Organization
+ apps/eggp/Main.hs view
@@ -0,0 +1,515 @@+{-# LANGUAGE  BlockArguments #-}+{-# LANGUAGE  TupleSections #-}+{-# LANGUAGE  MultiWayIf #-}+{-# LANGUAGE  OverloadedStrings #-}+{-# LANGUAGE  BangPatterns #-}+{-# LANGUAGE  TypeSynonymInstances, FlexibleInstances #-}++module Main where ++import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Simplify+import Algorithm.EqSat.Build+import Algorithm.EqSat.Queries+import Algorithm.EqSat.Info+import Algorithm.EqSat.DB+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection+import Algorithm.SRTree.Opt+import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.))+import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_)+import Control.Monad.State.Strict+import qualified Data.IntMap.Strict as IM+import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.Maybe (fromJust, isNothing, isJust)+import Data.SRTree+import Data.SRTree.Datasets+import Data.SRTree.Eval+import Data.SRTree.Random (randomTree)+import Data.SRTree.Print+import Options.Applicative as Opt hiding (Const)+import Random+import System.Random+import qualified Data.HashSet as Set+import Data.List ( sort, maximumBy, intercalate, sortOn, intersperse, nub )+import Data.IntSet (IntSet)+import qualified Data.IntSet as IntSet+import qualified Data.Sequence as FingerTree+import Data.Function ( on )+import qualified Data.Foldable as Foldable+import qualified Data.IntMap as IntMap+import List.Shuffle ( shuffle )+import Algorithm.SRTree.NonlinearOpt+import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS++import Debug.Trace+import Algorithm.EqSat (runEqSat,applySingleMergeOnlyEqSat)++import GHC.IO (unsafePerformIO)+import Control.Scheduler +import Control.Monad.IO.Unlift+import Data.SRTree (convertProtectedOps)++import Util++egraphGP :: [(DataSet, DataSet)] -> [DataSet] -> Args -> StateT EGraph (StateT StdGen IO) ()+egraphGP dataTrainVals dataTests args = do+  when ((not.null) (_loadFrom args)) $ (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg)++  pop <- replicateM (_nPop args) $ do ec <- insertRndExpr (_maxSize args) rndTerm rndNonTerm >>= canonical+                                      updateIfNothing fitFun ec+                                      pure ec+  insertTerms+  --runEqSat myCost rewritesParams 1+  --cleanDB+  evaluateUnevaluated fitFun+  pop' <- Prelude.mapM canonical pop+  when (_trace args) $ printPop pop'+  let m = (_nPop args) `div` (_maxSize args)++  finalPop <- iterateFor (_gens args) pop' $ \it ps' -> do+    -- ps <- Prelude.mapM canonical ps'+    --let chunks = 1+    --    (parts, r) = _nPop args `divMod` chunks -- chunks of 50+    --    genChunk = Prelude.mapM canonical ps' >>= \ps -> replicateM chunks (tournament ps) >>= Prelude.mapM combine+    --    genPart  = genChunk >>= \chunk -> (runEqSat myCost rewritesSimple 1 >> cleanDB) >> pure chunk+    --newPop' <- (<>) <$> (concat <$> (replicateM parts genPart)) <*> (concat <$> replicateM r genPart)+    --parents <- replicateM (_nPop args - if (_moo args) then (_maxSize args) else 0) (tournament ps)+    --newPop' <- Prelude.mapM combine parents+    --runEqSat myCost rewritesSimple 1 --applySingleMergeOnlyEqSat myCost rewritesSimple+    --cleanDB+    newPop' <- replicateM (_nPop args) (evolve ps')+    --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB+    --newPop' <- Prelude.mapM (\eId -> canonical eId >>= \eId' -> (updateIfNothing fitFun eId' >> pure eId')) newPop''+    --Prelude.mapM_ (updateIfNothing fitFun) newPop'++    totSz <- gets (Map.size . _eNodeToEClass) -- (IntMap.size . _eClass)+    let full = totSz > max maxMem (_nPop args)+    --io . print $ totSz+    when full (cleanEGraph >> cleanDB)++    newPop <- if (_moo args)+                then do+                        let n_paretos = (_nPop args) `div` (_maxSize args)+                        pareto <- concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n 2)+                        let remainder = _nPop args - length pareto+                        lft <- if full+                                  then getTopFitEClassThat remainder (const True)+                                  else pure $ Prelude.take remainder newPop'+                        Prelude.mapM canonical (pareto <> lft)+                else if full+                       then getTopFitEClassThat (_nPop args) (const True)+                       else Prelude.mapM canonical newPop'+    when (_trace args) $ printPop newPop+    pure newPop++  io $ putStrLn "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"+  if (_printPareto args)+    then paretoFront fitFun (_maxSize args) printExpr+    else printBest fitFun printExpr+  when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode )+  where+    maxSize = (_maxSize args)+    maxMem = 2000000 -- running 1 iter of eqsat for each new individual will consume ~3GB+    fitFun = fitnessMV shouldReparam (_optRepeat args) (_optIter args) (_distribution args) dataTrainVals+    nonTerms   = parseNonTerms (_nonterminals args)+    (Sz2 _ nFeats) = MA.size (getX .fst . head $ dataTrainVals)+    params         = if _nParams args == -1 then [param 0] else Prelude.map param [0 .. _nParams args - 1]+    shouldReparam  = _nParams args == -1+    relabel        = if shouldReparam then relabelParams else relabelParamsOrder+    terms          = if _distribution args == ROXY+                          then (var 0 : params)+                          else [var ix | ix <- [0 .. nFeats-1]] -- <> params+    uniNonTerms = [t | t <- nonTerms, isUni t]+    binNonTerms = [t | t <- nonTerms, isBin t]+    isUni (Uni _ _)   = True+    isUni _           = False+    isBin (Bin _ _ _) = True+    isBin _           = False++    -- TODO: merge two or more egraphs+    cleanEGraph = do let nParetos = 10 -- (maxMem `div` 5) `div` _maxSize args+                     io . putStrLn $ "cleaning"+                     pareto <- (concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n nParetos))+                                 >>= Prelude.mapM canonical+                     infos  <- forM pareto (\c -> gets (_info . (IntMap.! c) . _eClass))+                     exprs  <- forM pareto getBestExpr+                     put emptyGraph+                     newIds <- fromTrees myCost $ Prelude.map relabel exprs+                     forM_ (Prelude.zip newIds (Prelude.reverse infos)) $ \(eId, info) ->+                         insertFitness eId (fromJust $ _fitness info) (_theta info)++    rndTerm    = do coin <- toss+                    if coin then Random.randomFrom terms else Random.randomFrom params+    rndNonTerm = Random.randomFrom nonTerms++    refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList >>= pure . 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 f = pure xs+    iterateFor n xs f = do xs' <- f n xs+                           iterateFor (n-1) xs' f++    evolve xs' = do xs <- Prelude.mapM canonical xs'+                    parents <- tournament xs+                    offspring <- combine parents+                    --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB+                    if _nParams args == 0+                       then runEqSat myCost rewritesWithConstant 1 >> cleanDB >> refitChanged+                       else runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged+                    canonical offspring >>= updateIfNothing fitFun+                    canonical offspring+                    --pure offspring++    tournament xs = do p1 <- applyTournament xs >>= canonical+                       p2 <- applyTournament xs >>= canonical+                       pure (p1, p2)++    applyTournament :: [EClassId] -> RndEGraph EClassId+    applyTournament xs = do challengers <- replicateM (_nTournament args) (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 (_pc args)+                         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 (\c -> (<maxSize-sz) <$> getSize c) cands+      candidates  <- filterM (\c -> doesNotExistGens mGrandParents (parent c)) 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 (_pm args)+                  if coin+                     then do pos <- rnd $ randomRange (0, sz-1)+                             tree <- mutAt pos maxSize Nothing p+                             fromTree myCost (relabel tree) >>= canonical+                     else pure p++    peel :: Fix SRTree -> SRTree ()+    peel (Fix (Bin op l r)) = Bin op () ()+    peel (Fix (Uni f t)) = 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 -- we chose to mutate the root+    mutAt 0 1        _             _ = rnd $ randomFrom terms -- we don't have size left+    mutAt 0 sizeLeft (Just parent) _ = do -- we reached the mutation place+      ec    <- insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical -- create a random expression with the size limit+      (Fix tree) <- getBestExpr ec           --+      root  <- getBestENode ec+      exist <- canonize (parent ec) >>= doesExist+      if exist+         -- the expression `parent ec` already exists, try to fix+         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+                 if null candidates+                     then pure $ Fix tree -- there's no candidate, so we failed and admit defeat+                     else do newToken <- rnd (randomFrom candidates)+                             pure . Fix $ replaceChildren (childrenOf tree) newToken++         else pure . Fix $ tree++    mutAt pos sizeLeft parent 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' >>= \t -> (Fix . Uni f) <$> mutAt (pos-1) (sizeLeft-1) (Just $ Uni f) t+          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'++    checkToken parent en' = do en <- canonize en'+                               mEc <- gets ((Map.!? en) . _eNodeToEClass)+                               case mEc of+                                    Nothing -> pure True+                                    Just ec -> do ec' <- canonical ec+                                                  ec'' <- canonize (parent ec')+                                                  not <$> doesExist ec''+    doesExist, doesNotExist :: ENode -> RndEGraph Bool+    doesExist en = gets ((Map.member en) . _eNodeToEClass)+    doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)++    doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool+    doesNotExistGens []              en = gets ((Map.notMember en) . _eNodeToEClass)+    doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass)+                                             if b+                                               then pure True+                                               else case mGrand of+                                                    Nothing -> pure False+                                                    Just gf -> do ec  <- gets ((Map.! en) . _eNodeToEClass)+                                                                  en' <- canonize (gf ec)+                                                                  doesNotExistGens grands en'++    printExpr :: Int -> EClassId -> RndEGraph ()+    printExpr ix ec = do+        thetas' <- gets (_theta . _info . (IM.! ec) . _eClass)+        bestExpr <- (if _simplify args then simplifyEqSatDefault else id) <$> getBestExpr ec+        let nParams = countParamsUniq bestExpr+            fromSz (MA.Sz x) = x +            nThetas  = Prelude.map (fromSz . MA.size) thetas'+        (_, thetas) <- if Prelude.any (/=nParams) nThetas || _simplify args+                        then fitFun bestExpr+                        else pure (1.0, thetas')++        forM_ (Prelude.zip3 dataTrainVals dataTests thetas) $ \((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 = _distribution args+                best'     = if shouldReparam then relabelParams bestExpr else relabelParamsOrder bestExpr+                expr      = paramsToConst (MA.toList theta) best'+                mse_train = mse x y best' theta+                mse_val   = mse x_val y_val best' theta+                mse_te    = mse x_te y_te best' theta+                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  = mdl distribution mYErr x y theta best'+                mdl_val    = mdl distribution mYErr_val x_val y_val theta best'+                mdl_te     = mdl distribution mYErr_te x_te y_te theta best'+                vals       = intercalate ","+                           $ Prelude.map show [mse_train, mse_val, mse_te+                                              , r2_train, r2_val, r2_te+                                              , nll_train, nll_val, nll_te+                                              , mdl_train, mdl_val, mdl_te]+                thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)+                showFun    = if _numpy args then "\"" <> showPython best' <> "\"" else showExpr expr+            io . putStrLn $ show ix <> "," <> showFun <> ","+                          <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)+                          <> "," <> vals++    printPop pop = forM_ pop $ \ecN'-> do+            ecN'' <- canonical ecN'+            _tree <- getBestExpr ecN''+            fi <- fromJust <$> getFitness ecN''+            thetas <- getTheta ecN''+            let thetaStr   = intercalate "_" $ Prelude.map (intercalate ";" . Prelude.map show . MA.toList) thetas+                tree = if _simplify args then simplifyEqSatDefault _tree else _tree+                showFun = if _numpy args then showPython else showExpr+            io . putStrLn $ showFun tree <> "," <> thetaStr <> "," <> show fi+            pure ()++    insertTerms =+        forM terms $ \t -> do fromTree myCost t >>= canonical++data Args = Args+  { _dataset      :: String,+    _testData     :: String,+    _gens         :: Int,+    _maxSize      :: Int,+    _split        :: Int,+    _printPareto  :: Bool,+    _trace        :: Bool,+    _distribution :: Distribution,+    _optIter      :: Int,+    _optRepeat    :: Int,+    _nParams      :: Int,+    _nPop         :: Int,+    _nTournament  :: Int,+    _pc           :: Double,+    _pm           :: Double,+    _nonterminals :: String,+    _dumpTo       :: String,+    _loadFrom     :: String,+    _moo          :: Bool,+    _numpy        :: Bool,+    _simplify     :: Bool+  }+  deriving (Show)++-- parser of command line arguments+opt :: Parser Args+opt = Args+   <$> strOption+       ( long "dataset"+       <> short 'd'+       <> metavar "INPUT-FILE"+       <> help "CSV dataset." )+  <*> strOption+       ( long "test"+       <> short 't'+       <> value ""+       <> showDefault+       <> help "test data")+   <*> option auto+      ( long "generations"+      <> short 'g'+      <> metavar "GENS"+      <> showDefault+      <> value 100+      <> help "Number of generations." )+  <*> option auto+       ( long "maxSize"+       <> short 's'+       <> help "max-size." )+  <*> option auto+       ( long "split"+       <> short 'k'+       <> value 1 +       <> showDefault+       <> help "k-split ratio training-validation")+  <*> switch+       ( long "print-pareto"+       <> help "print Pareto front instead of best found expression")+  <*> switch+       ( long "trace"+       <> help "print all evaluated expressions.")+  <*> option auto+       ( long "loss"+       <> value MSE+       <> showDefault+       <> help "loss function: MSE, Gaussian, Poisson, Bernoulli.")+  <*> option auto+       ( long "opt-iter"+       <> value 30+       <> showDefault+       <> help "number of iterations in parameter optimization.")+  <*> option auto+       ( long "opt-retries"+       <> value 1+       <> showDefault+       <> help "number of retries of parameter fitting.")+  <*> option auto+       ( long "number-params"+       <> value (-1)+       <> showDefault+       <> help "maximum number of parameters in the model. If this argument is absent, the number is bounded by the maximum size of the expression and there will be no repeated parameter.")+  <*> option auto+       ( long "nPop"+       <> value 100+       <> showDefault+       <> help "population size (Default: 100).")+  <*> option auto+       ( long "tournament-size"+       <> value 2+       <> showDefault+       <> help "tournament size.")+  <*> option auto+       ( long "pc"+       <> value 1.0+       <> showDefault+       <> help "probability of crossover.")+  <*> option auto+       ( long "pm"+       <> value 0.3+       <> showDefault+       <> help "probability of mutation.")+  <*> strOption+       ( long "non-terminals"+       <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+       <> showDefault+       <> help "set of non-terminals to use in the search."+       )+  <*> strOption+       ( long "dump-to"+       <> value ""+       <> showDefault+       <> help "dump final e-graph to a file."+       )+  <*> strOption+       ( long "load-from"+       <> value ""+       <> showDefault+       <> help "load initial e-graph from a file."+       )+  <*> switch+       ( long "moo"+       <> help "replace the current population with the pareto front instead of replacing it with the generated children."+       )+  <*> switch+       ( long "to-numpy"+       <> help "outputs the expressions using a numpy format."+       )+  <*> switch+       ( long "simplify"+       <> help "simplify the expressions before displaying them."+       )+main :: IO ()+main = do+  args <- execParser opts+  g    <- getStdGen+  let datasets = words (_dataset args)+  dataTrains' <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets +  dataTests   <- if null (_testData args)+                  then pure dataTrains'+                  else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args)++  let (dataTrainVals, g') = runState (Prelude.mapM (`splitData` (_split args)) dataTrains') g+      alg = evalStateT (egraphGP dataTrainVals dataTests args) emptyGraph+  evalStateT alg g'+  where+    opts = Opt.info (opt <**> helper)+            ( fullDesc <> progDesc "An implementation of GP with modified crossover and mutation\+                                   \ operators designed to exploit equality saturation and e-graphs.\+                                   \ https://arxiv.org/abs/2501.17848\n"+           <> header "eggp - E-graph Genetic Programming for Symbolic Regression." )
+ apps/eggp/Random.hs view
@@ -0,0 +1,54 @@+module Random where ++import System.Random +import Control.Monad.State.Strict+import Control.Monad+import Data.SRTree +import Data.SRTree.Eval+import Data.Massiv.Array as MA++type Rng a = StateT StdGen IO a ++toss :: Rng Bool+toss = state random+{-# INLINE toss #-}++tossBiased :: Double -> Rng Bool+tossBiased p = do r <- state random +                  pure (r < p)++randomVal :: Rng Double+randomVal = state random++randomRange :: (Ord val, Random val) => (val, val) -> Rng val+randomRange rng = state (randomR rng)+{-# INLINE randomRange #-}++randomFrom :: [a] -> Rng a+randomFrom funs = do n <- randomRange (0, length funs - 1)+                     pure $ funs !! n+{-# INLINE randomFrom #-}++randomVec :: Int -> Rng PVector+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))++randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree)+randomTree minDepth maxDepth maxSize genTerm genNonTerm grow +  | noSpaceLeft = genTerm+  | needNonTerm = genRecursion +  | otherwise   = do r <- toss+                     if r +                       then genTerm +                       else genRecursion+  where +    noSpaceLeft = maxDepth <= 1 || maxSize <= 2+    needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2++    genRecursion = do +        node <- genNonTerm+        case node of +          Uni f _    -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow +          Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow+                           r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow +                           pure . Fix  $ Bin op l r+{-# INLINE randomTree #-}
+ apps/eggp/Util.hs view
@@ -0,0 +1,266 @@+{-# LANGUAGE  BlockArguments #-}+{-# LANGUAGE  TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Random+import Algorithm.SRTree.Likelihoods+import Data.SRTree.Print+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries++import Debug.Trace++type RndEGraph a = EGraphST (StateT StdGen IO) a+type DataSet = (SRMatrix, PVector, Maybe PVector)++csvHeader :: String+csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"++io :: IO a -> RndEGraph a+io = lift . lift+{-# INLINE io #-}+rnd :: StateT StdGen IO a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _)     = 1+myCost (Const _)   = 1+myCost (Param _)   = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t)   = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()+while p arg prog = do when (p arg) do arg' <- prog arg+                                      while p arg' prog++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =+  if isNaN val -- || isNaN tr+    then (-(1/0), theta) -- infinity+    else (val, theta)+  where+    --tree          = relabelParams _tree+    nParams       = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+    (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+    --tr            = evalF x y mYErr+    val           = evalF x_val y_val mYErr_val++--{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nRep nIter distribution dataTrain dataVal tree = do+    let nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+    thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+    let fits = maximumBy (compare `on` fst) $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal tree) thetaOrigs+    pure fits+--{-# INLINE fitnessFunRep #-}+++fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [PVector])+fitnessMV shouldReparam nRep nIter distribution dataTrainsVals _tree = do+  let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree+  response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv tree+  pure (minimum (Prelude.map fst response), Prelude.map snd response)++-- helper query functions+-- TODO: move to egraph lib+getFitness :: EClassId -> RndEGraph (Maybe Double)+getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: EClassId -> RndEGraph ([PVector])+getTheta c = gets (_theta . _info . (IM.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: EClassId -> RndEGraph Int+getSize c = gets (_size . _info . (IM.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}++getBestFitness :: RndEGraph (Maybe Double)+getBestFitness = do+    bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+    gets (_fitness . _info . (IM.! bec) . _eClass)++-- TODO: move to dataset lib+chunksOf :: Int -> [e] -> [[e]]+chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))+ where+  splitter :: [e] -> ([e] -> a -> a) -> a -> a+  splitter [] _ n = n+  splitter l c n = l `c` splitter (Prelude.drop i l) c n+  build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]+  build g = g (:) []++splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)+splitData (x, y, mYErr) k = do+  if k == 1+    then pure ((x, y, mYErr), (x, y, mYErr))+    else do+      ixs' <- (state . shuffle) [0 .. sz-1]+      let ixs = chunksOf k ixs'++      let (x_tr, x_te) = getX ixs x+          (y_tr, y_te) = getY ixs y+          mY = fmap (getY ixs) mYErr+          (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)+      pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))+  where+    (MA.Sz sz) = MA.size y+    comp_x     = MA.getComp x+    comp_y     = MA.getComp y++    getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)+    getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]+                    in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+                       , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+                       )+    getY :: [[Int]] -> PVector -> (PVector, PVector)+    getY ixs ys  = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+                   , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+                   )++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b++parseNonTerms :: String -> [SRTree ()]+parseNonTerms = Prelude.map toNonTerm . splitOn ","+  where+    binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+    uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+    toNonTerm xs' = let xs = Prelude.map toLower xs'+                    in case binTerms Map.!? xs of+                          Just op -> Bin op () ()+                          Nothing -> case uniTerms Map.!? xs of+                                          Just f -> Uni f ()+                                          Nothing -> error $ "invalid non-terminal " <> show xs++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId+insertExpr t fitFun = do+    ecId <- fromTree myCost t >>= canonical+    (f, p) <- fitFun t+    insertFitness ecId f p+    io . putStrLn $ "Best fit global: " <> show f+    pure ecId+  where powabs l r  = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+      mf <- getFitness ec+      case mf of+        Nothing -> do+          t <- getBestExpr ec+          (f, p) <- fitFun t+          insertFitness ec f p+          pure True+        Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+                    if not (null ecIds)+                          then do rndId' <- rnd $ randomFrom ecIds+                                  rndId  <- canonical rndId'+                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+                                  case constType of+                                    NotConst -> pure $ Just rndId+                                    _        -> pure Nothing+                          else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)++getBestExprWithSize n =+        do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+           if (not (null ec))+            then do+              bestFit <- getFitness $ head ec+              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+              pure [(head ec, bestFit)]+            else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+      do grow <- rnd toss+         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+         t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+         fromTree myCost t >>= canonical++refit fitFun ec = do+  t <- getBestExpr ec+  (f, p) <- fitFun t+  insertFitness ec f p++--printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest fitFun printExprFun = do+      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+      bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)+      --refit fitFun bec+      --io.print $ "should be " <> show bestFit+      printExprFun 0 bec++--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))+    where+    go :: Int -> Int -> Double -> RndEGraph ()+    go n ix f+        | n > maxSize = pure ()+        | otherwise   = do+            ecList <- getBestExprWithSize n+            if not (null ecList)+                then do let (ec, mf) = head ecList+                            improved = fromJust mf > f+                        ec' <- canonical ec+                        when improved $ refit fitFun ec' >> printExprFun ix ec'+                        go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))+                else go (n+1) ix f++evaluateUnevaluated fitFun = do+          ec <- gets (IntSet.toList . _unevaluated . _eDB)+          forM_ ec $ \c -> do+              t <- getBestExpr c+              (f, p) <- fitFun t+              insertFitness c f p++evaluateRndUnevaluated fitFun = do+          ec <- gets (IntSet.toList . _unevaluated . _eDB)+          c <- rnd . randomFrom $ ec+          t <- getBestExpr c+          (f, p) <- fitFun t+          insertFitness c f p+          pure c
− apps/egraphGP/Main.hs 
@@ -1,486 +0,0 @@-{-# LANGUAGE  BlockArguments #-}-{-# LANGUAGE  TupleSections #-}-{-# LANGUAGE  MultiWayIf #-}-{-# LANGUAGE  OverloadedStrings #-}-{-# LANGUAGE  BangPatterns #-}-{-# LANGUAGE  TypeSynonymInstances, FlexibleInstances #-}--module Main where --import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Simplify-import Algorithm.EqSat.Build-import Algorithm.EqSat.Queries-import Algorithm.EqSat.Info-import Algorithm.EqSat.DB-import Algorithm.SRTree.Likelihoods-import Algorithm.SRTree.ModelSelection-import Algorithm.SRTree.Opt-import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.))-import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_)-import Control.Monad.State.Strict-import qualified Data.IntMap.Strict as IM-import qualified Data.Map.Strict as Map-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.Maybe (fromJust, isNothing, isJust)-import Data.SRTree-import Data.SRTree.Datasets-import Data.SRTree.Eval-import Data.SRTree.Random (randomTree)-import Data.SRTree.Print-import Options.Applicative as Opt hiding (Const)-import Random-import System.Random-import qualified Data.HashSet as Set-import Data.List ( sort, maximumBy, intercalate, sortOn, intersperse, nub )-import Data.IntSet (IntSet)-import qualified Data.IntSet as IntSet-import qualified Data.Sequence as FingerTree-import Data.Function ( on )-import qualified Data.Foldable as Foldable-import qualified Data.IntMap as IntMap-import List.Shuffle ( shuffle )-import Algorithm.SRTree.NonlinearOpt-import Data.Binary ( encode, decode )-import qualified Data.ByteString.Lazy as BS--import Debug.Trace-import Algorithm.EqSat (runEqSat,applySingleMergeOnlyEqSat)--import GHC.IO (unsafePerformIO)-import Control.Scheduler -import Control.Monad.IO.Unlift-import Data.SRTree (convertProtectedOps)--import Util--egraphGP :: DataSet -> DataSet -> DataSet -> Args -> StateT EGraph (StateT StdGen IO) ()-egraphGP dataTrain dataVal dataTest args = do-  when ((not.null) (_loadFrom args)) $ (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg)--  pop <- replicateM (_nPop args) $ do ec <- insertRndExpr (_maxSize args) rndTerm rndNonTerm >>= canonical-                                      updateIfNothing fitFun ec-                                      pure ec-  insertTerms-  --runEqSat myCost rewritesParams 1-  --cleanDB-  evaluateUnevaluated fitFun-  pop' <- Prelude.mapM canonical pop-  when (_trace args) $ printPop pop'-  let m = (_nPop args) `div` (_maxSize args)--  finalPop <- iterateFor (_gens args) pop' $ \it ps' -> do-    -- ps <- Prelude.mapM canonical ps'-    --let chunks = 1-    --    (parts, r) = _nPop args `divMod` chunks -- chunks of 50-    --    genChunk = Prelude.mapM canonical ps' >>= \ps -> replicateM chunks (tournament ps) >>= Prelude.mapM combine-    --    genPart  = genChunk >>= \chunk -> (runEqSat myCost rewritesSimple 1 >> cleanDB) >> pure chunk-    --newPop' <- (<>) <$> (concat <$> (replicateM parts genPart)) <*> (concat <$> replicateM r genPart)-    --parents <- replicateM (_nPop args - if (_moo args) then (_maxSize args) else 0) (tournament ps)-    --newPop' <- Prelude.mapM combine parents-    --runEqSat myCost rewritesSimple 1 --applySingleMergeOnlyEqSat myCost rewritesSimple-    --cleanDB-    newPop' <- replicateM (_nPop args) (evolve ps')-    --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB-    --newPop' <- Prelude.mapM (\eId -> canonical eId >>= \eId' -> (updateIfNothing fitFun eId' >> pure eId')) newPop''-    --Prelude.mapM_ (updateIfNothing fitFun) newPop'--    totSz <- gets (IntMap.size . _eClass)-    let full = False -- totSz > max maxMem (_nPop args)-    when full cleanEGraph--    newPop <- if (_moo args)-                then do-                        let n_paretos = (_nPop args) `div` (_maxSize args)-                        pareto <- concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n 2)-                        let remainder = _nPop args - length pareto-                        lft <- if full-                                  then getTopFitEClassThat remainder (const True)-                                  else pure $ Prelude.take remainder newPop'-                        Prelude.mapM canonical (pareto <> lft)-                else if full-                       then getTopFitEClassThat (_nPop args) (const True)-                       else Prelude.mapM canonical newPop'-    when (_trace args) $ printPop newPop---    pure newPop--  io $ putStrLn "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"-  if (_printPareto args)-    then paretoFront fitFun (_maxSize args) printExpr-    else printBest fitFun printExpr-  when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode )-  where-    maxSize = (_maxSize args)-    maxMem = 10000 -- running 1 iter of eqsat for each new individual will consume ~3GB-    fitFun = fitnessFunRep (_optRepeat args) (_optIter args) (_distribution args) dataTrain dataVal-    nonTerms   = parseNonTerms (_nonterminals args)-    (Sz2 _ nFeats) = MA.size (getX dataTrain)-    terms          = if _distribution args == ROXY-                          then [var 0, param 0]-                          else [var ix | ix <- [0 .. nFeats-1]] <> [param 0]-    uniNonTerms = [t | t <- nonTerms, isUni t]-    binNonTerms = [t | t <- nonTerms, isBin t]-    isUni (Uni _ _) = True-    isUni _         = False-    isBin (Bin _ _ _) = True-    isBin _           = False--    -- TODO: merge two or more egraphs-    cleanEGraph = do let nParetos = (maxMem `div` 5) `div` _maxSize args-                     pareto <- (concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n nParetos))-                                 >>= Prelude.mapM canonical-                     infos  <- forM pareto (\c -> gets (_info . (IntMap.! c) . _eClass))-                     exprs  <- forM pareto getBestExpr-                     put emptyGraph-                     newIds <- fromTrees myCost exprs-                     forM_ (Prelude.zip newIds (Prelude.reverse infos)) $ \(eId, info) ->-                         insertFitness eId (fromJust $ _fitness info) (fromJust $ _theta info)--    rndTerm    = Random.randomFrom terms-    rndNonTerm = Random.randomFrom nonTerms--    refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList >>= pure . 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 f = pure xs-    iterateFor n xs f = do xs' <- f n xs-                           iterateFor (n-1) xs' f--    evolve xs' = do xs <- Prelude.mapM canonical xs'-                    parents <- tournament xs-                    offspring <- combine parents-                    --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB-                    runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged-                    canonical offspring >>= updateIfNothing fitFun-                    canonical offspring-                    --pure offspring--    tournament xs = do p1 <- applyTournament xs >>= canonical-                       p2 <- applyTournament xs >>= canonical-                       pure (p1, p2)--    applyTournament :: [EClassId] -> RndEGraph EClassId-    applyTournament xs = do challengers <- replicateM (_nTournament args) (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 (_pc args)-                         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 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 (\c -> (<maxSize-sz) <$> getSize c) cands-      candidates  <- filterM (\c -> doesNotExistGens mGrandParents (parent c)) 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 (_pm args)-                  if coin-                     then do pos <- rnd $ randomRange (0, sz-1)-                             tree <- mutAt pos maxSize Nothing p-                             fromTree myCost tree >>= canonical-                     else pure p--    peel :: Fix SRTree -> SRTree ()-    peel (Fix (Bin op l r)) = Bin op () ()-    peel (Fix (Uni f t)) = 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 -- we chose to mutate the root-    mutAt 0 1        _             _ = rnd $ randomFrom terms -- we don't have size left-    mutAt 0 sizeLeft (Just parent) _ = do -- we reached the mutation place-      ec    <- insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical -- create a random expression with the size limit-      (Fix tree) <- getBestExpr ec           ---      root  <- getBestENode ec-      exist <- canonize (parent ec) >>= doesExist-      if exist-         -- the expression `parent ec` already exists, try to fix-         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-                 if null candidates-                     then pure $ Fix tree -- there's no candidate, so we failed and admit defeat-                     else do newToken <- rnd (randomFrom candidates)-                             pure . Fix $ replaceChildren (childrenOf tree) newToken--         else pure . Fix $ tree--    mutAt pos sizeLeft parent 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' >>= \t -> (Fix . Uni f) <$> mutAt (pos-1) (sizeLeft-1) (Just $ Uni f) t-          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'--    checkToken parent en' = do en <- canonize en'-                               mEc <- gets ((Map.!? en) . _eNodeToEClass)-                               case mEc of-                                    Nothing -> pure True-                                    Just ec -> do ec' <- canonical ec-                                                  ec'' <- canonize (parent ec')-                                                  not <$> doesExist ec''-    doesExist, doesNotExist :: ENode -> RndEGraph Bool-    doesExist en = gets ((Map.member en) . _eNodeToEClass)-    doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)--    doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool-    doesNotExistGens []              en = gets ((Map.notMember en) . _eNodeToEClass)-    doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass)-                                             if b-                                               then pure True-                                               else case mGrand of-                                                    Nothing -> pure False-                                                    Just gf -> do ec  <- gets ((Map.! en) . _eNodeToEClass)-                                                                  en' <- canonize (gf ec)-                                                                  doesNotExistGens grands en'--    printExpr :: Int -> EClassId -> RndEGraph ()-    printExpr ix ec = do-        theta' <- gets (fromJust . _theta . _info . (IM.! ec) . _eClass)-        bestExpr <- getBestExpr ec-        let nParams = countParams bestExpr-            (MA.Sz nTheta)  = MA.size theta'-        (_, theta) <- if (nParams /= nTheta)-                        then fitFun bestExpr-                        else pure (1.0, theta')--        let (x, y, mYErr) = dataTrain-            (x_val, y_val, mYErr_val) = dataVal-            (x_te, y_te, mYErr_te) = dataTest-            distribution = _distribution args-            best'     = relabelParams bestExpr-            expr      = paramsToConst (MA.toList theta) best'-            mse_train = mse x y best' theta-            mse_val   = mse x_val y_val best' theta-            mse_te    = mse x_te y_te best' theta-            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  = mdl distribution mYErr x y theta best'-            mdl_val    = mdl distribution mYErr_val x_val y_val theta best'-            mdl_te     = mdl distribution mYErr_te x_te y_te theta best'-            vals       = intercalate ","-                       $ Prelude.map show [mse_train, mse_val, mse_te-                                          , r2_train, r2_val, r2_te-                                          , nll_train, nll_val, nll_te-                                          , mdl_train, mdl_val, mdl_te]-            thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)-        io . putStrLn $ show ix <> "," <> showExpr expr <> ","-                      <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)-                      <> "," <> vals--    printPop pop = forM_ pop $ \ecN'-> do-            ecN'' <- canonical ecN'-            _tree <- getBestExpr ecN''-            fi <- fromJust <$> getFitness ecN''-            theta <- fromJust <$> getTheta ecN''-            let thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)-            io . putStrLn $ showExpr _tree <> "," <> thetaStr <> "," <> show fi-            pure ()--    insertTerms =-        forM terms $ \t -> do fromTree myCost t >>= canonical--data Args = Args-  { _dataset      :: String,-    _testData     :: String,-    _gens         :: Int,-    _maxSize      :: Int,-    _split        :: Int,-    _printPareto  :: Bool,-    _trace        :: Bool,-    _distribution :: Distribution,-    _optIter      :: Int,-    _optRepeat    :: Int,-    _nPop         :: Int,-    _nTournament  :: Int,-    _pc           :: Double,-    _pm           :: Double,-    _nonterminals :: String,-    _dumpTo       :: String,-    _loadFrom     :: String,-    _moo          :: Bool-  }-  deriving (Show)---- parser of command line arguments-opt :: Parser Args-opt = Args-   <$> strOption-       ( long "dataset"-       <> short 'd'-       <> metavar "INPUT-FILE"-       <> help "CSV dataset." )-  <*> strOption-       ( long "test"-       <> short 't'-       <> value ""-       <> showDefault-       <> help "test data")-   <*> option auto-      ( long "generations"-      <> short 'g'-      <> metavar "GENS"-      <> showDefault-      <> value 100-      <> help "Number of generations." )-  <*> option auto-       ( long "maxSize"-       <> short 's'-       <> help "max-size." )-  <*> option auto-       ( long "split"-       <> short 'k'-       <> value 1 -       <> showDefault-       <> help "k-split ratio training-validation")-  <*> switch-       ( long "print-pareto"-       <> help "print Pareto front instead of best found expression")-  <*> switch-       ( long "trace"-       <> help "print all evaluated expressions.")-  <*> option auto-       ( long "distribution"-       <> value Gaussian-       <> showDefault-       <> help "distribution of the data.")-  <*> option auto-       ( long "opt-iter"-       <> value 30-       <> showDefault-       <> help "number of iterations in parameter optimization.")-  <*> option auto-       ( long "opt-retries"-       <> value 1-       <> showDefault-       <> help "number of retries of parameter fitting.")-  <*> option auto-       ( long "nPop"-       <> value 100-       <> showDefault-       <> help "population size (Default: 100).")-  <*> option auto-       ( long "tournament-size"-       <> value 2-       <> showDefault-       <> help "tournament size.")-  <*> option auto-       ( long "pc"-       <> value 1.0-       <> showDefault-       <> help "probability of crossover.")-  <*> option auto-       ( long "pm"-       <> value 0.3-       <> showDefault-       <> help "probability of mutation.")-  <*> strOption-       ( long "non-terminals"-       <> value "Add,Sub,Mul,Div,PowerAbs,Recip"-       <> showDefault-       <> help "set of non-terminals to use in the search."-       )-  <*> strOption-       ( long "dump-to"-       <> value ""-       <> showDefault-       <> help "dump final e-graph to a file."-       )-  <*> strOption-       ( long "load-from"-       <> value ""-       <> showDefault-       <> help "load initial e-graph from a file."-       )-  <*> switch-       ( long "moo"-       <> help "replace the current population with the pareto front instead of replacing it with the generated children."-       )--main :: IO ()-main = do-  args <- execParser opts-  g    <- getStdGen-  dataTrain' <- loadTrainingOnly (_dataset args) True-  dataTest   <- if null (_testData args)-                  then pure dataTrain'-                  else loadTrainingOnly (_testData args) True-  let ((dataTrain, dataVal), g') = runState (splitData dataTrain' $ _split args) g-      alg = evalStateT (egraphGP dataTrain dataVal dataTest args) emptyGraph-  evalStateT alg g'-  where-    opts = Opt.info (opt <**> helper)-            ( fullDesc <> progDesc "An implementation of GP with modified crossover and mutation\-                                   \ operators designed to exploit equality saturation and e-graphs."-           <> header "GPEgg - Genetic Programming for Symbolic Regression using e-graphs." )
− apps/egraphGP/Random.hs 
@@ -1,54 +0,0 @@-module Random where --import System.Random -import Control.Monad.State.Strict-import Control.Monad-import Data.SRTree -import Data.SRTree.Eval-import Data.Massiv.Array as MA--type Rng a = StateT StdGen IO a --toss :: Rng Bool-toss = state random-{-# INLINE toss #-}--tossBiased :: Double -> Rng Bool-tossBiased p = do r <- state random -                  pure (r < p)--randomVal :: Rng Double-randomVal = state random--randomRange :: (Ord val, Random val) => (val, val) -> Rng val-randomRange rng = state (randomR rng)-{-# INLINE randomRange #-}--randomFrom :: [a] -> Rng a-randomFrom funs = do n <- randomRange (0, length funs - 1)-                     pure $ funs !! n-{-# INLINE randomFrom #-}--randomVec :: Int -> Rng PVector-randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))--randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree)-randomTree minDepth maxDepth maxSize genTerm genNonTerm grow -  | noSpaceLeft = genTerm-  | needNonTerm = genRecursion -  | otherwise   = do r <- toss-                     if r -                       then genTerm -                       else genRecursion-  where -    noSpaceLeft = maxDepth <= 1 || maxSize <= 2-    needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2--    genRecursion = do -        node <- genNonTerm-        case node of -          Uni f _    -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow -          Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow-                           r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow -                           pure . Fix  $ Bin op l r-{-# INLINE randomTree #-}
− apps/egraphGP/Util.hs 
@@ -1,258 +0,0 @@-{-# LANGUAGE  BlockArguments #-}-{-# LANGUAGE  TupleSections #-}--module Util where--import qualified Data.Map.Strict as Map-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.SRTree-import Data.SRTree.Eval-import Algorithm.SRTree.Opt-import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Build-import Algorithm.EqSat.Info--import Algorithm.SRTree.NonlinearOpt-import System.Random-import Random-import Algorithm.SRTree.Likelihoods---import Algorithm.SRTree.ModelSelection---import Algorithm.SRTree.Opt-import qualified Data.IntMap.Strict as IM-import Control.Monad.State.Strict-import Control.Monad ( when, replicateM, forM, forM_ )-import Data.Maybe ( fromJust )-import Data.List ( maximumBy )-import Data.Function ( on )-import List.Shuffle ( shuffle )-import Data.List.Split ( splitOn )-import Data.Char ( toLower )-import qualified Data.IntSet as IntSet-import Data.SRTree.Datasets-import Algorithm.EqSat.Queries--type RndEGraph a = EGraphST (StateT StdGen IO) a-type DataSet = (SRMatrix, PVector, Maybe PVector)--csvHeader :: String-csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"--io :: IO a -> RndEGraph a-io = lift . lift-{-# INLINE io #-}-rnd :: StateT StdGen IO a -> RndEGraph a-rnd = lift-{-# INLINE rnd #-}--myCost :: SRTree Int -> Int-myCost (Var _)     = 1-myCost (Const _)   = 1-myCost (Param _)   = 1-myCost (Bin _ l r) = 2 + l + r-myCost (Uni _ t)   = 3 + t--while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()-while p arg prog = do when (p arg) do arg' <- prog arg-                                      while p arg' prog--fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)-fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =-  if isNaN val -- || isNaN tr-    then (-(1/0), theta) -- infinity-    else (val, theta)-  where-    --tree          = relabelParams _tree-    nParams       = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig-    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta-    --tr            = evalF x y mYErr-    val           = evalF x_val y_val mYErr_val----{-# INLINE fitnessFun #-}--fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)-fitnessFunRep nRep nIter distribution dataTrain dataVal _tree = do-    let tree = relabelParams _tree-        nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)-    let fits = maximumBy (compare `on` fst) $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal tree) thetaOrigs-    pure fits---{-# INLINE fitnessFunRep #-}----- helper query functions--- TODO: move to egraph lib-getFitness :: EClassId -> RndEGraph (Maybe Double)-getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)-{-# INLINE getFitness #-}-getTheta :: EClassId -> RndEGraph (Maybe PVector)-getTheta c = gets (_theta . _info . (IM.! c) . _eClass)-{-# INLINE getTheta #-}-getSize :: EClassId -> RndEGraph Int-getSize c = gets (_size . _info . (IM.! c) . _eClass)-{-# INLINE getSize #-}-isSizeOf :: (Int -> Bool) -> EClass -> Bool-isSizeOf p = p . _size . _info-{-# INLINE isSizeOf #-}--getBestFitness :: RndEGraph (Maybe Double)-getBestFitness = do-    bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)-    gets (_fitness . _info . (IM.! bec) . _eClass)---- TODO: move to dataset lib-chunksOf :: Int -> [e] -> [[e]]-chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))- where-  splitter :: [e] -> ([e] -> a -> a) -> a -> a-  splitter [] _ n = n-  splitter l c n = l `c` splitter (Prelude.drop i l) c n-  build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]-  build g = g (:) []--splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)-splitData (x, y, mYErr) k = do-  if k == 1-    then pure ((x, y, mYErr), (x, y, mYErr))-    else do-      ixs' <- (state . shuffle) [0 .. sz-1]-      let ixs = chunksOf k ixs'--      let (x_tr, x_te) = getX ixs x-          (y_tr, y_te) = getY ixs y-          mY = fmap (getY ixs) mYErr-          (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)-      pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))-  where-    (MA.Sz sz) = MA.size y-    comp_x     = MA.getComp x-    comp_y     = MA.getComp y--    getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)-    getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]-                    in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]-                       , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]-                       )-    getY :: [[Int]] -> PVector -> (PVector, PVector)-    getY ixs ys  = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]-                   , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]-                   )--getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)-getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)--getX :: DataSet -> SRMatrix-getX (a, _, _) = a--getTarget :: DataSet -> PVector-getTarget (_, b, _) = b--getError :: DataSet -> Maybe PVector-getError (_, _, c) = c--loadTrainingOnly fname b = getTrain <$> loadDataset fname b--parseNonTerms :: String -> [SRTree ()]-parseNonTerms = Prelude.map toNonTerm . splitOn ","-  where-    binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]-    uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]-    toNonTerm xs' = let xs = Prelude.map toLower xs'-                    in case binTerms Map.!? xs of-                          Just op -> Bin op () ()-                          Nothing -> case uniTerms Map.!? xs of-                                          Just f -> Uni f ()-                                          Nothing -> error $ "invalid non-terminal " <> show xs---- RndEGraph utils--- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val-insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, PVector)) -> RndEGraph EClassId-insertExpr t fitFun = do-    ecId <- fromTree myCost t >>= canonical-    (f, p) <- fitFun t-    insertFitness ecId f p-    io . putStrLn $ "Best fit global: " <> show f-    pure ecId-  where powabs l r  = Fix (Bin PowerAbs l r)--updateIfNothing fitFun ec = do-      mf <- getFitness ec-      case mf of-        Nothing -> do-          t <- getBestExpr ec-          (f, p) <- fitFun t-          insertFitness ec f p-          pure True-        Just _ -> pure False--pickRndSubTree :: RndEGraph (Maybe EClassId)-pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)-                    if not (null ecIds)-                          then do rndId' <- rnd $ randomFrom ecIds-                                  rndId  <- canonical rndId'-                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)-                                  case constType of-                                    NotConst -> pure $ Just rndId-                                    _        -> pure Nothing-                          else pure Nothing--getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)--getBestExprWithSize n =-        do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical-           if (not (null ec))-            then do-              bestFit <- getFitness $ head ec-              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)-              pure [(head ec, bestFit)]-            else pure []--insertRndExpr maxSize rndTerm rndNonTerm =-      do grow <- rnd toss-         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])-         t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow-         fromTree myCost t >>= canonical--refit fitFun ec = do-  t <- getBestExpr ec-  (f, p) <- fitFun t-  insertFitness ec f p----printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-printBest fitFun printExprFun = do-      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical-      bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)-      --refit fitFun bec-      --io.print $ "should be " <> show bestFit-      printExprFun 0 bec----paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))-    where-    go :: Int -> Int -> Double -> RndEGraph ()-    go n ix f-        | n > maxSize = pure ()-        | otherwise   = do-            ecList <- getBestExprWithSize n-            if not (null ecList)-                then do let (ec, mf) = head ecList-                            improved = fromJust mf > f-                        ec' <- canonical ec-                        when improved $ refit fitFun ec' >> printExprFun ix ec'-                        go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))-                else go (n+1) ix f--evaluateUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          forM_ ec $ \c -> do-              t <- getBestExpr c-              (f, p) <- fitFun t-              insertFitness c f p--evaluateRndUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          c <- rnd . randomFrom $ ec-          t <- getBestExpr c-          (f, p) <- fitFun t-          insertFitness c f p-          pure c
− apps/egraphSearch/Main.hs 
@@ -1,325 +0,0 @@-{-# LANGUAGE  BlockArguments #-}-{-# LANGUAGE  TupleSections #-}-{-# LANGUAGE  MultiWayIf #-}-{-# LANGUAGE  OverloadedStrings #-}-{-# LANGUAGE  BangPatterns #-}-{-# LANGUAGE  TypeSynonymInstances, FlexibleInstances #-}--module Main where --import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Simplify-import Algorithm.EqSat.Build-import Algorithm.EqSat.Queries-import Algorithm.EqSat.Info-import Algorithm.EqSat.DB-import Algorithm.SRTree.Likelihoods-import Algorithm.SRTree.ModelSelection-import Control.Monad ( forM_, forM, when )-import Control.Monad.State.Strict-import qualified Data.IntMap.Strict as IM-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.Maybe ( fromJust, isNothing )-import Data.SRTree-import Data.SRTree.Print ( showExpr )-import Options.Applicative as Opt hiding (Const)-import Random-import System.Random-import Data.List ( intercalate )-import qualified Data.IntSet as IntSet-import Algorithm.EqSat (runEqSat)-import Data.Binary ( encode, decode )-import qualified Data.ByteString.Lazy as BS-import Debug.Trace-import qualified Data.HashSet as Set-import Control.Lens (over)--import Util--data Alg = OnlyRandom | BestFirst deriving (Show, Read, Eq)--egraphSearch :: DataSet -> DataSet -> DataSet -> Args -> StateT EGraph (StateT StdGen IO) ()--- terms nEvals maxSize printPareto printTrace slowIter slowRep =-egraphSearch dataTrain dataVal dataTest args = do-  if null (_loadFrom args)-    then do ecFst <- insertRndExpr (_maxSize args) rndTerm rndNonTerm2-            --ecFst <- insertBestExpr -- use only to debug-            updateIfNothing fitFun ecFst-            insertTerms-            evaluateUnevaluated fitFun-            runEqSat myCost rewritesParams 1-            cleanDB-    else (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg)-  nCls <- gets (IM.size . _eClass)-  nUnev <- gets (IntSet.size . _unevaluated . _eDB)-  let nEvs = nCls - nUnev--  while ((<(_gens args)) . snd) (10, nEvs) $-    \(radius, nEvs) ->-      do-       nCls  <- gets (IM.size . _eClass)-       --nUnev <- gets (IntSet.size . _unevaluated . _eDB)-       --let nEvs = nCls - nUnev -- WARNING: see if it affects results--       ecN <- case (_alg args) of-                    OnlyRandom -> do let ratio = fromIntegral nEvs / fromIntegral nCls-                                     b <- rnd (tossBiased ratio)-                                     ec <- if b && ratio > 0.99-                                              then insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical-                                              else do mEc <- pickRndSubTree -- evaluateRndUnevaluated fitFun >>= canonical-                                                      case mEc of-                                                           Nothing ->insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical-                                                           Just ec' -> pure ec'-                                     pure ec-                    BestFirst  -> do-                      ecsPareto <- getParetoEcsUpTo 50 (_maxSize args)-                      ecPareto     <- combineFrom ecsPareto >>= canonical-                      curFitPareto <- getFitness ecPareto--                      if isNothing curFitPareto-                        then pure ecPareto-                        else do ecsBest    <- getTopFitEClassThat 100 (isSizeOf (<(_maxSize args)))-                                ecBest     <- combineFrom ecsBest >>= canonical-                                curFitBest <- getFitness ecBest-                                if isNothing curFitBest-                                  then pure ecBest-                                  else do ee <- pickRndSubTree-                                          case ee of-                                            Nothing -> insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical-                                            Just c  -> pure c--       -- when upd $-       ecN' <- canonical ecN-       upd <- updateIfNothing fitFun ecN'-       when upd $ runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged---       when (upd && (_trace args))-         do-            ecN'' <- canonical ecN'-            _tree <- getBestExpr ecN''-            fi <- negate . fromJust <$> getFitness ecN''-            theta <- fromJust <$> getTheta ecN''-            let thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)-            io . putStrLn $ showExpr _tree <> "," <> thetaStr <> "," <> show fi-            pure ()-       let radius' = if (not upd) then (max 10 $ min (500 `div` (_maxSize args)) (radius+1)) else (max 10 $ radius-1)-           nEvs'    = nEvs + if upd then 1 else 0-       pure (radius', nEvs')-  io $ putStrLn csvHeader-  if (_printPareto args)-    then paretoFront (_maxSize args) printExpr-    else printBest printExpr-  when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode )-  where-    fitFun = fitnessFunRep (_optRepeat args) (_optIter args) (_distribution args) dataTrain dataVal--    refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList-                      modify' $ over (eDB . refits) (const Set.empty)-                      forM_ ids $ \ec -> do t <- getBestExpr ec-                                            (f, p) <- fitFun t-                                            insertFitness ec f p--    combineFrom [] = pure 0 -- this is the first terminal and it will always be already evaluated-    combineFrom ecs = do-        nt  <- rnd rndNonTerm-        p1  <- rnd (randomFrom ecs)-        p2  <- rnd (randomFrom ecs)-        l1  <- rnd (randomFrom [2..(_maxSize args)-2]) -- sz 10: [2..8]--        e1  <- randomChildFrom p1 l1 >>= canonical-        ml  <- gets (_size . _info . (IM.! e1) . _eClass)-        l2  <- rnd (randomFrom [1..((_maxSize args) - ml - 1)]) -- maxSize - maxSize + 2 - 2= 0 -- sz 10: [1..7] (2) / [1..1] (8)-        e2  <- randomChildFrom p2 l2 >>= canonical-        case nt of-          Uni Id ()    -> canonical e1-          Uni f ()     -> add myCost (Uni f e1) >>= canonical-          Bin op () () -> do b <- rnd toss-                             if b-                              then add myCost (Bin op e1 e2) >>= canonical-                              else add myCost (Bin op e2 e1) >>= canonical-          _            -> canonical e1 -- it is a terminal, should it happen?--    randomChildFrom ec' maxL = do-      p <- rnd toss -- whether to go deeper or return this level-      ec <- canonical ec'-      l <- gets (_size . _info . (IM.! ec) . _eClass )--      if p || l > maxL-          then do --enodes <- gets (_eNodes . (IM.! ec) . _eClass)-                  enode  <- gets (_best . _info . (IM.! ec) . _eClass) -- we should return the best otherwise we may build larger exprs-                  case enode of-                      Uni _ eci     -> randomChildFrom eci maxL-                      Bin _ ecl ecr -> do coin <- rnd toss-                                          if coin-                                            then randomChildFrom ecl maxL-                                            else randomChildFrom ecr maxL-                      _ -> pure ec -- this shouldn't happen unless maxL==0-          else pure ec--    nonTerms   = parseNonTerms (_nonterminals args)-    --[ Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () (), Bin PowerAbs () (),  Uni Recip ()]-    (Sz2 _ nFeats) = MA.size (getX dataTrain)-    terms          = if _distribution args == ROXY-                          then [var 0, param 0]-                          else [var ix | ix <- [0 .. nFeats-1]] <> [param 0]--    rndTerm    = Random.randomFrom terms-    rndNonTerm = Random.randomFrom $ (Uni Id ()) : nonTerms-    rndNonTerm2 = Random.randomFrom nonTerms--    insertTerms =-        forM terms $ \t -> do fromTree myCost t >>= canonical--    printExpr :: Int -> EClassId -> RndEGraph ()-    printExpr ix ec = do -        theta' <- gets (fromJust . _theta . _info . (IM.! ec) . _eClass)-        bestExpr <- getBestExpr ec-        let nParams = countParams bestExpr-            (MA.Sz nTheta)  = MA.size theta'-        (_, theta) <- if (nParams /= nTheta)-                        then fitFun bestExpr-                        else pure (1.0, theta')--        let (x, y, mYErr) = dataTrain-            (x_val, y_val, mYErr_val) = dataVal-            (x_te, y_te, mYErr_te) = dataTest-            distribution = _distribution args-            best'     = relabelParams bestExpr-            expr      = paramsToConst (MA.toList theta) best'-            mse_train = mse x y best' theta-            mse_val   = mse x_val y_val best' theta-            mse_te    = mse x_te y_te best' theta-            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  = mdl distribution mYErr x y theta best'-            mdl_val    = mdl distribution mYErr_val x_val y_val theta best'-            mdl_te     = mdl distribution mYErr_te x_te y_te theta best'-            vals       = intercalate ","-                       $ Prelude.map show [mse_train, mse_val, mse_te-                                          , r2_train, r2_val, r2_te-                                          , nll_train, nll_val, nll_te-                                          , mdl_train, mdl_val, mdl_te]-            thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)-        io . putStrLn $ show ix <> "," <> showExpr expr <> ","-                      <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)-                      <> "," <> vals---data Args = Args-  { _dataset      :: String,-    _testData     :: String,-    _gens         :: Int,-    _alg          :: Alg,-    _maxSize      :: Int,-    _split        :: Int,-    _printPareto  :: Bool,-    _trace        :: Bool,-    _distribution :: Distribution,-    _optIter      :: Int,-    _optRepeat    :: Int,-    _nonterminals :: String,-    _dumpTo       :: String,-    _loadFrom     :: String-  }-  deriving (Show)---- parser of command line arguments-opt :: Parser Args-opt = Args-   <$> strOption-       ( long "dataset"-       <> short 'd'-       <> metavar "INPUT-FILE"-       <> help "CSV dataset." )-  <*> strOption-       ( long "test"-       <> short 't'-       <> value ""-       <> showDefault-       <> help "test data")-   <*> option auto-      ( long "generations"-      <> short 'g'-      <> metavar "GENS"-      <> showDefault-      <> value 100-      <> help "Number of generations." )-   <*> option auto-       ( long "algorithm"-       <> short 'a'-       <> metavar "ALG"-       <> help "Algorithm." )-  <*> option auto-       ( long "maxSize"-       <> short 's'-       <> help "max-size." )-  <*> option auto-       ( long "split"-       <> short 'k'-       <> value 1 -       <> showDefault-       <> help "k-split ratio training-validation")-  <*> switch-       ( long "print-pareto"-       <> help "print Pareto front instead of best found expression")-  <*> switch-       ( long "trace"-       <> help "print all evaluated expressions.")-  <*> option auto-       ( long "distribution"-       <> value Gaussian-       <> showDefault-       <> help "distribution of the data.")-  <*> option auto-       ( long "opt-iter"-       <> value 30-       <> showDefault-       <> help "number of iterations in parameter optimization.")-  <*> option auto-       ( long "opt-retries"-       <> value 1-       <> showDefault-       <> help "number of retries of parameter fitting.")-  <*> strOption-       ( long "non-terminals"-       <> value "Add,Sub,Mul,Div,PowerAbs,Recip"-       <> showDefault-       <> help "set of non-terminals to use in the search."-       )-  <*> strOption-       ( long "dump-to"-       <> value ""-       <> showDefault-       <> help "dump final e-graph to a file."-       )-  <*> strOption-       ( long "load-from"-       <> value ""-       <> showDefault-       <> help "load initial e-graph from a file."-       )--main :: IO ()-main = do-  args <- execParser opts-  g    <- getStdGen-  dataTrain' <- loadTrainingOnly (_dataset args) True-  dataTest   <- if null (_testData args)-                  then pure dataTrain'-                  else loadTrainingOnly (_testData args) True-  let ((dataTrain, dataVal), g') = runState (splitData dataTrain' $ _split args) g-      alg = evalStateT (egraphSearch dataTrain dataVal dataTest args) emptyGraph-  evalStateT alg g'--  where-    opts = Opt.info (opt <**> helper)-            ( fullDesc <> progDesc "Symbolic Regression search algorithm\-                                   \ exploiting the potentials of equality saturation\-                                   \ and e-graphs."-           <> header "SymREgg - symbolic regression with e-graphs."-            )
− apps/egraphSearch/Random.hs 
@@ -1,54 +0,0 @@-module Random where --import System.Random -import Control.Monad.State.Strict-import Control.Monad-import Data.SRTree -import Data.SRTree.Eval-import Data.Massiv.Array as MA--type Rng a = StateT StdGen IO a --toss :: Rng Bool-toss = state random-{-# INLINE toss #-}--tossBiased :: Double -> Rng Bool-tossBiased p = do r <- state random -                  pure (r < p)--randomVal :: Rng Double-randomVal = state random--randomRange :: (Ord val, Random val) => (val, val) -> Rng val-randomRange rng = state (randomR rng)-{-# INLINE randomRange #-}--randomFrom :: [a] -> Rng a-randomFrom funs = do n <- randomRange (0, length funs - 1)-                     pure $ funs !! n-{-# INLINE randomFrom #-}--randomVec :: Int -> Rng PVector-randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))--randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree)-randomTree minDepth maxDepth maxSize genTerm genNonTerm grow -  | noSpaceLeft = genTerm-  | needNonTerm = genRecursion -  | otherwise   = do r <- toss-                     if r -                       then genTerm -                       else genRecursion-  where -    noSpaceLeft = maxDepth <= 1 || maxSize <= 2-    needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2--    genRecursion = do -        node <- genNonTerm-        case node of -          Uni f _    -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow -          Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow-                           r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow -                           pure . Fix  $ Bin op l r-{-# INLINE randomTree #-}
− apps/egraphSearch/Util.hs 
@@ -1,254 +0,0 @@-{-# LANGUAGE  BlockArguments #-}-{-# LANGUAGE  TupleSections #-}--module Util where--import qualified Data.Map.Strict as Map-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.SRTree-import Data.SRTree.Eval-import Algorithm.SRTree.Opt-import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Build-import Algorithm.EqSat.Info--import Algorithm.SRTree.NonlinearOpt-import System.Random-import Random-import Algorithm.SRTree.Likelihoods---import Algorithm.SRTree.ModelSelection---import Algorithm.SRTree.Opt-import qualified Data.IntMap.Strict as IM-import Control.Monad.State.Strict-import Control.Monad ( when, replicateM, forM, forM_ )-import Data.Maybe ( fromJust )-import Data.List ( maximumBy )-import Data.Function ( on )-import List.Shuffle ( shuffle )-import Data.List.Split ( splitOn )-import Data.Char ( toLower )-import qualified Data.IntSet as IntSet-import Data.SRTree.Datasets-import Algorithm.EqSat.Queries--type RndEGraph a = EGraphST (StateT StdGen IO) a-type DataSet = (SRMatrix, PVector, Maybe PVector)--csvHeader :: String-csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"--io :: IO a -> RndEGraph a-io = lift . lift-{-# INLINE io #-}-rnd :: StateT StdGen IO a -> RndEGraph a-rnd = lift-{-# INLINE rnd #-}--myCost :: SRTree Int -> Int-myCost (Var _)     = 1-myCost (Const _)   = 1-myCost (Param _)   = 1-myCost (Bin _ l r) = 2 + l + r-myCost (Uni _ t)   = 3 + t--while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()-while p arg prog = do when (p arg) do arg' <- prog arg-                                      while p arg' prog--fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)-fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) _tree thetaOrig =-  if isNaN val || isNaN tr-    then (-(1/0), theta) -- infinity-    else (val, theta) -- (min tr val, theta)-  where-    tree          = relabelParams _tree-    nParams       = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig-    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta-    tr            = evalF x y mYErr-    val           = evalF x_val y_val mYErr_val--{-# INLINE fitnessFun #-}--fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)-fitnessFunRep nRep nIter distribution dataTrain dataVal _tree = do-    let tree = relabelParams _tree-        nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0-    thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)-    let fits = Prelude.map (fitnessFun nIter distribution dataTrain dataVal _tree) thetaOrigs-    pure (maximumBy (compare `on` fst) fits)-{-# INLINE fitnessFunRep #-}----fitnessMV :: Int -> Int -> Distribution -> [DataSet] -> [DataSet] -> Fix SRTree -> RndEGraph (Double, [PVector])---fitnessMV nRep nIter distribution dataTrains dataVals _tree = do---  response <- forM (zip dataTrains dataVals) $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv _tree---  pure (minimum (map fst response), map snd response)---- helper query functions--- TODO: move to egraph lib-getFitness :: EClassId -> RndEGraph (Maybe Double)-getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)-{-# INLINE getFitness #-}-getTheta :: EClassId -> RndEGraph (Maybe PVector)-getTheta c = gets (_theta . _info . (IM.! c) . _eClass)-{-# INLINE getTheta #-}-getSize :: EClassId -> RndEGraph Int-getSize c = gets (_size . _info . (IM.! c) . _eClass)-{-# INLINE getSize #-}-isSizeOf :: (Int -> Bool) -> EClass -> Bool-isSizeOf p = p . _size . _info-{-# INLINE isSizeOf #-}--getBestFitness :: RndEGraph (Maybe Double)-getBestFitness = do-    bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)-    gets (_fitness . _info . (IM.! bec) . _eClass)---- TODO: move to dataset lib-chunksOf :: Int -> [e] -> [[e]]-chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))- where-  splitter :: [e] -> ([e] -> a -> a) -> a -> a-  splitter [] _ n = n-  splitter l c n = l `c` splitter (Prelude.drop i l) c n-  build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]-  build g = g (:) []--splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)-splitData (x, y, mYErr) k = do-  if k == 1-    then pure ((x, y, mYErr), (x, y, mYErr))-    else do-      ixs' <- (state . shuffle) [0 .. sz-1]-      let ixs = chunksOf k ixs'--      let (x_tr, x_te) = getX ixs x-          (y_tr, y_te) = getY ixs y-          mY = fmap (getY ixs) mYErr-          (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)-      pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))-  where-    (MA.Sz sz) = MA.size y-    comp_x     = MA.getComp x-    comp_y     = MA.getComp y--    getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)-    getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]-                    in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]-                       , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]-                       )-    getY :: [[Int]] -> PVector -> (PVector, PVector)-    getY ixs ys  = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]-                   , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]-                   )--getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)-getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)--getX :: DataSet -> SRMatrix-getX (a, _, _) = a--getTarget :: DataSet -> PVector-getTarget (_, b, _) = b--getError :: DataSet -> Maybe PVector-getError (_, _, c) = c--loadTrainingOnly fname b = getTrain <$> loadDataset fname b--parseNonTerms :: String -> [SRTree ()]-parseNonTerms = Prelude.map toNonTerm . splitOn ","-  where-    binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]-    uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]-    toNonTerm xs' = let xs = Prelude.map toLower xs'-                    in case binTerms Map.!? xs of-                          Just op -> Bin op () ()-                          Nothing -> case uniTerms Map.!? xs of-                                          Just f -> Uni f ()-                                          Nothing -> error $ "invalid non-terminal " <> show xs---- RndEGraph utils--- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val-insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, PVector)) -> RndEGraph EClassId-insertExpr t fitFun = do-    ecId <- fromTree myCost t >>= canonical-    (f, p) <- fitFun t-    insertFitness ecId f p-    io . putStrLn $ "Best fit global: " <> show f-    pure ecId-  where powabs l r  = Fix (Bin PowerAbs l r)--updateIfNothing fitFun ec = do-      mf <- getFitness ec-      case mf of-        Nothing -> do-          t <- getBestExpr ec-          (f, p) <- fitFun t-          insertFitness ec f p-          pure True-        Just _ -> pure False--pickRndSubTree :: RndEGraph (Maybe EClassId)-pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)-                    if not (null ecIds)-                          then do rndId' <- rnd $ randomFrom ecIds-                                  rndId  <- canonical rndId'-                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)-                                  case constType of-                                    NotConst -> pure $ Just rndId-                                    _        -> pure Nothing-                          else pure Nothing--getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)--getBestExprWithSize n =-        do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical-           if (not (null ec))-            then do-              bestFit <- getFitness $ head ec-              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)-              (:[]) . (,bestP) . (,bestFit) . (,ec) <$> getBestExpr (head ec)-            else pure []--insertRndExpr maxSize rndTerm rndNonTerm =-      do grow <- rnd toss-         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])-         t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow-         fromTree myCost t >>= canonical--printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-printBest printExprFun = do-      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical-      printExprFun 0 bec--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()-paretoFront maxSize printExprFun = go 1 0 (-(1.0/0.0))-    where-    go :: Int -> Int -> Double -> RndEGraph ()-    go n ix f-        | n > maxSize = pure ()-        | otherwise   = do-            ecList <- getBestExprWithSize n-            if not (null ecList)-                then do let (((_, ec), mf), _) = head ecList-                            improved = fromJust mf > f-                        ec' <- traverse canonical ec-                        when improved $ printExprFun ix (head ec')-                        go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))-                else go (n+1) ix f--evaluateUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          forM_ ec $ \c -> do-              t <- getBestExpr c-              (f, p) <- fitFun t-              insertFitness c f p--evaluateRndUnevaluated fitFun = do-          ec <- gets (IntSet.toList . _unevaluated . _eDB)-          c <- rnd . randomFrom $ ec-          t <- getBestExpr c-          (f, p) <- fitFun t-          insertFitness c f p-          pure c
− apps/eqsatrepr/Main.hs 
@@ -1,302 +0,0 @@-{-# LANGUAGE OverloadedStrings #-}-module Main where--import Data.SRTree -import Algorithm.EqSat.Egraph-import Data.SRTree.Print -import qualified Data.Map as Map-import qualified Data.IntMap as IM-import Control.Monad.State.Strict-import System.Random-import Data.SRTree.Recursion ( cata )-import Control.Monad-import Control.Monad.Reader-import qualified Data.SRTree.Random as RT-import Data.List ( nub )-import Algorithm.EqSat.DB-import Algorithm.EqSat.Info-import Algorithm.EqSat.Build-import Algorithm.EqSat.Queries-import Algorithm.EqSat--isConstPt :: Pattern -> Map.Map ClassOrVar ClassOrVar -> EGraph -> Bool-isConstPt (VarPat c) subst eg =-    let cid = getInt $ subst Map.! (Right $ fromEnum c)-    in case (_consts . _info) (_eClass eg IM.! cid) of-         ConstVal x -> True-         _ -> False-isConstPt _ _ _ = False--notZero (VarPat c) subst eg =-  let cid = getInt $ subst Map.! (Right $ fromEnum c)-   in case (_consts . _info) (_eClass eg IM.! cid) of-         ConstVal x -> x /= 0-         _ -> True-notZero _ _ _ = True--rewriteBasic =-    [-      "x" * "x" :=> "x" ** 2-    , "x" * "y" :=> "y" * "x"-    , "x" + "y" :=> "y" + "x"-    , ("x" ** "y") * "x" :=> "x" ** ("y" + 1) :| isConstPt "y"-    -- , ("x" * "y") / "x" :=> "y"-    , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z")-    , ("x" + "y") + "z" :=> "x" + ("y" + "z")-    , ("x" + "y") - "z" :=> "x" + ("y" - "z")-    , ("x" * "y") * "z" :=> "x" * ("y" * "z")-    , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")-    , "x" - ("y" + "z") :=> ("x" - "y") - "z"-    , "x" - ("y" - "z") :=> ("x" - "y") + "z"-    , ("x" * "y") / "z" :=> ("x" / "z") * "y"-    , (("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y") :| isConstPt "w") :| isConstPt "z"-    , ((("x" * "y") + ("z" * "w")) :=> "x" * ("y" + ("z" / "x") * "w") :| isConstPt "x") :| isConstPt "z"-    , ((("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x") :| isConstPt "z"-    , ((("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x") :| isConstPt "z"-    ]--rewritesFun =-    [-      log (sqrt "x") :=> 0.5 * log "x"-    , log (exp "x")  :=> "x"-    , exp (log "x")  :=> "x"-    , "x" ** (1/2)   :=> sqrt "x"-    ,  log ("x" * "y") :=> log "x" + log "y"-    , log ("x" / "y") :=> log "x" - log "y"-    , log ("x" ** "y") :=> "y" * log "x"-    , sqrt ("y" * "x") :=> sqrt "y" * sqrt "x"-    , sqrt ("y" / "x") :=> sqrt "y" / sqrt "x"-    , abs ("x" * "y") :=> abs "x" * abs "y"-    ,  sqrt ("z" * ("x" - "y")) :=> sqrt (negate "z") * sqrt ("y" - "x")-    , sqrt ("z" * ("x" + "y")) :=> sqrt "z" * sqrt ("x" + "y")-    ]---- Rules that reduces redundant parameters-constReduction =-    [-      0 + "x" :=> "x"-    , "x" - 0 :=> "x"-    , 1 * "x" :=> "x"-    , 0 * "x" :=> 0-    , 0 / "x" :=> 0 :| notZero "x"-    , "x" - "x" :=> 0-    , "x" / "x" :=> 1 :| notZero "x"-    , "x" ** 1 :=> "x"-    , 0 ** "x" :=> 0-    , 1 ** "x" :=> 1-    , "x" * (1 / "x") :=> 1-    , 0 - "x" :=> negate "x"-    , "x" + negate "y" :=> "x" - "y"-    , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x"-    ]---x0 = var 0-x1 = var 1-x2 = var 2-x3 = var 3-x4 = var 4-x5 = var 5-x6 = var 6-x7 = var 7-x8 = var 8--trees :: [Fix SRTree]-trees = [  (4.059e-3 + (0.988153 * (((1.923901 * x1) * ((-1.228652 * x0) * (-0.278891 * x2))) * ((((((-0.35119 * x5) + (-0.354523 * x3)) - (-0.369148 * x6)) + ((0.342012 * x4) + (2.054e-2 * x2))) - ((0.349297 * x7) - (0.336081 * x8)))))))-        , ((14.316036 * (((((0.975231 * x4)) * (1.259663 * x3)) * (0.314221 * x0)) * (0.178249 * x2))))-        , (1.2 * (3.4 * x1 * 4.2 * x0) / ((3.2 * x2) * ((1.1 * x3) * (3.5 * x4))))-        , ((1.002563 * (((0.428416 * x1) * (2.554566 * x0)) / (((2.53743 * x2) * (2.327917 * x3)) * (2.320736 * x3)))))-        , (2.82238 + (3.092415 * (sin(log(abs(0.0))) * ((-0.162842 * x2) - (0.116404 * x1)))))-        , log(0.0) * ((1.2 * x2) - (0.116404 * x1))-        , ((x0 - x0) * x0)-        , (1 + 1) - 1-        , (x0 + x0) - x0-        , (x0/x0 + 1) - 1-        , (x0 * x0) / x0-        , sin(log(0.0))-        , -1 * exp(log(abs(-1.3 * (x1 - 1.2 * x2))))-        , -1 * exp(log(abs((1.3 * x1 + 1.56 * x2))))-        , -1 * exp(log(abs((-1.3 * x1 + 1.56 * x2))))-        , -1 * exp(log(abs(((0.256 * x3) + (-0.2561 * x2)))))-        , log(abs(-1.199026) * abs((x2 + (1.191617 * x3))))-        , log(abs((1.199026 * x2) + (-1.191617 * x3)))-        , 1 * x0-        , x0 * 1-        , x0 + x1-        , x1 + x0-        , x0 + sin(x1)-        , x0 * (1 + x1)-        , (1 + x1) * x0-        , x0 + x0 * x1-        , (x0 + x1) + 2-        , x0 + (x1 + 2)-        , x0 + (2 + x1)-        , log(abs(0)) + x0-        , abs(((1.3 * x1) + (-1.56 * x2))) * (-1.0)-        ]---myCost :: SRTree Int -> Int-myCost (Var _) = 1-myCost (Const _) = 1-myCost (Param _) = 1-myCost (Bin op l r) = 2 + l + r-myCost (Uni _ t) = 3 + t--rewrites = rewriteBasic <> constReduction <> rewritesFun--testEqSat :: Fix SRTree -> IO ()-testEqSat t = do-    let e = eqSat t rewrites myCost 30 `evalState` emptyGraph-    putStr $ (showExpr t) <> " == " <> (showExpr e) <> "\n"--testEqSats :: IO ()-testEqSats = mapM_ testEqSat trees----initialPop :: HyperParams -> Rng [Fix SRTree]-initialPop hyperparams = do-   let depths = [3 .. _maxDepth hyperparams]-   pop <- forM depths $ \md ->-           do let m = _popSize hyperparams `div` (_maxDepth hyperparams - 3 + 1)-                  g = take m $ cycle [True, False]-              mapM (randomTree hyperparams{ _maxDepth = md}) g-   pure (concat pop)-{-# INLINE initialPop #-}--data Method = Grow | Full--type Rng a = StateT StdGen IO a-type GenUni = Fix SRTree -> Fix SRTree-type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree--toss :: Rng Bool-toss = state random-{-# INLINE toss #-}--randomRange :: (Ord val, Random val) => (val, val) -> Rng val-randomRange rng = state (randomR rng)-{-# INLINE randomRange #-}--randomFrom :: [a] -> Rng a-randomFrom funs = do n <- randomRange (0, length funs - 1)-                     pure $ funs !! n-{-# INLINE randomFrom #-}--countNodes' :: Fix SRTree -> Int-countNodes' = cata alg-  where-    alg (Var _)     = 1-    alg (Param _)   = 1-    alg (Const _)   = 0-    alg (Bin _ l r) = 1 + l + r-    alg (Uni Abs t) = t-    alg (Uni _ t)   = 1 + t-{-# INLINE countNodes' #-}---randomTree :: HyperParams -> Bool -> Rng (Fix SRTree)-randomTree hp grow-  | depth <= 1 || size <= 2 = randomFrom term-  | (min_depth >= 0 || (depth > 2 && not grow)) && size > 2 = genNonTerm-  | otherwise = genTermOrNon-  where-    min_depth = _minDepth hp-    depth     = _maxDepth hp-    size      = _maxSize hp-    term      = _term hp-    nonterm   = _nonterm hp--    genNonTerm =-       do et <- randomFrom nonterm-          case et of-            Left uniT -> uniT <$> randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow-            Right binT -> do l <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow-                             r <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1 - countNodes' l} grow-                             pure (binT l r)-    genTermOrNon = do r <- toss-                      if r-                        then randomFrom term-                        else genNonTerm-{-# INLINE randomTree #-}--data HyperParams =-    HP { _minDepth  :: Int-       , _maxDepth  :: Int-       , _maxSize   :: Int-       , _popSize   :: Int-       , _tournSize :: Int-       , _pc        :: Double-       , _pm        :: Double-       , _term      :: [Fix SRTree]-       , _nonterm   :: [Either GenUni GenBin]-       }---countSubTrees = do ecs <- gets (IM.keys . _eClass) -                   subs <- mapM (\ec -> getAllExpressionsFrom ec >>= pure . length) ecs -                   pure $ sum subs -countRootTrees rs = do subs <- mapM (\ec -> getAllExpressionsFrom ec >>= pure . length) rs-                       pure $ sum subs--terms = [var 0, var 1, var 2, param 0, param 1, param 2, param 3]-nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> abs l ** r), Left (1/)]--calcRedundancy :: Int -> IO ()-calcRedundancy nPop = do-    let hp = HP 2 4 10 nPop 2 1.0 0.25 terms nonterms-        p  = RT.P [0, 1, 2, 3, 4, 5] (0, 3) (1, 3) [Log]-    g <- getStdGen-    pop <- (`evalStateT` g)  <$> replicateM nPop $ runReaderT (RT.randomTree 10) p-    let nSubsSingle = sum $ map (\p -> (fromTrees myCost [p] >> countSubTrees) `evalState` emptyGraph) pop -        myEqPop = do rs <- fromTrees myCost pop-                     let rsN = nub rs -                     cnt <- countSubTrees-                     pure (cnt, rsN)-        (nSubs, rsN) = myEqPop `evalState` emptyGraph -    putStr "Ratio of subtrees: "-    putStrLn $ show nSubsSingle <> "/" <> show nSubs <> " = " <> show (fromIntegral nSubsSingle / fromIntegral nSubs)-    let nSubsR = sum $ map (\p -> (fromTree myCost p >>= \r -> countRootTrees [r]) `evalState` emptyGraph) pop-        nSubsSingleR = (fromTrees myCost pop >>= countRootTrees) `evalState` emptyGraph-    putStr "Ratio of rooted trees: "-    putStrLn $ show nSubsSingleR <> "/" <> show nSubsR <> " = " <> show (fromIntegral nSubsSingleR / fromIntegral nSubsR)--main :: IO ()-main = do -    let t1 = var 0 + 12.0-        t2 = 3.2 * var 0-        t3 = 3.2 * var 0 / (var 0 + 12.0)-        t4 = var 0 + sin (var 0)-        t5 = 1.5 + exp 5.2-        egraphRun :: EGraphST IO ()-        egraphRun = do v <- fromTrees myCost [t3,t1,t2,t4]-                       roots <- findRootClasses-                       ecId  <- gets ((Map.! (Var 0)) . _eNodeToEClass)-                       calculateHeights -                       h <- gets (map _height . IM.elems . _eClass)-                       v <- gets (map (_consts . _info) . IM.elems . _eClass)-                       c <- gets (map (_cost . _info) . IM.elems . _eClass)-                       parents <- gets (_parents . (IM.! ecId) . _eClass)-                       exprs <- mapM getExpressionFrom roots-                       exprs' <- gets (IM.keys . _eClass) >>= mapM getExpressionFrom --                       lift $ do putStr "Parents of x0: "-                                 print parents -                                 putStrLn "\nexpressions from root: "-                                 mapM_ (putStrLn . showExpr) exprs-                                 putStrLn "\nexpressions from each e-class: "-                                 mapM_ (putStrLn . showExpr) exprs'-                                 putStrLn "heights: "-                                 mapM_ print h -- (print . _height) (IM.elems $ _eClass eg')-                                 putStrLn "values: "-                                 mapM_ print v -- (print . _consts . _info) (IM.elems $ _eClass eg')-                                 putStrLn "costs: "-                                 mapM_ print c -- (print . _cost . _info) (IM.elems $ _eClass eg')-        nPop = 10000-        hp = HP 3 7 100 nPop 2 1.0 0.25 terms nonterms-        p  = RT.P [0] (-3, 3) (-3, 3) []-    egraphRun `evalStateT` emptyGraph-    g <- getStdGen-    pop <- evalStateT (initialPop hp) g-    mapM_ (\nP -> putStr "pop " >> print nP >> calcRedundancy nP >> putStrLn "") [100, 200, 500, 1000, 5000, 10000, 20000, 100000]
apps/rEGGression/Commands.hs view
@@ -12,7 +12,7 @@ import qualified Data.IntMap.Strict as IntMap import qualified Data.IntSet as IntSet import Control.Monad.State.Strict-import Control.Monad ( forM_ )+import Control.Monad ( forM_, filterM ) import Data.Char ( toUpper ) import qualified Data.Map as Map import qualified Data.HashSet as Set@@ -44,12 +44,12 @@ import qualified Data.ByteString.Lazy as BS  import Util-+import Debug.Trace  -- * Parsing  -- top 5 by fitness|mdl [less than 5 params, less than 10 nodes] data Command  = Top Int Filter Criteria PatStr-              | Distribution FilterDist (Maybe Limit)+              | Distribution FilterDist (Maybe Limit) CriteriaDist Int Int               -- below these will not be a parsable command               | Report EClassId ArgOpt               | Optimize EClassId Int ArgOpt@@ -64,9 +64,10 @@ type Filter = EClass -> Bool -- pattern? type FilterDist = Int -> Bool data Criteria = ByFitness | ByDL deriving Eq+data CriteriaDist = ByCount | ByAvgFit data Limit = Limit Int Bool deriving Show data PatStr = PatStr String Bool | AntiPatStr String Bool | NoPat-type ArgOpt = (Distribution, DataSet, DataSet)+type ArgOpt = (Distribution, [DataSet], [DataSet])  -- top 10 with <=10|=10 size with <=4 parameters by fitness|dl matching pat -- report id@@ -99,8 +100,28 @@                                else filters'                stripSp                limit   <- listToMaybe <$> many' parseLimit-               pure $ Distribution (getAll . mconcat filters) limit+               stripSp+               by'     <- listToMaybe <$> many' parseCriteriaDL+               stripSp +               least' <- listToMaybe <$> many' parseLeast+               stripSp +               top'   <- listToMaybe <$> many' parseTopDist+               let by = case by' of +                          Nothing -> ByCount+                          Just b  -> b+                   least = case least' of +                             Nothing -> 1 +                             Just l  -> l +                   top   = case top' of +                             Nothing -> 1000+                             Just t  -> t+               pure $ Distribution (getAll . mconcat filters) limit by least top  +parseLeast = stringCI "with at least " >> decimal +parseTopDist = stringCI "from top " >> decimal +parseCriteriaDL = (stringCI "by count" >> pure ByCount)+              <|> (stringCI "by fitness" >> pure ByAvgFit)+ parseFilter = do stringCI "with"                  stripSp                  field <- parseSz <|> parseCost <|> parseParams@@ -120,8 +141,8 @@ parseCost = stringCI "cost" >> pure (_cost . _info) parseParams = stringCI "parameters" >> pure (mbLen . _theta . _info)    where-      mbLen Nothing = 0-      mbLen (Just ps) = MA.unSz $ MA.size ps+      mbLen [] = 0+      mbLen ps = MA.unSz $ MA.size $ Prelude.head ps parseCmp = do op <- parseLEQ <|> parseLT <|> parseEQ <|> parseGEQ <|> parseGT               stripSp               n <- decimal@@ -172,7 +193,7 @@ run (Top n filters criteria withPat) = do    let (pat', getFun, isParents) =           case withPat of-            PatStr p parent     -> (p, if criteria == ByFitness then getTopFitEClassIn    else getTopDLEClassIn, parent)+            PatStr p parent     -> (p, if criteria == ByFitness then getTopFitEClassIn else getTopDLEClassIn, parent)             AntiPatStr p parent -> (p, if criteria == ByFitness then getTopFitEClassNotIn else getTopDLEClassNotIn, parent)     let etree = parsePat $ B.pack pat'@@ -181,29 +202,35 @@      Right pat -> do         ecs' <- egraph $ (Prelude.map fromLeft . Prelude.filter isLeft . Prelude.map snd) <$> match pat         ecs  <- egraph $ Prelude.mapM canonical ecs'-                           >>= getParents isParents+                          >>= removeNotTrivial (lenPat pat)+                          >>= getParents isParents filters+         ids  <- egraph $ getFun n filters ecs         printSimpleMultiExprs (reverse $ nub ids) -run (Distribution pSz mLimit) = do-  ee <- egraph $ IntSet.toList . IntSet.fromList <$> getAllEvaluatedEClasses+run (Distribution pSz mLimit by least top) = do+  ee <- egraph $ IntSet.toList . IntSet.fromList <$> getTopFitEClassThat top (const True) -- getAllEvaluatedEClasses   allPats <- egraph $ getAllPatternsFrom pSz Map.empty ee   let (n, isAsc) = case mLimit of                      Nothing -> (Map.size allPats, True)                      Just (Limit sz asc) -> (sz, asc)+      predCount = (if isAsc then fst else negate . fst) . snd+      predAvgFit = (if isAsc then snd else negate . snd) . snd   printMultiCounts (Prelude.take n-                   $ sortOn (if isAsc then snd else negate . snd)+                   $ case by of +                       ByCount -> sortOn predCount+                       ByAvgFit -> sortOn predAvgFit                    $ Map.toList-                   $ Map.filterWithKey (\k v -> k /= VarPat 'A' && pSz (lenPat k))+                   $ Map.filterWithKey (\k (v,_) -> v >= least && k /= VarPat 'A' && pSz (lenPat k))                    allPats)  run (Report eid (dist, trainData, testData)) = egraph $ printExpr trainData testData dist eid -run (Optimize eid nIters (dist, trainData@(x, y, mYErr), testData)) = do -- dist trainData testData+run (Optimize eid nIters (dist, trainDatas, testData)) = do -- dist trainData testData    t <- egraph $ relabelParams <$> getBestExpr eid-   (f, theta) <- egraph $ fitnessFunRep nIters dist trainData t-   egraph $ insertFitness eid f theta-   let mdl_train  = mdl dist mYErr x y theta t+   (f, thetas) <- egraph $ fitnessMV nIters dist trainDatas t+   egraph $ insertFitness eid f thetas+   let mdl_train  = Prelude.maximum $ Prelude.map (\(theta, (x, y, mYErr)) -> mdl dist mYErr x y theta t) $ Prelude.zip thetas trainDatas    egraph $ insertDL eid mdl_train    printSimpleMultiExprs [eid] @@ -217,7 +244,7 @@ run (Subtrees eid) = do    isValid <- egraph $ gets ((IntMap.member eid) . _eClass)    if isValid-     then do ids <- egraph $ getAllChildEClasses eid+     then do ids <- egraph $ getAllChildBestEClasses eid              printSimpleMultiExprs ids      else io.putStrLn $ "Invalid id." @@ -265,7 +292,8 @@                            let (tree, ps) = if convertParam then floatConstsToParam tree' else (tree', theta)                                theta      = if convertParam then if dist==MSE then ps <> params else ps else params                            eid <- fromTree myCost tree >>= canonical-                           insertFitness eid f $ MA.fromList MA.Seq theta+                           -- TODO: how to import MvSR?+                           insertFitness eid f $ [MA.fromList MA.Seq theta]                            runEqSat myCost rewritesParams 1                            cleanDB @@ -288,7 +316,8 @@                            let (tree, ps) = if convertParam then floatConstsToParam tree' else (tree', theta)                                theta      = if convertParam then if dist==MSE then ps <> params else ps else params                            eid <- fromTree myCost tree >>= canonical-                           insertFitness eid f $ MA.fromList MA.Seq theta+                           -- TODO: how to import MvSR?+                           insertFitness eid f $ [MA.fromList MA.Seq theta]                            runEqSat myCost rewritesParams 1                            cleanDB getFormat :: String -> SRAlgs@@ -314,44 +343,45 @@     toTuple [eq, t, f] = (eq, t, f)     toTuple xss = error $ show xss -getParents False ecs = pure ecs-getParents True  ecs = IntSet.toList <$> getParentsOf (IntSet.fromList ecs) 500000 (IntSet.fromList ecs)+getParents False _ ecs = pure ecs+getParents True  p ecs = IntSet.toList <$> getParentsOf p (IntSet.fromList ecs) 300000 ecs  isBest (e', en') = do e <- canonical e'                       best <- gets (_best . _info . (IntMap.! e) . _eClass) >>= canonize                       en <- canonize en'                       pure (en == best) -getParentsOf :: IntSet.IntSet -> Int -> IntSet.IntSet -> RndEGraph IntSet.IntSet-getParentsOf visited n queue | IntSet.size queue >= n = pure queue-getParentsOf visited n queue =-   do ecs          <- Prelude.mapM canonical (IntSet.toList queue)-      parents'     <- IntSet.unions <$> Prelude.mapM canonizeParents ecs-      uneval       <- gets (_unevaluated . _eDB)-      grandParents <- getParentsOf (IntSet.union visited parents') (n-1) parents'-      pure (IntSet.filter (not . (`IntSet.member` uneval)) $ (queue <> grandParents))-   where-      canonizeParents ec = do parents <- gets (_parents . (IntMap.! ec) . _eClass)-                                          >>= Prelude.mapM (\(e, en) -> isBest (e, en) >>= \b -> pure (e, en, b)) . Set.toList-                                          >>= pure . Set.fromList--                              pure $ IntSet.fromList . Prelude.map (\(e, _, _) -> e) . Set.toList $-                                    Set.filter (\(e, en, b) -> b && ec `Prelude.elem` (childrenOf en)-                                                   && not (e `IntSet.member` visited)-                                                ) parents+getParentsOf :: (EClass -> Bool) -> IntSet.IntSet -> Int -> [EClassId] -> RndEGraph IntSet.IntSet+getParentsOf p visited n queue | IntSet.size visited >= n || null queue = pure visited+getParentsOf p visited n queue =+   do parents'     <- IntSet.unions <$> Prelude.mapM (\e -> canonical e >>= canonizeParents) queue +      grandParents <- getParentsOf p ((visited <> parents')) n (IntSet.toList parents')+      pure (visited <> grandParents)+   where+      filterUneval uneval = IntSet.filter (`IntSet.notMember` uneval)+      isNew ec (e, en) = ec `Prelude.elem` (childrenOf en) && (e `IntSet.notMember` visited)+      canonizeParents ec = do ecl <- gets ((IntMap.! ec) . _eClass)+                              let parents' = Set.toList . Set.filter (isNew ec) $ _parents ecl+                              parents <- Prelude.map fst <$> filterM isBest parents'+                              pure (IntSet.fromList parents)  isLeft (Left _)   = True isLeft _          = False fromLeft (Left x) = x fromLeft _        = undefined -getAllPatternsFrom :: Monad m => (Int -> Bool) -> Map.Map Pattern Int -> [EClassId] -> EGraphST m (Map.Map Pattern Int)-getAllPatternsFrom pSz counts []     = pure counts-getAllPatternsFrom pSz counts (x:xs) = do pats <- Map.fromListWith (+) . Prelude.map (,1) . Prelude.filter (pSz . lenPat) <$> getAllPatterns (<=max_pat) x-                                          getAllPatternsFrom pSz (Map.unionWith (+) pats counts) xs-  where max_pat = 10+addTuple (a, b) (c, d) = (a+c, b+d) +getAllPatternsFrom :: (Int -> Bool) -> Map.Map Pattern (Int, Double) -> [EClassId] -> EGraphST IO (Map.Map Pattern (Int, Double))+getAllPatternsFrom pSz counts []     = pure $ Map.map (\(v1, v2) -> (v1, v2/fromIntegral v1)) counts+getAllPatternsFrom pSz counts (x:xs) = do fit' <- getFitness x +                                          case fit' of +                                            Nothing -> getAllPatternsFrom pSz counts xs+                                            Just fit -> do+                                                         pats <- Map.map (,fit) <$> getAllPatterns pSz x+                                                         getAllPatternsFrom pSz (Map.unionWith addTuple pats counts) xs+ relabelVarPat :: Pattern -> Pattern relabelVarPat t = alg t `evalState` 65    where@@ -372,22 +402,33 @@                     >>= getEvaluated   pure (pat, IntSet.size ecs) -getEvaluated ecs = getParentsOf (IntSet.fromList ecs) 500000 (IntSet.fromList ecs)+getEvaluated ecs = getParentsOf (const True) (IntSet.fromList ecs) 500000 ecs -getAllPatterns :: Monad m => (Int -> Bool) -> EClassId -> EGraphST m [Pattern]+getAllPatterns :: Monad m => (Int -> Bool) -> EClassId -> EGraphST m (Map.Map Pattern Int) getAllPatterns pSz eid = do    eid' <- canonical eid    best <- gets (_best . _info . (IntMap.! eid') . _eClass)    case best of-      Var ix     -> pure [VarPat 'A', Fixed (Var ix)]-      Param ix   -> pure [VarPat 'A', Fixed (Param ix)]-      Const x    -> pure [VarPat 'A', Fixed (Const x)]-      Uni f t    -> do pats <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz t-                       pure (VarPat 'A' : [relabelVarPat $ Fixed (Uni f t') | t' <- pats])-      Bin op l r | l==r -> do pats <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz l-                              pure (VarPat 'A' : [relabelVarPat $ Fixed (Bin op l' l') | l' <- pats])-                  | otherwise -> do patsL <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz l-                                    patsR <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz r-                                    pure (VarPat 'A' : [relabelVarPat $ Fixed (Bin op l' r') | l' <- patsL, r' <- patsR])-+      Var ix     -> pure $ Map.fromList [(VarPat 'A', 1), (Fixed (Var ix), 1)]+      Param ix   -> pure $ Map.fromList [(VarPat 'A', 1), (Fixed (Param ix), 1)]+      Const x    -> pure $ Map.fromList [(VarPat 'A', 1), (Fixed (Const x), 1)]+      Uni f t    -> do pats <- Map.filterWithKey (\k _ -> (pSz . lenPat) k) <$> getAllPatterns pSz t+                       pure $ Map.insertWith (+) (VarPat 'A') 1 +                            $ Map.mapKeysWith (+) (\t' -> Fixed (Uni f t')) pats+      Bin op l r | l==r -> do pats <- Map.filterWithKey (\k _ -> (pSz . lenPat) k) <$> getAllPatterns pSz l+                              pure $ Map.insertWith (+) (VarPat 'A') 1 $ Map.mapKeysWith (+) (\t' -> Fixed (Bin op t' t')) pats+                  | otherwise -> do patsL <- Map.filterWithKey (\k _ -> (pSz . lenPat) k) <$> getAllPatterns pSz l+                                    patsR <- Map.filterWithKey (\k _ -> (pSz . lenPat) k) <$> getAllPatterns pSz r+                                    pure $ Map.fromList $ (VarPat 'A', 1) : [(relabelVarPat $ Fixed (Bin op l' r'), min vl vr) | (l', vl) <- Map.toList patsL, (r', vr) <- Map.toList patsR] +isNotTrivial :: Monad m => Int -> EClassId -> EGraphST m Bool+isNotTrivial n ec = do+  c <- gets (_consts . _info . (IntMap.! ec) . _eClass)+  m <- gets (_size . _info . (IntMap.! ec) . _eClass)+  pure (c == NotConst && m >= n)+removeNotTrivial :: Monad m => Int -> [EClassId] -> EGraphST m [EClassId]+removeNotTrivial n [] = pure []+removeNotTrivial n (ec:ecs) = do+  b <- isNotTrivial n ec+  ecs' <- removeNotTrivial n ecs+  pure $ if b then (ec:ecs') else ecs'
apps/rEGGression/Main.hs view
@@ -143,14 +143,14 @@   let cmd = parseCmd parseDist (B.pack $ unwords args)   runIfRight cmd -reportCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+reportCmd :: Distribution -> [DataSet] -> [DataSet] -> [String] -> Repl () reportCmd _ _ _ [] = helpCmd ["report"] reportCmd dist trainData testData args =   case readMaybe @Int (head args) of     Nothing -> io.putStrLn $ "The id must be an integer."     Just n  -> run (Report n (dist, trainData, testData)) -optimizeCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+optimizeCmd :: Distribution -> [DataSet] -> [DataSet] -> [String] -> Repl () optimizeCmd _ _ _ [] = helpCmd ["optimize"] optimizeCmd dist trainData testData args =   case readMaybe @Int (head args) of@@ -164,7 +164,7 @@                         Nothing -> io.putStrLn $ "The argument must be an integer."                         Just n  -> run (Subtrees n) -insertCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+insertCmd :: Distribution -> [DataSet] -> [DataSet] -> [String] -> Repl () insertCmd dist trainData testData [] = helpCmd ["insert"] insertCmd dist trainData testData args = do   let etree = parseSR TIR "" False $ B.pack (unwords args)@@ -296,10 +296,14 @@ main = do   args <- execParser opts   g <- getStdGen-  (dataTrain, varnames) <- loadTrainingOnly (_dataset args) True-  (dataTest, _)  <- if null (_testData args)-                     then pure (dataTrain, "")-                     else loadTrainingOnly (_testData args) True+  let datasets = words (_dataset args)+  dataTrainsWP <- Prelude.mapM (flip loadTrainingOnly True) datasets+  let dataTrains = Prelude.map fst dataTrainsWP+      varnames   = snd . head $ dataTrainsWP++  dataTests  <- if null (_testData args)+                  then pure dataTrains+                  else Prelude.map fst <$> (Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args))   eg <- if (not.null) (_loadFrom args)            then decode <$> BS.readFile (_loadFrom args)            else if (not. null) (_parseCSV args)@@ -309,10 +313,10 @@       dist = _distribution args       funs = [ helpCmd              , topCmd-             , reportCmd dist dataTrain dataTest-             , optimizeCmd dist dataTrain dataTest+             , reportCmd dist dataTrains dataTests+             , optimizeCmd dist dataTrains dataTests              , subtreesCmd-             , insertCmd dist dataTrain dataTest+             , insertCmd dist dataTrains dataTests              , countPatCmd              , distCmd              , paretoCmd@@ -323,7 +327,7 @@       cmdMap = Map.fromList $ Prelude.zip commands funs        repl = evalRepl (const $ pure ">>> ") (cmd cmdMap) [] Nothing Nothing (Word comp) ini final-      crDB = if _calcDL args then (createDB >> fillDL dist dataTrain >> rebuildAllRanges) else (createDB >> rebuildAllRanges)+      crDB = if _calcDL args then (createDBBest >> fillDL dist dataTrains >> rebuildAllRanges) else (createDBBest >> rebuildAllRanges)   if (not.null) (_convertFromTo args)      then convert (_convertFromTo args) (_to args) varnames      else do when (_calcDL args) $ putStrLn "Calculating DL..."
apps/rEGGression/Util.hs view
@@ -88,13 +88,18 @@     pure (fitnessFun nIter distribution dataTrain tree thetaOrigs) {-# INLINE fitnessFunRep #-} +fitnessMV :: Int -> Distribution -> [DataSet] -> Fix SRTree -> RndEGraph (Double, [PVector])+fitnessMV nIter distribution dataTrainsVals _tree = do+  response <- forM dataTrainsVals $ \dt -> fitnessFunRep nIter distribution dt _tree+  pure (minimum (Prelude.map fst response), Prelude.map snd response) + -- helper query functions -- TODO: move to egraph lib getFitness :: EClassId -> RndEGraph (Maybe Double) getFitness c = gets (_fitness . _info . (IM.! c) . _eClass) {-# INLINE getFitness #-}-getTheta :: EClassId -> RndEGraph (Maybe PVector)+getTheta :: EClassId -> RndEGraph ([PVector]) getTheta c = gets (_theta . _info . (IM.! c) . _eClass) {-# INLINE getTheta #-} getSize :: EClassId -> RndEGraph Int@@ -126,35 +131,45 @@  loadTrainingOnly fname b = getTrain <$> loadDataset fname b -printExpr :: DataSet -> DataSet -> Distribution -> EClassId -> RndEGraph ()+mvFun fun thetas datasets = Prelude.map (\(theta, (x,y,e)) -> fun x y e theta)+                          $ Prelude.zip thetas datasets++printExpr :: [DataSet] -> [DataSet] -> Distribution -> EClassId -> RndEGraph () printExpr dataTrain dataTest distribution ec = do-        theta <- fromJust <$> getTheta ec+        thetas <- getTheta ec          bestExpr <- getBestExpr ec-        let (x, y, mYErr) = dataTrain-            (x_te, y_te, mYErr_te) = dataTest-            best'     = relabelParams bestExpr-            expr      = paramsToConst (MA.toList theta) best'-            mse_train = mse x y best' theta-            mse_te    = mse x_te y_te best' theta-            r2_train  = r2 x y best' theta-            r2_te     = r2 x_te y_te best' theta-            nll_train  = nll distribution mYErr x y best' theta-            nll_te     = nll distribution mYErr_te x_te y_te best' theta-            mdl_train  = mdl distribution mYErr x y theta best'-            mdl_te     = mdl distribution mYErr_te x_te y_te theta best'-            thetaStr   = intercalate ", " $ Prelude.map show (MA.toList theta)+        let --(x, y, mYErr) = dataTrain+            --(x_te, y_te, mYErr_te) = dataTest+            best'       = relabelParams bestExpr++            mseMV x y e theta = printf "%.4e" $ mse x y best' theta+            r2MV  x y e theta = printf "%.4e" $ r2 x y best' theta+            nllMV x y e theta = printf "%.4e" $ nll distribution e x y best' theta+            mdlMV x y e theta = printf "%.4e" $ mdl distribution e x y theta best'++            -- expr        = paramsToConst (MA.toList theta) best'+            mse_trains  = intercalate "; " $ mvFun mseMV thetas dataTrain+            mse_tes     = intercalate "; " $ mvFun mseMV thetas dataTest+            r2_trains   = intercalate "; " $ mvFun r2MV thetas dataTrain+            r2_tes      = intercalate "; " $ mvFun r2MV thetas dataTest+            nll_trains  = intercalate "; " $ mvFun nllMV thetas dataTrain+            nll_tes     = intercalate "; " $ mvFun nllMV thetas dataTest+            mdl_trains  = intercalate "; " $ mvFun mdlMV thetas dataTrain+            mdl_tes     = intercalate "; " $ mvFun mdlMV thetas dataTest+            thetaStr    = intercalate "; " $ Prelude.map (intercalate ", " . Prelude.map show . MA.toList) thetas+        insertDL ec $ Prelude.maximum $ Prelude.map (\(theta, (x, y, mYerr)) -> mdl distribution mYerr x y theta best') $ Prelude.zip thetas dataTrain          lift . putStr $ "Evaluation metrics for expression (" <> (show ec) <> "): "         lift . putStr $ setSGRCode [SetConsoleIntensity BoldIntensity]         lift . putStrLn $ showExpr best'         lift . putStr $ setSGRCode [Reset]-        lift . putStrLn $ "# of nodes\t" <> show (countNodes $ convertProtectedOps expr)+        lift . putStrLn $ "# of nodes\t" <> show (countNodes $ convertProtectedOps best')         lift . putStrLn $ "params:\t[" <> thetaStr <> "]" -        let rows = [ rowG ["MSE", printf "%.4f" mse_train, printf "%.4f" mse_te]-                   , rowG ["R^2", printf "%.4f" r2_train, printf "%.4f" r2_te]-                   , rowG ["nll", printf "%.4f" nll_train, printf "%.4f" nll_te]-                   , rowG ["DL", printf "%.4f" mdl_train, printf "%.4f" mdl_te]+        let rows = [ rowG ["MSE", mse_trains, mse_tes]+                   , rowG ["R^2", r2_trains, r2_tes]+                   , rowG ["nll", nll_trains, nll_tes]+                   , rowG ["DL",  mdl_trains, mdl_tes]                    ]             columnsReport = [def, numCol, numCol]             headerReport = titlesH $ Prelude.map bold ["Metric", "Training", "Test"]@@ -167,19 +182,19 @@                          dl  <- egraph $ getDL eid                          let fit' = case fit of                                       Nothing -> "--"-                                      Just f  -> printf "%.4f" f+                                      Just f  -> printf "%.4e" f                              p' = case p of-                                    Nothing -> "--"-                                    Just ps -> "[" <> intercalate ", " (Prelude.map (printf "%.4f") (MA.toList ps)) <> "]"+                                    [] -> "--"+                                    pss -> intercalate ";" $ Prelude.map (\ps -> "[" <> intercalate ", " (Prelude.map (printf "%.4e") (MA.toList ps)) <> "]") pss                              dl' = case dl of                                     Nothing -> "--"-                                    Just d  -> printf "%.4f" d+                                    Just d  -> printf "%.4e" d                           pure $ colsAllG center [[show eid], justifyText 50 $ showExpr t, [fit'], justifyText 50 p', [show sz], [dl']] -printCounts (pat, cnt) = do+printCounts (pat, (cnt, avgfit)) = do   let spat = showPat pat-  pure $ colsAllG center [justifyText 50 spat, [show cnt]]+  pure $ colsAllG center [justifyText 50 spat, [show cnt], [printf "%.4e" avgfit]]   where     showPat (Fixed (Var ix)) = 'x' : show ix     showPat (Fixed (Param ix)) = 't' : show ix@@ -192,7 +207,7 @@                                 io.putStrLn $ tableString (columnHeaderTableS columns unicodeS headerSimple rows)  printMultiCounts cnts = do rows <- forM cnts printCounts-                           io.putStrLn $ tableString (columnHeaderTableS [fixedLeftCol 50, numCol] unicodeS headerCount rows)+                           io.putStrLn $ tableString (columnHeaderTableS [fixedLeftCol 50, numCol, numCol] unicodeS headerCount rows)  bold s = formatted (setSGRCode [SetConsoleIntensity BoldIntensity]) (plain s) (setSGRCode [Reset]) @@ -200,11 +215,11 @@ headerSimple = titlesH $ Prelude.map (bold) ["Id", "Expression", "Fitness", "Parameters", "Size", "DL"] columns = [numCol, fixedLeftCol 50, numCol, fixedLeftCol 50, numCol, numCol] headerCount :: HeaderSpec LineStyle (Formatted String)-headerCount = titlesH $ Prelude.map bold ["Pattern", "Count"]+headerCount = titlesH $ Prelude.map bold ["Pattern", "Count", "Avg. Fitness"]  -- RndEGraph utils -- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val-insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, PVector)) -> RndEGraph EClassId+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId insertExpr t fitFun = do     ecId <- fromTree myCost t >>= canonical     (f, p) <- fitFun t@@ -261,12 +276,12 @@               (f, p) <- fitFun t               insertFitness c f p -fillDL dist (x, y, mYErr) = do+fillDL dist datasets = do   ecs <- getAllEvaluatedEClasses   forM_ ecs $ \ec -> do-    theta <- fromJust <$> getTheta ec+    thetas <- getTheta ec     bestExpr <- relabelParams <$> getBestExpr ec-    if MA.size theta /= countParams bestExpr-       then (lift . putStrLn) $ "Wrong number of parameters in " <> showExpr bestExpr <> ": " <> show theta <> "   " <> show ec-       else do let mdl_train  = mdl dist mYErr x y theta bestExpr-               insertDL ec mdl_train+    if MA.size (head thetas) /= countParams bestExpr+       then (lift . putStrLn) $ "Wrong number of parameters in " <> showExpr bestExpr <> ": " <> show (head thetas) <> "   " <> show ec+       else do let mdl_trains = Prelude.map (\(theta, (x, y, mYerr)) -> mdl dist mYerr x y theta bestExpr) $ Prelude.zip thetas datasets+               insertDL ec $ Prelude.maximum mdl_trains
apps/srtools/Args.hs view
@@ -174,4 +174,4 @@   str >>= \s -> mkReader ("unsupported distribution " <> s) id (capitalize s)   where     capitalize ""     = ""-    capitalize (c:cs) = toUpper c : map toLower cs+    capitalize (c:cs) = toUpper c : if length cs == 2 then map toUpper cs else map toLower cs
+ apps/symregg/Main.hs view
@@ -0,0 +1,481 @@+{-# LANGUAGE  BlockArguments #-}+{-# LANGUAGE  TupleSections #-}+{-# LANGUAGE  MultiWayIf #-}+{-# LANGUAGE  OverloadedStrings #-}+{-# LANGUAGE  BangPatterns #-}+{-# LANGUAGE  TypeSynonymInstances, FlexibleInstances #-}++module Main where ++import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Simplify+import Algorithm.EqSat.Build+import Algorithm.EqSat.Queries+import Algorithm.EqSat.Info+import Algorithm.EqSat.DB+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection+import Control.Monad ( forM_, forM, when, filterM )+import Control.Monad.State.Strict+import qualified Data.IntMap.Strict as IM+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.Maybe ( fromJust, isNothing )+import Data.SRTree+import Data.SRTree.Print ( showExpr )+import Options.Applicative as Opt hiding (Const)+import Random+import System.Random+import Data.List ( intercalate )+import qualified Data.IntSet as IntSet+import Algorithm.EqSat (runEqSat)+import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS+import Debug.Trace+import qualified Data.HashSet as Set+import Control.Lens (over)+import qualified Data.Map.Strict as Map++import Util++data Alg = OnlyRandom | BestFirst deriving (Show, Read, Eq)++egraphSearch :: [(DataSet, DataSet)] -> [DataSet] -> Args -> StateT EGraph (StateT StdGen IO) ()+-- terms nEvals maxSize printPareto printTrace slowIter slowRep =+egraphSearch dataTrainVals dataTests args = do+  if null (_loadFrom args)+    then do ecFst <- insertRndExpr (_maxSize args) rndTerm rndNonTerm2+            --ecFst <- insertBestExpr -- use only to debug+            updateIfNothing fitFun ecFst+            insertTerms+            evaluateUnevaluated fitFun+            runEqSat myCost (if _nParams args == 0 then rewritesWithConstant else rewritesParams) 1+            cleanDB+    else (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg)+  nCls <- gets (IM.size . _eClass)+  nUnev <- gets (IntSet.size . _unevaluated . _eDB)+  let nEvs = nCls - nUnev++  while ((<(_gens args)) . snd) (10, nEvs) $+    \(radius, nEvs) ->+      do+       nCls  <- gets (IM.size . _eClass)+       --nUnev <- gets (IntSet.size . _unevaluated . _eDB)+       --let nEvs = nCls - nUnev -- WARNING: see if it affects results++       ecN <- case (_alg args) of+                    OnlyRandom -> do let ratio = fromIntegral nEvs / fromIntegral nCls+                                     b <- rnd (tossBiased ratio)+                                     ec <- if b && ratio > 0.99+                                              then insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+                                              else do mEc <- pickRndSubTree -- evaluateRndUnevaluated fitFun >>= canonical+                                                      case mEc of+                                                           Nothing ->insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+                                                           Just ec' -> pure ec'+                                     pure ec+                    BestFirst  -> do+                      ecsPareto <- getParetoEcsUpTo 50 (_maxSize args)+                      ecPareto     <- combineFrom ecsPareto >>= canonical+                      curFitPareto <- getFitness ecPareto++                      if isNothing curFitPareto+                        then pure ecPareto+                        else do ecsBest    <- getTopFitEClassThat 100 (isSizeOf (<(_maxSize args)))+                                ecBest     <- combineFrom ecsBest >>= canonical+                                curFitBest <- getFitness ecBest+                                if isNothing curFitBest+                                  then pure ecBest+                                  else do ee <- pickRndSubTree+                                          case ee of+                                            Nothing -> insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+                                            Just c  -> pure c++       -- when upd $+       ecN' <- canonical ecN+       upd <- updateIfNothing fitFun ecN'+       when upd $ runEqSat myCost (if _nParams args == 0 then rewritesWithConstant else rewritesParams) 1 >> cleanDB >> refitChanged+++       when (upd && (_trace args))+         do+            ecN'' <- canonical ecN'+            _tree <- getBestExpr ecN''+            fi <- negate . fromJust <$> getFitness ecN''+            thetas <- getTheta ecN''+            let thetaStr   = intercalate "_" $ Prelude.map (intercalate ";" . Prelude.map show . MA.toList) thetas+            io . putStrLn $ showExpr _tree <> "," <> thetaStr <> "," <> show fi+            pure ()+       let radius' = if (not upd) then (max 10 $ min (500 `div` (_maxSize args)) (radius+1)) else (max 10 $ radius-1)+           nEvs'    = nEvs + if upd then 1 else 0+       pure (radius', nEvs')+  io $ putStrLn csvHeader+  if (_printPareto args)+    then paretoFront (_maxSize args) printExpr+    else printBest printExpr+  when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode )+  where+    maxSize = _maxSize args+    relabel        = if (_nParams args == -1) then relabelParams else relabelParamsOrder+    fitFun = fitnessMV (_optRepeat args) (_optIter args) (_distribution args) dataTrainVals++    refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList+                      modify' $ over (eDB . refits) (const Set.empty)+                      forM_ ids $ \ec -> do t <- getBestExpr ec+                                            (f, p) <- fitFun t+                                            insertFitness ec f p++    combineFrom [] = pure 0+    combineFrom ecs = do+      p1  <- rnd (randomFrom ecs) >>= canonical+      p2  <- rnd (randomFrom ecs) >>= canonical+      coin <- rnd toss+      if coin+         then crossover p1 p2 >>= canonical+         else mutate p1 >>= canonical++    combineFrom' [] = pure 0 -- this is the first terminal and it will always be already evaluated+    combineFrom' ecs = do+        nt  <- rnd rndNonTerm+        p1  <- rnd (randomFrom ecs)+        p2  <- rnd (randomFrom ecs)+        l1  <- rnd (randomFrom [2..(_maxSize args)-2]) -- sz 10: [2..8]++        e1  <- randomChildFrom p1 l1 >>= canonical+        ml  <- gets (_size . _info . (IM.! e1) . _eClass)+        l2  <- rnd (randomFrom [1..((_maxSize args) - ml - 1)]) -- maxSize - maxSize + 2 - 2= 0 -- sz 10: [1..7] (2) / [1..1] (8)+        e2  <- randomChildFrom p2 l2 >>= canonical+        case nt of+          Uni Id ()    -> canonical e1+          Uni f ()     -> add myCost (Uni f e1) >>= canonical+          Bin op () () -> do b <- rnd toss+                             if b+                              then add myCost (Bin op e1 e2) >>= canonical+                              else add myCost (Bin op e2 e1) >>= canonical+          _            -> canonical e1 -- it is a terminal, should it happen?++    randomChildFrom ec' maxL = do+      p <- rnd toss -- whether to go deeper or return this level+      ec <- canonical ec'+      l <- gets (_size . _info . (IM.! ec) . _eClass )++      if p || l > maxL+          then do --enodes <- gets (_eNodes . (IM.! ec) . _eClass)+                  enode  <- gets (_best . _info . (IM.! ec) . _eClass) -- we should return the best otherwise we may build larger exprs+                  case enode of+                      Uni _ eci     -> randomChildFrom eci maxL+                      Bin _ ecl ecr -> do coin <- rnd toss+                                          if coin+                                            then randomChildFrom ecl maxL+                                            else randomChildFrom ecr maxL+                      _ -> pure ec -- this shouldn't happen unless maxL==0+          else pure ec++    nonTerms   = parseNonTerms (_nonterminals args)+    --[ Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () (), Bin PowerAbs () (),  Uni Recip ()]+    (Sz2 _ nFeats) = MA.size (getX . fst . head $ dataTrainVals)+    terms          = if _distribution args == ROXY+                          then [var 0, param 0]+                          else [var ix | ix <- [0 .. nFeats-1]]+                               <> if _nParams args == -1+                                     then [param 0]+                                     else Prelude.map param [0 .. _nParams args - 1]+    rndTerm    = Random.randomFrom terms+    rndNonTerm = Random.randomFrom $ (Uni Id ()) : nonTerms+    rndNonTerm2 = Random.randomFrom nonTerms+    uniNonTerms = Prelude.filter isUni nonTerms+    binNonTerms = Prelude.filter isBin nonTerms+    isUni (Uni _ _) = True+    isUni _         = False+    isBin (Bin _ _ _) = True+    isBin _           = False++    insertTerms =+        forM terms $ \t -> do fromTree myCost t >>= canonical++    printExpr :: Int -> EClassId -> RndEGraph ()+    printExpr ix ec = do +        thetas' <- gets (_theta . _info . (IM.! ec) . _eClass)+        bestExpr <- getBestExpr ec+        let nParams = countParams bestExpr+            fromSz (MA.Sz x) = x +            nThetas = Prelude.map (fromSz . MA.size) thetas'+        (_, thetas) <- if Prelude.any (/=nParams) nThetas+                        then fitFun bestExpr+                        else pure (1.0, thetas')++        forM_ (Prelude.zip3 dataTrainVals dataTests thetas) $ \((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 = _distribution args+                best'     = relabelParams bestExpr+                expr      = paramsToConst (MA.toList theta) best'+                mse_train = mse x y best' theta+                mse_val   = mse x_val y_val best' theta+                mse_te    = mse x_te y_te best' theta+                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  = mdl distribution mYErr x y theta best'+                mdl_val    = mdl distribution mYErr_val x_val y_val theta best'+                mdl_te     = mdl distribution mYErr_te x_te y_te theta best'+                vals       = intercalate ","+                           $ Prelude.map show [mse_train, mse_val, mse_te+                                              , r2_train, r2_val, r2_te+                                              , nll_train, nll_val, nll_te+                                              , mdl_train, mdl_val, mdl_te]+                thetaStr   = intercalate ";" $ Prelude.map show (MA.toList theta)+            io . putStrLn $ show ix <> "," <> showExpr expr <> ","+                          <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)+                          <> "," <> vals++    -- From eggp+    crossover p1 p2 = do sz <- getSize p1+                         if sz == 1+                          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 (\c -> (<maxSize-sz) <$> getSize c) cands+      candidates  <- filterM (\c -> doesNotExistGens mGrandParents (parent c)) 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+                  pos <- rnd $ randomRange (0, sz-1)+                  tree <- mutAt pos maxSize Nothing p+                  fromTree myCost (relabel tree) >>= canonical++    peel :: Fix SRTree -> SRTree ()+    peel (Fix (Bin op l r)) = Bin op () ()+    peel (Fix (Uni f t)) = 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 -- we chose to mutate the root+    mutAt 0 1        _             _ = rnd $ randomFrom terms -- we don't have size left+    mutAt 0 sizeLeft (Just parent) _ = do -- we reached the mutation place+      ec    <- insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical -- create a random expression with the size limit+      (Fix tree) <- getBestExpr ec           --+      root  <- getBestENode ec+      exist <- canonize (parent ec) >>= doesExist+      if exist+          -- the expression `parent ec` already exists, try to fix+          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+                  if null candidates+                      then pure $ Fix tree -- there's no candidate, so we failed and admit defeat+                      else do newToken <- rnd (randomFrom candidates)+                              pure . Fix $ replaceChildren (childrenOf tree) newToken++          else pure . Fix $ tree++    mutAt pos sizeLeft parent 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' >>= \t -> (Fix . Uni f) <$> mutAt (pos-1) (sizeLeft-1) (Just $ Uni f) t+          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'++    checkToken parent en' = do  en <- canonize en'+                                mEc <- gets ((Map.!? en) . _eNodeToEClass)+                                case mEc of+                                    Nothing -> pure True+                                    Just ec -> do ec' <- canonical ec+                                                  ec'' <- canonize (parent ec')+                                                  not <$> doesExist ec''+    doesExist, doesNotExist :: ENode -> RndEGraph Bool+    doesExist en = gets ((Map.member en) . _eNodeToEClass)+    doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)++    doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool+    doesNotExistGens []              en = gets ((Map.notMember en) . _eNodeToEClass)+    doesNotExistGens (mGrand:grands) en = do  b <- gets ((Map.notMember en) . _eNodeToEClass)+                                              if b+                                                then pure True+                                                else case mGrand of+                                                    Nothing -> pure False+                                                    Just gf -> do ec  <- gets ((Map.! en) . _eNodeToEClass)+                                                                  en' <- canonize (gf ec)+                                                                  doesNotExistGens grands en'++data Args = Args+  { _dataset      :: String,+    _testData     :: String,+    _gens         :: Int,+    _alg          :: Alg,+    _maxSize      :: Int,+    _split        :: Int,+    _printPareto  :: Bool,+    _trace        :: Bool,+    _distribution :: Distribution,+    _optIter      :: Int,+    _optRepeat    :: Int,+    _nParams      :: Int,+    _nonterminals :: String,+    _dumpTo       :: String,+    _loadFrom     :: String+  }+  deriving (Show)++-- parser of command line arguments+opt :: Parser Args+opt = Args+   <$> strOption+       ( long "dataset"+       <> short 'd'+       <> metavar "INPUT-FILE"+       <> help "CSV dataset." )+  <*> strOption+       ( long "test"+       <> short 't'+       <> value ""+       <> showDefault+       <> help "test data")+   <*> option auto+      ( long "generations"+      <> short 'g'+      <> metavar "GENS"+      <> showDefault+      <> value 100+      <> help "Number of generations." )+   <*> option auto+       ( long "algorithm"+       <> short 'a'+       <> metavar "ALG"+       <> help "Algorithm." )+  <*> option auto+       ( long "maxSize"+       <> short 's'+       <> help "max-size." )+  <*> option auto+       ( long "split"+       <> short 'k'+       <> value 1 +       <> showDefault+       <> help "k-split ratio training-validation")+  <*> switch+       ( long "print-pareto"+       <> help "print Pareto front instead of best found expression")+  <*> switch+       ( long "trace"+       <> help "print all evaluated expressions.")+  <*> option auto+       ( long "loss"+       <> value MSE+       <> showDefault+       <> help "distribution of the data.")+  <*> option auto+       ( long "opt-iter"+       <> value 30+       <> showDefault+       <> help "number of iterations in parameter optimization.")+  <*> option auto+       ( long "opt-retries"+       <> value 1+       <> showDefault+       <> help "number of retries of parameter fitting.")+  <*> option auto+       ( long "number-params"+       <> value (-1)+       <> showDefault+       <> help "maximum number of parameters in the model. If this argument is absent, the number is bounded by the maximum size of the expression and there will be no repeated parameter.")+  <*> strOption+       ( long "non-terminals"+       <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+       <> showDefault+       <> help "set of non-terminals to use in the search."+       )+  <*> strOption+       ( long "dump-to"+       <> value ""+       <> showDefault+       <> help "dump final e-graph to a file."+       )+  <*> strOption+       ( long "load-from"+       <> value ""+       <> showDefault+       <> help "load initial e-graph from a file."+       )++main :: IO ()+main = do+  args <- execParser opts+  g    <- getStdGen+  let datasets = words (_dataset args)+  dataTrains' <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets +  dataTests   <- if null (_testData args)+                  then pure dataTrains'+                  else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args)++  let (dataTrainVals, g') = runState (Prelude.mapM (`splitData` (_split args)) dataTrains') g+      alg = evalStateT (egraphSearch dataTrainVals dataTests args) emptyGraph+  evalStateT alg g'++  where+    opts = Opt.info (opt <**> helper)+            ( fullDesc <> progDesc "Symbolic Regression search algorithm\+                                   \ exploiting the potentials of equality saturation\+                                   \ and e-graphs."+           <> header "SymRegg - symbolic regression with e-graphs."+            )
+ apps/symregg/Random.hs view
@@ -0,0 +1,54 @@+module Random where ++import System.Random +import Control.Monad.State.Strict+import Control.Monad+import Data.SRTree +import Data.SRTree.Eval+import Data.Massiv.Array as MA++type Rng a = StateT StdGen IO a ++toss :: Rng Bool+toss = state random+{-# INLINE toss #-}++tossBiased :: Double -> Rng Bool+tossBiased p = do r <- state random +                  pure (r < p)++randomVal :: Rng Double+randomVal = state random++randomRange :: (Ord val, Random val) => (val, val) -> Rng val+randomRange rng = state (randomR rng)+{-# INLINE randomRange #-}++randomFrom :: [a] -> Rng a+randomFrom funs = do n <- randomRange (0, length funs - 1)+                     pure $ funs !! n+{-# INLINE randomFrom #-}++randomVec :: Int -> Rng PVector+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))++randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree)+randomTree minDepth maxDepth maxSize genTerm genNonTerm grow +  | noSpaceLeft = genTerm+  | needNonTerm = genRecursion +  | otherwise   = do r <- toss+                     if r +                       then genTerm +                       else genRecursion+  where +    noSpaceLeft = maxDepth <= 1 || maxSize <= 2+    needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2++    genRecursion = do +        node <- genNonTerm+        case node of +          Uni f _    -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow +          Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow+                           r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow +                           pure . Fix  $ Bin op l r+{-# INLINE randomTree #-}
+ apps/symregg/Util.hs view
@@ -0,0 +1,255 @@+{-# LANGUAGE  BlockArguments #-}+{-# LANGUAGE  TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Random+import Algorithm.SRTree.Likelihoods+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_, filterM )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries++type RndEGraph a = EGraphST (StateT StdGen IO) a+type DataSet = (SRMatrix, PVector, Maybe PVector)++csvHeader :: String+csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"++io :: IO a -> RndEGraph a+io = lift . lift+{-# INLINE io #-}+rnd :: StateT StdGen IO a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _)     = 1+myCost (Const _)   = 1+myCost (Param _)   = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t)   = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()+while p arg prog = do when (p arg) do arg' <- prog arg+                                      while p arg' prog++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) _tree thetaOrig =+  if isNaN val || isNaN tr+    then (-(1/0), theta) -- infinity+    else (val, theta) -- (min tr val, theta)+  where+    tree          = relabelParams _tree+    nParams       = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+    (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+    evalF a b c   = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+    tr            = evalF x y mYErr+    val           = evalF x_val y_val mYErr_val++{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nRep nIter distribution dataTrain dataVal _tree = do+    let tree = relabelParams _tree+        nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+    thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+    let fits = Prelude.map (fitnessFun nIter distribution dataTrain dataVal _tree) thetaOrigs+    pure (maximumBy (compare `on` fst) fits)+{-# INLINE fitnessFunRep #-}++fitnessMV :: Int -> Int -> Distribution -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [PVector])+fitnessMV nRep nIter distribution dataTrainsVals _tree = do+  response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv _tree+  pure (minimum (Prelude.map fst response), Prelude.map snd response)++-- helper query functions+-- TODO: move to egraph lib+getFitness :: EClassId -> RndEGraph (Maybe Double)+getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: EClassId -> RndEGraph [PVector]+getTheta c = gets (_theta . _info . (IM.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: EClassId -> RndEGraph Int+getSize c = gets (_size . _info . (IM.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}++getBestFitness :: RndEGraph (Maybe Double)+getBestFitness = do+    bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+    gets (_fitness . _info . (IM.! bec) . _eClass)++-- TODO: move to dataset lib+chunksOf :: Int -> [e] -> [[e]]+chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))+ where+  splitter :: [e] -> ([e] -> a -> a) -> a -> a+  splitter [] _ n = n+  splitter l c n = l `c` splitter (Prelude.drop i l) c n+  build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]+  build g = g (:) []++splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)+splitData (x, y, mYErr) k = do+  if k == 1+    then pure ((x, y, mYErr), (x, y, mYErr))+    else do+      ixs' <- (state . shuffle) [0 .. sz-1]+      let ixs = chunksOf k ixs'++      let (x_tr, x_te) = getX ixs x+          (y_tr, y_te) = getY ixs y+          mY = fmap (getY ixs) mYErr+          (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)+      pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))+  where+    (MA.Sz sz) = MA.size y+    comp_x     = MA.getComp x+    comp_y     = MA.getComp y++    getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)+    getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]+                    in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+                       , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+                       )+    getY :: [[Int]] -> PVector -> (PVector, PVector)+    getY ixs ys  = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+                   , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+                   )++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b++parseNonTerms :: String -> [SRTree ()]+parseNonTerms = Prelude.map toNonTerm . splitOn ","+  where+    binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+    uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+    toNonTerm xs' = let xs = Prelude.map toLower xs'+                    in case binTerms Map.!? xs of+                          Just op -> Bin op () ()+                          Nothing -> case uniTerms Map.!? xs of+                                          Just f -> Uni f ()+                                          Nothing -> error $ "invalid non-terminal " <> show xs++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId+insertExpr t fitFun = do+    ecId <- fromTree myCost t >>= canonical+    (f, p) <- fitFun t+    insertFitness ecId f p+    io . putStrLn $ "Best fit global: " <> show f+    pure ecId+  where powabs l r  = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+      mf <- getFitness ec+      case mf of+        Nothing -> do+          t <- getBestExpr ec+          (f, p) <- fitFun t+          insertFitness ec f p+          pure True+        Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+                    if not (null ecIds)+                          then do rndId' <- rnd $ randomFrom ecIds+                                  rndId  <- canonical rndId'+                                  constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+                                  case constType of+                                    NotConst -> pure $ Just rndId+                                    _        -> pure Nothing+                          else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)++getBestExprWithSize n =+        do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+           if (not (null ec))+            then do+              bestFit <- getFitness $ head ec+              bestP   <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+              (:[]) . (,bestP) . (,bestFit) . (,ec) <$> getBestExpr (head ec)+            else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+      do grow <- rnd toss+         n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+         t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+         fromTree myCost t >>= canonical++printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest printExprFun = do+      bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+      printExprFun 0 bec++paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront maxSize printExprFun = go 1 0 (-(1.0/0.0))+    where+    go :: Int -> Int -> Double -> RndEGraph ()+    go n ix f+        | n > maxSize = pure ()+        | otherwise   = do+            ecList <- getBestExprWithSize n+            if not (null ecList)+                then do let (((_, ec), mf), _) = head ecList+                            improved = fromJust mf > f+                        ec' <- traverse canonical ec+                        when improved $ printExprFun ix (head ec')+                        go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))+                else go (n+1) ix f++evaluateUnevaluated fitFun = do+          ec <- gets (IntSet.toList . _unevaluated . _eDB)+          forM_ ec $ \c -> do+              t <- getBestExpr c+              (f, p) <- fitFun t+              insertFitness c f p++evaluateRndUnevaluated fitFun = do+          ec <- gets (IntSet.toList . _unevaluated . _eDB)+          c <- rnd . randomFrom $ ec+          t <- getBestExpr c+          (f, p) <- fitFun t+          insertFitness c f p+          pure c+
apps/tinygp/GP.hs view
@@ -17,7 +17,7 @@ import Data.Massiv.Array qualified as M import Debug.Trace ( traceShow, trace ) import Util-import Data.List ( intercalate )+import Data.List ( intercalate, maximumBy ) import qualified Data.Vector.Mutable as MV  data Method = Grow | Full | BTC@@ -27,7 +27,7 @@ type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree type FitFun = Individual -> Rng Individual -data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: PVector }+data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: [PVector] }  instance Show Individual where      show (Individual t f p) = showExpr t <> "," <> show f <> "," <> show p @@ -111,8 +111,8 @@ randomIndividual hyperparams fitFun grow = do     t <- randomTree hyperparams grow      let p = countParams t-    theta' <- replicateM p (randomRange (-1,1))-    fitFun $ Individual t 0.0 (M.fromList compMode theta' :: PVector)+    --theta' <- replicateM p (randomRange (-1,1))+    fitFun $ Individual t 0.0 [] -- (M.fromList compMode theta' :: PVector)     --pure ind     --if isInfinite (_fit ind)     --   then randomIndividual hyperparams fitFun grow @@ -127,19 +127,25 @@               mapM (randomIndividual hyperparams{ _maxDepth = md} fitFun) g    pure (V.concat pop) -fitness :: SRMatrix -> PVector -> Individual -> Rng Individual-fitness x y ind = do+fitnessMV :: Distribution -> [(SRMatrix, PVector, Maybe PVector)] -> Individual -> Rng Individual+fitnessMV dist datas ind = do+  fs <- forM datas (fitness dist ind)+  let fitOpt = minimum $ map fst fs+  pure ind{_fit = fitOpt, _params = map snd fs}++fitness :: Distribution -> Individual ->  (SRMatrix, PVector, Maybe PVector) -> Rng (Double, PVector)+fitness dist ind (x, y, e) = do     let tree = relabelParams $ _tree ind         p    = countParams tree     theta1' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))     theta2' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))-    let (theta1, f1, _) = minimizeNLL MSE Nothing 50 x y tree theta1'-        (theta2, f2, _) = minimizeNLL MSE Nothing 50 x y tree theta2'-        fit1 = negate f1-        fit2 = negate f2+    let (theta1, f1, _) = minimizeNLL dist e 50 x y tree theta1'+        (theta2, f2, _) = minimizeNLL dist e 50 x y tree theta2'+        fit1 = if isNaN f1 then (-1.0/0.0) else negate f1+        fit2 = if isNaN f1 then (-1.0/0.0) else negate f2         thetaOpt = if fit1 > fit2 then theta1 else theta2         fitOpt   = max fit1 fit2-    pure ind{_fit = fitOpt, _params = thetaOpt}+    pure (fitOpt, thetaOpt)   mutate :: HyperParams -> Individual -> Rng (Maybe Individual)@@ -150,7 +156,7 @@   t <- go p (_maxSize hp) (_tree ind)   --(t, b) <- go sz (_pm hp) (_tree ind)   if b <= _pm hp && countNodes t <= _maxSize hp-     then pure . Just $ Individual t 0.0 M.empty+     then pure . Just $ Individual t 0.0 []      else pure Nothing       where         go 0 msz t = randomTree hp{_maxSize = msz-1} True@@ -210,14 +216,14 @@         Nothing -> pure parent1         Just c  -> fitFun c -printFinal ind x y x_test y_test = do+printFinal dist ind dataTrains dataTests = do   let tree     = relabelParams $ _tree ind-      theta    = _params ind-      mseTrain = mse x y tree theta-      mseTest  = mse x_test y_test tree theta-      r2Train  = r2 x y tree theta-      r2Test   = r2 x_test y_test tree theta-      thetaStr = intercalate ";" $ map show $ M.toList theta+      thetas   = _params ind+      mseTrain = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTrains+      mseTest  = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTests+      r2Train  = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTrains+      r2Test   = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTests+      thetaStr = intercalate "_" $ map (intercalate ";" . map show . M.toList) thetas   putStrLn "id,Expression,theta,size,MSE_train,MSE_test,R2_train,R2_test"   putStr $ "0," <> showExpr tree <> "," <> thetaStr <> "," <> show (countNodes tree) <> "," <> show mseTrain <> "," <> show mseTest <> "," <> show r2Train <> "," <> show r2Test @@ -229,7 +235,7 @@                            putStr " - "                             putStr (show (_fit ind))                            putStr " "-                           print (M.toList $ _params ind)+                           print (map M.toList $ _params ind) {-# INLINE report #-}  evolution :: Int -> HyperParams -> FitFun -> Rng (Individual)
apps/tinygp/Main.hs view
@@ -1,6 +1,6 @@ module Main (main) where -import GP ( HyperParams(HP), fitness, evolution, printFinal )+import GP ( HyperParams(HP), fitnessMV, evolution, printFinal ) import Data.SRTree import System.Random ( getStdGen ) import Control.Monad.State.Strict ( evalStateT )@@ -8,6 +8,7 @@ import Options.Applicative import Data.Massiv.Array  import Util+import Algorithm.SRTree.Likelihoods  -- Data type to store command line arguments data Args = Args@@ -19,7 +20,8 @@     pc        :: Double,     pm        :: Double,     _nonterminals :: String,-    _nTournament  :: Int+    _nTournament  :: Int,+    _distribution :: Distribution   }   deriving (Show) @@ -81,6 +83,11 @@        <> showDefault        <> help "tournament size."        )+   <*> option auto+       ( long "distribution"+       <> value MSE+       <> showDefault+       <> help "distribution of the data.")  nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> Fix $ Bin PowerAbs l r), Left recip, Left log, Left exp, Left (\t -> Fix $ Uni SqrtAbs t)] @@ -88,13 +95,20 @@ main = do   args <- execParser opts   g <- getStdGen-  (x, y, _) <- loadTrainingOnly (dataset args) True-  (x_test, y_test, _) <- loadTrainingOnly (_testData args) True+  --(x, y, _) <- loadTrainingOnly (dataset args) True+  --(x_test, y_test, _) <- loadTrainingOnly (_testData args) True++  let datasets = words (dataset args)+  dataTrains <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets+  dataTests  <- if null (_testData args)+                  then pure dataTrains+                  else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args)+   let hp = HP 3 10 (_maxSize args) (popSize args) (_nTournament args) (pc args) (pm args) terms (parseNonTerms $ _nonterminals args)-      (Sz2 _ nFeats) = size x+      (Sz2 _ nFeats) = size . getX $ head dataTrains       terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]]-  best <- evalStateT (evolution (gens args) hp (fitness x y)) g-  printFinal best x y x_test y_test+  best <- evalStateT (evolution (gens args) hp (fitnessMV (_distribution args) dataTrains)) g+  printFinal (_distribution args) best dataTrains dataTests   where     opts = info (opt <**> helper)             ( fullDesc <> progDesc "Very simple example of GP using SRTree."
src/Algorithm/EqSat/Build.hs view
@@ -263,6 +263,12 @@               gets (_patDB . _eDB) {-# INLINE createDB #-} +createDBBest :: Monad m => EGraphST m DB+createDBBest = do modify' $ over (eDB . patDB) (const Map.empty)+                  ecls <- gets (Prelude.map (\(eId, ec) -> (_best (_info ec), eId)) . IntMap.toList . _eClass)+                  mapM_ (uncurry addToDB) ecls+                  gets (_patDB . _eDB)+ -- | `addToDB` adds an e-node and e-class id to the database addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB () addToDB enode eid = do@@ -457,11 +463,21 @@ getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId] getAllChildEClasses eId' = do   eId <- canonical eId'-  nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)-  nub <$> go eId+  IntSet.toList <$> go [eId] IntSet.empty    where-    go :: Monad m => Int -> EGraphST m [Int]+    hasNoTerminal :: [ENode] -> Bool+    hasNoTerminal = all (not . null . childrenOf) +    getNodes :: Monad m => EClassId -> EGraphST m [ENode]+    getNodes n = gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)++    go :: Monad m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet+    go [] visited = pure visited+    go queue visited = do +        nodes <- concatMap childrenOf . concat . filter hasNoTerminal <$> mapM getNodes queue+        eids <- filter (\e -> e `IntSet.notMember` visited) <$> (mapM canonical nodes)+        go eids (visited `IntSet.union` IntSet.fromList queue)+            {-     go n = do nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)               let hasTerminal = any (null . childrenOf) nodes               eids <- mapM canonical $ concatMap childrenOf nodes@@ -469,7 +485,23 @@                 then pure [n]                 else do eids' <- mapM go eids                         pure ((n : eids) <> concat eids')+                        -} {-# INLINE getAllChildEClasses #-}++getAllChildBestEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildBestEClasses eId' = do+  eId <- canonical eId'+  nub <$> go eId++  where+    go :: Monad m => Int -> EGraphST m [Int]+    go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass)+              let hasTerminal = (null . childrenOf) node+              eids <- mapM canonical $ childrenOf node+              if hasTerminal+                then pure [n]+                else do eids' <- mapM go eids+                        pure ((n : eids) <> concat eids')  -- | returns a random expression rooted at e-class `eId` getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree)
src/Algorithm/EqSat/Egraph.hs view
@@ -201,7 +201,7 @@                         , _consts  :: Consts                         , _fitness :: Maybe Double    -- NOTE: this cannot be NaN                         , _dl      :: Maybe Double-                        , _theta   :: Maybe PVector+                        , _theta   :: [PVector]                         , _size    :: Int                         -- , _properties :: Property                         -- TODO: include evaluation of expression from this e-class
src/Algorithm/EqSat/Info.hs view
@@ -90,7 +90,7 @@      enode' <- canonize enode      cost   <- calculateCost costFun enode'      sz <- sum <$> mapM (\ecId -> gets (_size . _info . (IntMap.! ecId) . _eClass)) (childrenOf enode')-     pure $ EData cost enode' consts Nothing Nothing Nothing (sz+1)+     pure $ EData cost enode' consts Nothing Nothing [] (sz+1)  getChildrenMinHeight :: Monad m => ENode -> EGraphST m Int getChildrenMinHeight enode = do@@ -167,13 +167,13 @@     evalOp' (ConstVal x) (ConstVal y) = ConstVal $ evalOp op x y     evalOp' _            _            = NotConst -insertFitness :: Monad m => EClassId -> Double -> PVector -> EGraphST m ()+insertFitness :: Monad m => EClassId -> Double -> [PVector] -> EGraphST m () insertFitness eId' fit params = do   eId <- canonical eId'   ec <- gets ((IntMap.! eId) . _eClass)   let oldFit  = _fitness . _info $ ec   --when (oldFit < Just fit) $ do-  let newInfo = (_info ec){_fitness = Just fit, _theta = Just params}+  let newInfo = (_info ec){_fitness = Just fit, _theta = params}       newEc   = ec{_info = newInfo}       sz = _size newInfo   modify' $ over eClass (IntMap.insert eId newEc)
src/Algorithm/EqSat/Queries.hs view
@@ -28,6 +28,7 @@ import Data.Sequence ( Seq(..) ) import qualified Data.Sequence as FingerTree import qualified Data.Foldable as Foldable+import Data.SRTree (childrenOf)  import Debug.Trace @@ -98,6 +99,7 @@                                        else if ecId `Set.member` ecs && p ec                                               then go (m-1) (ecId:bests) t                                               else go m bests t+ getTopECLassNotIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopECLassNotIn b n p ecs' = do   let f = if b then _fitRangeDB else _dlRangeDB@@ -151,7 +153,7 @@ getTopFitEClassThat  = getTopECLassThat True getTopDLEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopDLEClassThat   = getTopECLassThat False-getTopFitEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassIn :: Monad m =>  Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopFitEClassIn    = getTopECLassIn True getTopDLEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopDLEClassIn     = getTopECLassIn False@@ -188,3 +190,4 @@                                           then root                                           else (x, eid) :<| root)                                         xs -- (Prelude.filter ((/= eid) . snd) xs)+
src/Algorithm/EqSat/Simplify.hs view
@@ -12,7 +12,7 @@ -- Module containing the algebraic rules and simplification function. -- ------------------------------------------------------------------------------module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple ) where+module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple, rewritesWithConstant ) where  import Algorithm.EqSat (eqSat, applySingleMergeOnlyEqSat) import Algorithm.EqSat.Egraph
src/Algorithm/SRTree/ConfidenceIntervals.hs view
@@ -115,7 +115,7 @@     getResStdError row = sqrt $ (A.!.!) row $ A.fromList compMode $ map (row A.!.!) covs     resStdErr          = map getResStdError jac -predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat $ take 10 xss'+predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat xss' -- $ take 10 xss'   where     yhat     = A.toList $ predFun xss     theta'   = A.toStorableVector theta@@ -134,6 +134,7 @@  -- inverse function of the distributions  inverseDist :: Floating p => Distribution -> p -> p+inverseDist MSE y = y inverseDist Gaussian y  = y inverseDist Bernoulli y = log (y/(1-y)) inverseDist Poisson y   = log y
src/Algorithm/SRTree/Likelihoods.hs view
@@ -128,7 +128,7 @@     s       = theta M.! (p' - 1)     (Sz m') = M.size ys      (Sz p') = M.size theta-    nParams = countParams t+    nParams = countParamsUniq t     m       = fromIntegral m'     p       = fromIntegral p' @@ -147,7 +147,7 @@ -- y log phi + (1-y) log (1 - phi), assuming y \in {0,1} nll Bernoulli _ xss ys tree theta   | notValid ys = error "For Bernoulli distribution the output must be either 0 or 1."-  | otherwise   = negate . M.sum $ delay ys * yhat - log (M.map (1+) $ exp yhat)+  | otherwise   = M.sum $ (M.map (1-) (delay ys)) * yhat + log (M.map (1+) $ exp (M.map negate yhat))   where     (Sz m)   = M.size ys     yhat     = evalTree xss theta tree@@ -171,7 +171,7 @@     (Sz p')      = M.size theta     (Sz2 m n)    = M.size xss     p            = fromIntegral p'-    num_params   = countParams tree+    num_params   = countParamsUniq tree      x0           = xss <! 0     logX         = xss <! 1@@ -224,13 +224,13 @@ buildNLL Gaussian m tree =  (square(tree - var (-1)) / square (param p)) + log ((square (param p)))   where     square = Fix . Uni Square-    p = countParams tree+    p = countParamsUniq tree buildNLL HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2)) buildNLL Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree buildNLL Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree buildNLL ROXY m tree = neglogP   where-    p = countParams tree+    p = countParamsUniq tree     f = log (abs tree) / log 10     fprime = deriveByVar 0 tree / (log 10 * tree) * var 0 * log 10     logX         = var 1@@ -458,6 +458,7 @@     res    = delay ys - phi      (phi, phi') = case dist of+                    MSE       -> (yhat, M.replicate cmp (Sz m) 1)                     Gaussian  -> (yhat, M.replicate cmp (Sz m) 1)                     Bernoulli -> (logistic yhat, phi*(M.replicate cmp (Sz m) 1 - phi))                     Poisson   -> (exp yhat, phi)
src/Algorithm/SRTree/Opt.hs view
@@ -35,7 +35,7 @@   | n == 0     = (t0, f, 0)   | otherwise  = (t_opt', nll dist mYerr xss ys tree t_opt', nEvs)   where-    tree'      = buildNLL dist (fromIntegral m) $ relabelParams $ tree -- convertProtectedOps+    tree'      = buildNLL dist (fromIntegral m) tree -- convertProtectedOps     t0'        = toStorableVector t0     treeArr    = IntMap.toAscList $ tree2arr tree'     j2ix       = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..]@@ -70,7 +70,7 @@   | n == 0     = t0   | otherwise  = t_opt'   where-    tree'      = buildNLL dist (fromIntegral m) $ relabelParams tree+    tree'      = buildNLL dist (fromIntegral m) tree -- relabelParams     t0'        = toStorableVector t0     treeArr    = IntMap.toAscList $ tree2arr tree'     j2ix       = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..]
src/Data/SRTree.hs view
@@ -26,11 +26,13 @@          , countVarNodes          , countConsts          , countParams+         , countParamsUniq          , countOccurrences          , countUniqueTokens          , numberOfVars          , getIntConsts          , relabelParams+         , relabelParamsOrder          , relabelVars          , constsToParam          , floatConstsToParam@@ -57,11 +59,13 @@          , countVarNodes          , countConsts          , countParams+         , countParamsUniq          , countOccurrences          , countUniqueTokens          , numberOfVars          , getIntConsts          , relabelParams+         , relabelParamsOrder          , relabelVars          , constsToParam          , floatConstsToParam
src/Data/SRTree/Internal.hs view
@@ -32,11 +32,13 @@          , countVarNodes          , countConsts          , countParams+         , countParamsUniq          , countOccurrences          , countUniqueTokens          , numberOfVars          , getIntConsts          , relabelParams+         , relabelParamsOrder          , relabelVars          , constsToParam          , floatConstsToParam@@ -47,11 +49,13 @@          )          where -import Control.Monad.State (MonadState (get), State, evalState, modify)+import Control.Monad.State (MonadState (get), State, evalState, modify, put) import Data.SRTree.Recursion (Fix (..), cata, cataM) import qualified Data.Set as S import Data.String (IsString (..)) import Text.Read (readMaybe)+import qualified Data.IntMap as IntMap+import Data.List ( nub )  -- | Tree structure to be used with Symbolic Regression algorithms. -- This structure is a fixed point of a n-ary tree. @@ -356,6 +360,20 @@       alg (Bin _ l r) = 0 + l + r {-# INLINE countParams #-} +-- | Count the unique occurrences of `Param` nodes+--+-- >>> countParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0"+-- 2+countParamsUniq :: Fix SRTree -> Int+countParamsUniq t = length . nub $ cata alg t+  where+      alg Var {} = []+      alg (Param ix) = [ix]+      alg Const {} = []+      alg (Uni _ t) = t+      alg (Bin _ l r) = l <> r+{-# INLINE countParamsUniq #-}+ -- | Count the number of const nodes -- -- >>> countConsts $ "x0"* 2 + 3 * sin "x0"@@ -446,6 +464,34 @@       alg :: SRTree (Fix SRTree) -> State Int (Fix SRTree)       alg (Var ix)    = pure $ var ix       alg (Param ix)  = do iy <- get; modify (+1); pure (param iy)+      alg (Const c)   = pure $ Fix $ Const c+      alg (Uni f t)   = pure $ Fix (Uni f t)+      alg (Bin f l r) = pure $ Fix (Bin f l r)++-- | Reorder the labels of the parameters indices+--+-- >>> showExpr . relabelParamsOrder $ "x0" + "t1" * sin ("t3" + "x1") - "t1"+-- "x0" + "t0" * sin ("t1" + "x1") - "t0"+relabelParamsOrder :: Fix SRTree -> Fix SRTree+relabelParamsOrder t = cataM leftToRight alg t `evalState` (IntMap.empty, 0)+  where+      -- | leftToRight (left to right) defines the sequence of processing+      leftToRight (Uni f mt)    = Uni f <$> mt;+      leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+      leftToRight (Var ix)      = pure (Var ix)+      leftToRight (Param ix)    = pure (Param ix)+      leftToRight (Const c)     = pure (Const c)++      -- | any time we reach a Param ix, it replaces ix with current state+      -- and increments one to the state.+      alg :: SRTree (Fix SRTree) -> State (IntMap.IntMap Int, Int) (Fix SRTree)+      alg (Var ix)    = pure $ var ix+      alg (Param ix)  = do (m, iy) <- get+                           if IntMap.member ix m+                              then pure (param $ m IntMap.! ix)+                              else do let m' = IntMap.insert ix iy m+                                      put (m', iy+1)+                                      pure (param iy)       alg (Const c)   = pure $ Fix $ Const c       alg (Uni f t)   = pure $ Fix (Uni f t)       alg (Bin f l r) = pure $ Fix (Bin f l r)
src/Data/SRTree/Print.hs view
@@ -103,7 +103,7 @@   where     alg = \case       Var ix        -> concat ["x[:, ", show ix, "]"]-      Param ix      -> concat ["t[:, ", show ix, "]"]+      Param ix      -> concat ["t[", show ix, "]"]       Const c       -> show c       Bin Power l r -> concat [l, " ** ", r]       Bin op l r    -> concat ["(", l, " ", showOp op, " ", r, ")"]
srtree.cabal view
@@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack  name:           srtree-version:        2.0.0.4+version:        2.0.1.2 synopsis:       A general library to work with Symbolic Regression expression trees. description:    A Symbolic Regression Tree data structure to work with mathematical expressions with support to first order derivative and simplification; category:       Math, Data, Data Structures@@ -57,8 +57,6 @@   hs-source-dirs:       src   ghc-options: -fwarn-incomplete-patterns -threaded-  extra-lib-dirs:-      /usr/local/lib   extra-libraries:       nlopt   build-depends:@@ -75,7 +73,7 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <1.1+    , massiv >=1.0.4.1 && <1.1     , mtl >=2.2 && <2.4     , random >=1.2 && <1.4     , scheduler >=2.0.0.1 && <3@@ -89,14 +87,14 @@     , zlib >=0.6.3 && <0.8   default-language: Haskell2010 -executable egraphGP+executable eggp   main-is: Main.hs   other-modules:       Random       Util       Paths_srtree   hs-source-dirs:-      apps/egraphGP+      apps/eggp   ghc-options: -threaded -rtsopts -with-rtsopts=-N   build-depends:       attoparsec >=0.14.4 && <0.15@@ -112,7 +110,7 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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@@ -128,54 +126,19 @@     , zlib >=0.6.3 && <0.8   default-language: Haskell2010 -executable egraphSearch+executable reggression   main-is: Main.hs   other-modules:+      Commands       Random       Util       Paths_srtree   hs-source-dirs:-      apps/egraphSearch-  ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall-  build-depends:-      attoparsec >=0.14.4 && <0.15-    , attoparsec-expr >=0.1.1.2 && <0.2-    , base >=4.19 && <5-    , 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.0 && <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-    , srtree-    , 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.*-    , vector >=0.12 && <0.14-    , zlib >=0.6.3 && <0.8-  default-language: Haskell2010--executable eqsatrepr-  main-is: Main.hs-  other-modules:-      Paths_srtree-  hs-source-dirs:-      apps/eqsatrepr+      apps/rEGGression   ghc-options: -threaded -rtsopts -with-rtsopts=-N   build-depends:-      attoparsec >=0.14.4 && <0.15+      ansi-terminal >=1.1.2 && <1.2+    , attoparsec >=0.14.4 && <0.15     , attoparsec-expr >=0.1.1.2 && <0.2     , base >=4.19 && <5     , binary >=0.8.9.1 && <0.9@@ -188,13 +151,16 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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+    , repline >=0.4.2.0 && <0.5     , scheduler >=2.0.0.1 && <3     , split >=0.2.5 && <0.3     , srtree     , statistics >=0.16.2.1 && <0.17+    , table-layout >=1.0.0.1 && <1.1     , transformers >=0.6.1.0 && <0.7     , unliftio >=0.2.10 && <1     , unliftio-core >=0.2.1 && <1@@ -203,19 +169,15 @@     , zlib >=0.6.3 && <0.8   default-language: Haskell2010 -executable reggression+executable srsimplify   main-is: Main.hs   other-modules:-      Commands-      Random-      Util       Paths_srtree   hs-source-dirs:-      apps/rEGGression+      apps/srsimplify   ghc-options: -threaded -rtsopts -with-rtsopts=-N   build-depends:-      ansi-terminal >=1.1.2 && <1.2-    , attoparsec >=0.14.4 && <0.15+      attoparsec >=0.14.4 && <0.15     , attoparsec-expr >=0.1.1.2 && <0.2     , base >=4.19 && <5     , binary >=0.8.9.1 && <0.9@@ -228,16 +190,14 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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-    , repline >=0.4.2.0 && <0.5     , scheduler >=2.0.0.1 && <3     , split >=0.2.5 && <0.3     , srtree     , statistics >=0.16.2.1 && <0.17-    , table-layout >=1.0.0.1 && <1.1     , transformers >=0.6.1.0 && <0.7     , unliftio >=0.2.10 && <1     , unliftio-core >=0.2.1 && <1@@ -246,12 +206,15 @@     , zlib >=0.6.3 && <0.8   default-language: Haskell2010 -executable srsimplify+executable srtools   main-is: Main.hs   other-modules:+      Args+      IO+      Report       Paths_srtree   hs-source-dirs:-      apps/srsimplify+      apps/srtools   ghc-options: -threaded -rtsopts -with-rtsopts=-N   build-depends:       attoparsec >=0.14.4 && <0.15@@ -267,7 +230,7 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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@@ -283,16 +246,15 @@     , zlib >=0.6.3 && <0.8   default-language: Haskell2010 -executable srtools+executable symregg   main-is: Main.hs   other-modules:-      Args-      IO-      Report+      Random+      Util       Paths_srtree   hs-source-dirs:-      apps/srtools-  ghc-options: -threaded -rtsopts -with-rtsopts=-N+      apps/symregg+  ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall   build-depends:       attoparsec >=0.14.4 && <0.15     , attoparsec-expr >=0.1.1.2 && <0.2@@ -307,12 +269,12 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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+    , split >=0.2.5     , srtree     , statistics >=0.16.2.1 && <0.17     , transformers >=0.6.1.0 && <0.7@@ -347,7 +309,7 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <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@@ -387,7 +349,7 @@     , ieee754 >=0.8.0 && <0.9     , lens >=5.2.3 && <5.4     , list-shuffle >=1.0.0.1 && <1.1-    , massiv >=1.0.4.0 && <1.1+    , massiv >=1.0.4.1 && <1.1     , mtl >=2.2 && <2.4     , random >=1.2 && <1.4     , scheduler >=2.0.0.1 && <3