diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,9 @@ # Changelog for srtree +## 2.0.1.6 + +- Added Fractional Bayes model selection + ## 2.0.1.5 - Fix `refit` to only replace the fitness if it improves the fitness diff --git a/apps/srtools/Args.hs b/apps/srtools/Args.hs --- a/apps/srtools/Args.hs +++ b/apps/srtools/Args.hs @@ -24,6 +24,7 @@ , toScreen :: Bool , useProfile :: Bool , simple :: Bool + , sigma :: Double , alpha :: Double , ptype :: PType } deriving Show @@ -122,6 +123,12 @@ <*> switch ( long "simple" <> help "If set, calculates only SSE.") + <*> option auto + ( long "sigma" + <> metavar "SIGMA" + <> showDefault + <> value 0.001 + <> help "Estimation of error for Guassian distribution.") <*> option auto ( long "alpha" <> metavar "ALPHA" diff --git a/apps/srtools/IO.hs b/apps/srtools/IO.hs --- a/apps/srtools/IO.hs +++ b/apps/srtools/IO.hs @@ -14,7 +14,7 @@ import qualified Data.SRTree.Print as P import Data.SRTree.Eval ( compMode ) -import Args ( Args(outfile, alpha,dist,niter) ) +import Args ( Args(outfile, alpha,dist,niter,sigma) ) import Report import Data.SRTree.Recursion ( cata ) @@ -44,7 +44,10 @@ -> (BasicInfo, SSE, SSE, Info, (BasicStats, [CI], [CI], [CI], [CI])) processTree args seed dset t ix = (basic, sseOrig, sseOpt, info, cis) where - (tree, theta0) = floatConstsToParam t + (tree, theta0') = floatConstsToParam t + theta0 = if dist args == Gaussian + then theta0' <> [sigma args] + else theta0' basic = getBasicStats args seed dset tree theta0 ix treeVal = case (_xVal dset, _yVal dset) of @@ -64,7 +67,10 @@ -> (BasicInfo, SSE, SSE) processTreeSimple args seed dset t ix = (basic, sseOrig, sseOpt) where - (tree, theta0) = floatConstsToParam t + (tree, theta0') = floatConstsToParam t + theta0 = if dist args == Gaussian + then theta0' <> [sigma args] + else theta0' basic = getBasicStats args seed dset tree theta0 ix treeVal = case (_xVal dset, _yVal dset) of diff --git a/apps/srtools/Main.hs b/apps/srtools/Main.hs --- a/apps/srtools/Main.hs +++ b/apps/srtools/Main.hs @@ -14,13 +14,14 @@ args <- execParser opts g <- getStdGen (dset, varnames, tgname) <- getDataset args + let seed = if rseed args < 0 then g else mkStdGen (rseed args) varnames' = map unpack $ split ',' $ pack varnames withInput (infile args) (from args) varnames False (simpl args) >>= if toScreen args - then printResultsScreen args seed dset varnames' tgname -- full report on screne + then printResultsScreen args seed dset varnames' tgname -- full report on screen else if simple args then printResultsSimple args seed dset varnames' -- csv file else printResults args seed dset varnames' -- csv file diff --git a/apps/srtools/Report.hs b/apps/srtools/Report.hs --- a/apps/srtools/Report.hs +++ b/apps/srtools/Report.hs @@ -185,10 +185,16 @@ (Nothing, _) -> (xTr, yTr) (_, Nothing) -> (xTr, yTr) (Just a, Just b) -> (a, b) - (tOpt, thetaOpt) = floatConstsToParam tree + (tOpt, thetaOpt_nosig) = floatConstsToParam tree + thetaOpt = if dist args == Gaussian + then thetaOpt_nosig <> [sigma args] + else thetaOpt_nosig thetaOpt' = A.fromList compMode thetaOpt - (tOptVal, thetaOptVal) = floatConstsToParam treeVal + (tOptVal, thetaOptVal_nosig) = floatConstsToParam treeVal + thetaOptVal = if dist args == Gaussian + then thetaOptVal_nosig <> [sigma args] + else thetaOptVal_nosig thetaOptVal' = A.fromList compMode thetaOptVal dist' = dist args diff --git a/src/Algorithm/EqSat.hs b/src/Algorithm/EqSat.hs --- a/src/Algorithm/EqSat.hs +++ b/src/Algorithm/EqSat.hs @@ -137,7 +137,7 @@ -- if nothing changed, return if it == 1 || (eNodes' == eNodes && eClasses' == eClasses) then pure (True, it) - else if IntMap.size eClasses' > 500 -- maximum allowed number of e-classes. TODO: customize + else if IntMap.size eClasses' > 1500 -- maximum allowed number of e-classes. TODO: customize then pure (False, it) else go (it-1) sch' diff --git a/src/Algorithm/EqSat/Build.hs b/src/Algorithm/EqSat/Build.hs --- a/src/Algorithm/EqSat/Build.hs +++ b/src/Algorithm/EqSat/Build.hs @@ -30,6 +30,7 @@ import qualified Data.Map.Strict as Map import qualified Data.HashSet as Set import Control.Monad.State.Strict +import Control.Monad.Identity import Data.SRTree.Recursion (cataM) import Algorithm.EqSat.Info import qualified Data.IntSet as IntSet @@ -271,7 +272,13 @@ -- | `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 +addToDB enode' eid = do + eid' <- canonical eid + isConst <- gets (_consts . _info . (IntMap.! eid') . _eClass) + let enode = case isConst of + ConstVal x -> Const x + ParamIx x -> Param x + _ -> enode' let ids = eid : childrenOf enode -- we will add the e-class id and the children ids op = getOperator enode -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB trie <- gets ((Map.!? op) . _patDB . _eDB) -- gets the entry for op, if it exists @@ -390,7 +397,12 @@ fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) [] {-# INLINE fromTrees #-} +countParamsEg :: EGraph -> EClassId -> Int +countParamsEg eg rt = countParams . runIdentity $ getBestExpr rt `evalStateT` eg +countParamsUniqEg :: EGraph -> EClassId -> Int +countParamsUniqEg eg rt = countParamsUniq . runIdentity $ getBestExpr rt `evalStateT` eg + -- | gets the best expression given the default cost function getBestExpr :: Monad m => EClassId -> EGraphST m (Fix SRTree) getBestExpr eid = do eid' <- canonical eid @@ -459,6 +471,42 @@ ts <- go ns pure (t:ts) {-# INLINE getAllExpressionsFrom #-} + +getNExpressionsFrom :: Monad m => Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom n eId' = getNExpressionsFrom' n 15 eId' + +getNExpressionsFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom' _ 0 _ = pure [] +getNExpressionsFrom' n d eId' = do + eId <- canonical eId' + nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass) + (concat <$> go n d nodes) + where + isTerm (Var _) = True + isTerm (Const _) = True + isTerm (Param _) = True + isTerm _ = False + toTree (Var ix) = Fix $ Var ix + toTree (Const x) = Fix $ Const x + toTree (Param ix) = Fix $ Param ix + toTree _ = undefined + + go n' _ [] = pure [] + go n' 0 ts = pure [] + go n' d (node:ns) = do + tt <- Prelude.map Fix <$> case node of + Bin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l + r' <- getNExpressionsFrom' n' (d-1) r + pure $ Prelude.take n [Bin op li ri | li <- l', ri <- r'] + Uni f t -> Prelude.map (Uni f) <$> getNExpressionsFrom' n' (d-1) t + Var ix -> pure [Var ix] + Const x -> pure [Const x] + Param ix -> pure [Param ix] + let n'' = n' - length tt + if n'' <= 0 + then pure [tt] + else do ts <- go n'' (d-1) ns + pure (tt:ts) getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId] getAllChildEClasses eId' = do diff --git a/src/Algorithm/EqSat/Egraph.hs b/src/Algorithm/EqSat/Egraph.hs --- a/src/Algorithm/EqSat/Egraph.hs +++ b/src/Algorithm/EqSat/Egraph.hs @@ -53,6 +53,7 @@ type EGraphST m a = StateT EGraph m a type Cost = Int type CostFun = SRTree Cost -> Cost +type ECache = IntMap.IntMap PVector instance Hashable ENode where hashWithSalt n enode = hashWithSalt n (encodeEnode enode) diff --git a/src/Algorithm/EqSat/SearchSRCache.hs b/src/Algorithm/EqSat/SearchSRCache.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/SearchSRCache.hs @@ -0,0 +1,244 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Search +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Support functions for search symbolic expressions with e-graphs +-- +----------------------------------------------------------------------------- + +module Algorithm.EqSat.SearchSRCache where + +import Data.SRTree +import Data.SRTree.Datasets +import System.Random +import Control.Monad.State.Strict +import Algorithm.EqSat.Egraph +import Algorithm.SRTree.Likelihoods +import qualified Data.IntMap as IM +import qualified Data.IntSet as IntSet +import qualified Data.SRTree.Random as Random +import Data.Function ( on ) +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.NonlinearOpt +import Control.Monad ( when, replicateM, forM, forM_ ) +import Algorithm.EqSat.Egraph +import Algorithm.SRTree.Opt +import Algorithm.EqSat.Info +import Algorithm.EqSat.Build +import Data.Maybe ( fromJust ) +import Data.SRTree.Random +import Algorithm.EqSat.Queries +import Data.List ( maximumBy ) +import qualified Data.Map.Strict as Map +import Control.Monad.Identity + +import Debug.Trace + +-- Environment of an e-graph with support to random generator and IO +type RndEGraph a = EGraphST (StateT StdGen (StateT [ECache] IO)) a + +io :: IO a -> RndEGraph a +io = lift . lift . lift +{-# INLINE io #-} +getCache :: StateT [ECache] IO a -> RndEGraph a +getCache = lift . lift +rnd :: StateT StdGen (StateT [ECache] 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 t +while p arg prog = do if (p arg) + then do arg' <- prog arg + while p arg' prog + else pure arg + +fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> EGraph -> EClassId -> ECache -> PVector -> (Double, PVector, ECache) +fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) egraph root cache thetaOrig = + if isNaN val -- || isNaN tr + then (-(1/0), theta,cache') -- infinity + else (val, theta, cache') + where + tree = runIdentity $ getBestExpr root `evalStateT` egraph + nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 + (theta, val, _, cache') = minimizeNLLEGraph VAR1 distribution mYErr nIter x y egraph root cache thetaOrig + evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta + -- val = evalF x_val y_val mYErr_val + +--{-# INLINE fitnessFun #-} + +fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> EClassId -> ECache -> RndEGraph (Double, PVector, ECache) +fitnessFunRep nRep nIter distribution dataTrain dataVal root cache = do + egraph <- get + let nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 + fst' (a, _, _) = a + thetaOrigs <- replicateM nRep (rnd $ randomVec nParams) + let fits = maximumBy (compare `on` fst') $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal egraph root cache) thetaOrigs + pure fits +--{-# INLINE fitnessFunRep #-} + + +fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector]) +fitnessMV shouldReparam nRep nIter distribution dataTrainsVals root = do + -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree + -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm' + caches <- getCache get + response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache + getCache $ put (Prelude.map trd response) + pure (minimum (Prelude.map fst' response), Prelude.map snd' response) + where fst' (a, _, _) = a + snd' (_, a, _) = a + trd (_, _, a) = a + +fitnessMVNoCache :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector]) +fitnessMVNoCache shouldReparam nRep nIter distribution dataTrainsVals root = do + -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree + -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm' + caches <- getCache get + response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache + pure (minimum (Prelude.map fst' response), Prelude.map snd' response) + where fst' (a, _, _) = a + snd' (_, a, _) = a + trd (_, _, a) = a + + + +-- 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 + 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 ec + 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) +getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize 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 ec + mf <- getFitness ec + case mf of + Nothing -> insertFitness ec f p + Just f' -> when (f > f') $ 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 [[String]] + go n ix f + | n > maxSize = pure [] + | otherwise = do + ecList <- getBestExprWithSize n + if not (null ecList) + then do let (ec, mf) = head ecList + f' = fromJust mf + improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f' + ec' <- canonical ec + if improved + then do refit fitFun ec' + t <- printExprFun ix ec' + ts <- go (n+1) (ix + if improved then 1 else 0) (max f f') + pure (t:ts) + else go (n+1) (ix + if improved then 1 else 0) (max f f') + 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 c + insertFitness c f p + +evaluateRndUnevaluated fitFun = do + ec <- gets (IntSet.toList . _unevaluated . _eDB) + c <- rnd . randomFrom $ ec + --t <- getBestExpr c + (f, p) <- fitFun c + insertFitness c f p + pure c + +-- | check whether an e-node exists or does not exist in the e-graph +doesExist, doesNotExist :: ENode -> RndEGraph Bool +doesExist en = gets ((Map.member en) . _eNodeToEClass) +doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass) + +-- | check whether the partial tree defined by a list of ancestors will create +-- a non-existent expression when combined with a certain e-node. +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' + +-- | check whether combining a partial tree `parent` with the e-node `en'` +-- will create a new expression +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'' diff --git a/src/Algorithm/EqSat/Simplify.hs b/src/Algorithm/EqSat/Simplify.hs --- a/src/Algorithm/EqSat/Simplify.hs +++ b/src/Algorithm/EqSat/Simplify.hs @@ -12,7 +12,7 @@ -- Module containing the algebraic rules and simplification function. -- ----------------------------------------------------------------------------- -module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple, rewritesWithConstant ) where +module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple, rewritesWithConstant, myCost ) where import Algorithm.EqSat (eqSat, applySingleMergeOnlyEqSat) import Algorithm.EqSat.Egraph @@ -103,6 +103,7 @@ --, ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x" --, (powabs "x" "y") * (powabs "x" "z") :=> powabs "x" ("y" + "x") , ("x" + "y") + "z" :=> "x" + ("y" + "z") + , ("x" + "y") - "z" :=> "x" + ("y" - "z") --, ("x" + "y") - "z" :=> "x" + ("y" - "z") -- TODO: check that I don't need that , ("x" * "y") * "z" :=> "x" * ("y" * "z") , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z") @@ -119,6 +120,9 @@ -- , "a" * (("x" * "y") + ("z" * "w")) :=> ("a" * "x") * ("y" + ("z" / "x") * "w") :| isConstPt "a" :| isConstPt "x" :| isConstPt "z" :| isNotZero "x" , (("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x" , (("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x" :| isConstPt "z" + , "x" * "x" :=> "x" ** 2 + , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 + , "x" ** 2 + "x" * "y" :=> "x" * ("x" + "y") -- , "x" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL -- , "x" + "y" * "z" :=> "y" * ("x" * "y" ** (-1) + "z") :| isNotZero "y" -- GABRIEL ] @@ -148,6 +152,7 @@ --, recip "x" :==: "x" ** (-1) -- GABRIEL --, "x" / "y" :==: "x" * "y" ** (-1) -- GABRIEL , abs "x" ** "y" :=> "x" ** "y" :| isEven "y" + , sqrt ("x" * "x") :=> abs "x" ] -- Rules that reduces redundant parameters diff --git a/src/Algorithm/SRTree/AD.hs b/src/Algorithm/SRTree/AD.hs --- a/src/Algorithm/SRTree/AD.hs +++ b/src/Algorithm/SRTree/AD.hs @@ -22,8 +22,10 @@ module Algorithm.SRTree.AD ( reverseModeArr + , reverseModeEGraph , reverseModeGraph , forwardModeUniqueJac + , evalCache ) where import Control.Monad (forM_, foldM, when) @@ -47,16 +49,228 @@ import Data.List ( foldl' ) import qualified Data.Vector.Storable as VS import Control.Scheduler -import Data.Maybe ( fromJust ) +import Data.Maybe ( fromJust, isJust ) +import Algorithm.EqSat.Egraph import Control.Monad.State.Strict +import Control.Monad.Identity --import UnliftIO.Async import qualified Data.Map.Strict as Map +evalCache :: SRMatrix -> EGraph -> ECache -> EClassId -> VS.Vector Double -> ECache +evalCache xss egraph cache root' theta = cache' + where + (Sz2 _ m') = M.size xss + m = Sz1 m' + root = canon root' + p = VS.length theta + comp = M.getComp xss + one :: Array S Ix1 Double + one = M.replicate comp m 1 + + canon rt = case _canonicalMap egraph IntMap.!? rt of + Nothing -> error "wrong canon" + Just rt' -> if rt == rt' then rt else canon rt' + + getNode rt' = let rt = canon rt' + cls = _eClass egraph IntMap.! rt + in (_best . _info) cls + + getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph + in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n' + + ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty) + where + evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalCached rt = insertKey rt + + insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + insertKey key' = do + let key = canon key' + isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst) + isCachedLocal <- gets ((key `IntMap.member`) . snd . fst) + when (not isCachedLocal && not isCachedGlobal) $ do + let node = getNode key + (ev, toLocal) <- evalKey node + modify' (insKey node ev toLocal) + getVal key + + evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalKey (Var ix) = pure $ (M.computeAs S $ xss evalKey n + Const v -> evalKey n + Param ix -> evalKey n + _ -> getFromCache rt + getFromCache rt = do + global <- gets ((IntMap.!? rt) . fst . fst) + local <- gets ((IntMap.!? rt) . snd . fst) + if | isJust global -> pure (fromJust global, False) + | isJust local -> pure (fromJust local, True) + | otherwise -> insertKey rt + +-- reverse mode applied directly on an e-graph. Supports caching. +-- assumes root points to the loss function, so for an expression +-- f(x) and the loss (y - (f(x))^2), root will point to "^" +reverseModeEGraph :: SRMatrix -> PVector -> Maybe PVector -> EGraph -> ECache -> EClassId -> VS.Vector Double -> (Array D Ix1 Double, VS.Vector Double) +reverseModeEGraph xss ys mYErr egraph cache root' theta = + (delay $ rootVal + , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]] + ) + where + rootVal = extractCache (cache'' IntMap.!? root', localcache' IntMap.!? root') + root = canon root' + yErr = fromJust mYErr + m = M.size ys + p = VS.length theta + comp = M.getComp xss + one :: Array S Ix1 Double + one = M.replicate comp m 1 + + canon rt = case _canonicalMap egraph IntMap.!? rt of + Nothing -> error "wrong canon" + Just rt' -> if rt == rt' then rt else canon rt' + + getNode rt' = let rt = canon rt' + cls = _eClass egraph IntMap.! rt + in (_best . _info) cls + + getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph + in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n' + + ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty) + where + evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalCached rt = insertKey rt + + insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + insertKey key' = do + let key = canon key' + isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst) + isCachedLocal <- gets ((key `IntMap.member`) . snd . fst) + when (not isCachedLocal && not isCachedGlobal) $ do + let node = getNode key + (ev, toLocal) <- evalKey node + modify' (insKey node ev toLocal) + getVal key + + evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalKey (Var ix) = pure $ if | ix == -1 -> (ys, False) + | ix == -2 -> (yErr, False) + | otherwise -> (M.computeAs S $ xss evalKey n + Const v -> evalKey n + Param ix -> evalKey n + _ -> getFromCache rt + getFromCache rt = do + global <- gets ((IntMap.!? rt) . fst . fst) + local <- gets ((IntMap.!? rt) . snd . fst) + if | isJust global -> pure (fromJust global, False) + | isJust local -> pure (fromJust local, True) + | otherwise -> insertKey rt + + extractCache (Nothing, Nothing) = error "no root info" + extractCache (Just r, _) = r + extractCache (_, Just r) = r + + ((cache'', localcache'), cachedGrad) = calcGrad root one `execState` ((cache', localcache), Map.empty) + + calcGrad :: Int -> Array S Ix1 Double -> State ((IntMap.IntMap (Array S Ix1 Double), IntMap.IntMap (Array S Ix1 Double)), Map.Map (SRTree Int) (Array S Ix1 Double)) () + calcGrad rt v = do let node = getNode rt + case node of + Bin op l r -> do xl <- fst <$> getVal l + xr <- fst <$> getVal r + (dl, dr) <- diff op v xl xr l r + calcGrad l dl + calcGrad r dr + Uni f t -> do x <- fst <$> getVal t + calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x)) + Param ix -> modify' (insertGrad v (Param ix)) + _ -> pure () + where + insertGrad v k ((a, b), g) = ((a, b), Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g) + + --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double) + diff Add dx fx gy l r = pure (dx, dx) + diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx) + diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx) + diff Div dx fx gy l r = do + let k = getId (Bin Div l r) + v <- fst <$> getVal k + pure (M.computeAs S $ M.zipWith (/) dx gy + , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy)) + diff Power dx fx gy l r = do + let k = getId (Bin Power l r) + v <- fst <$> getVal k + pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v) + + diff PowerAbs dx fx gy l r = do + let k = getId (Bin PowerAbs l r) + v <- fst <$> getVal k + let v2 = M.map abs fx + v3 = M.computeAs S $ M.zipWith (*) fx gy + pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2 + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v) + + diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx + dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl + in pure (M.computeAs S $ dxl, M.computeAs S $ dxy) + + fixNaN x = if isNaN x then 0 else x + + reverseModeGraph :: SRMatrix -> PVector -> Maybe PVector -> VS.Vector Double -> Fix SRTree -> (Array D Ix1 Double, VS.Vector Double) -reverseModeGraph xss ys mYErr theta tree = (delay $ cachedVal IntMap.! root +reverseModeGraph xss ys mYErr theta tree = (delay $ cachedVal' IntMap.! root , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]]) where yErr = fromJust mYErr @@ -87,9 +301,12 @@ graph (a, _, _, _) = a insKey key ev (a, b, c, d) = (Map.insert key d a, IntMap.insert d key b, IntMap.insert d ev c, d+1) + -- get the values from the cache getVal key (a, b, c, d) = c IntMap.! key + -- maps the the struct to an integer key getKey key (a, b, c, d) = a Map.! key + -- this tells the order in which we traverse the tree leftToRight (Uni f mt) = Uni f <$> mt; leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr leftToRight (Var ix) = pure (Var ix) diff --git a/src/Algorithm/SRTree/Likelihoods.hs b/src/Algorithm/SRTree/Likelihoods.hs --- a/src/Algorithm/SRTree/Likelihoods.hs +++ b/src/Algorithm/SRTree/Likelihoods.hs @@ -24,9 +24,11 @@ , nll , predict , buildNLL + , buildNLLEGraph , gradNLL , gradNLLArr , gradNLLGraph + , gradNLLEGraph , fisherNLL , getSErr , hessianNLL @@ -34,7 +36,7 @@ ) where -import Algorithm.SRTree.AD ( reverseModeArr, reverseModeGraph ) +import Algorithm.SRTree.AD ( reverseModeArr, reverseModeGraph, reverseModeEGraph ) import Data.Massiv.Array hiding (all, map, read, replicate, tail, take, zip) import qualified Data.Massiv.Array as M import qualified Data.Massiv.Array.Mutable as Mut @@ -50,7 +52,14 @@ import Debug.Trace import Data.SRTree.Print +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Simplify +import Algorithm.EqSat.Build +import Control.Monad.State.Strict +import Control.Monad.Identity +import Data.SRTree.Print + -- | Supported distributions for negative log-likelihood -- MSE refers to mean squared error -- HGaussian is Gaussian with heteroscedasticity, where the error should be provided @@ -122,10 +131,10 @@ -- | Gaussian distribution, theta must contain an additional parameter corresponding -- to variance. nll Gaussian mYerr xss ys t theta - | nParams == p' = error "For Gaussian distribution theta must contain the variance as its last value." + | nParams == (p'-1) = error "For Gaussian distribution theta must contain the variance as its last value." | otherwise = 0.5*(sse xss ys t theta / s + m*log (2*pi*s)) where - s = theta M.! (p' - 1) + s = sqrt $ mse xss ys t theta -- theta M.! (p' - 1) (Sz m') = M.size ys (Sz p') = M.size theta nParams = countParamsUniq t @@ -250,6 +259,73 @@ + s2 * (logX - mu_gauss) ** 2 ) / den +buildNLLEGraph MSE m egraph root = runIdentity $ addToEg `runStateT` egraph + where + addToEg :: EGraphST Identity EClassId + addToEg = do v <- add myCost (Var (-1)) + c1 <- add myCost (Const 2) + c2 <- add myCost (Const m) + x <- add myCost (Bin Sub root v) + y <- add myCost (Bin Power x c1) + add myCost (Bin Div y c2) + + +buildNLLEGraph Gaussian m egraph root = runIdentity (addToEg `runStateT` egraph) + where + p = countParamsUniqEg egraph root + addToEg :: EGraphST Identity EClassId + addToEg = do v <- add myCost (Var (-1)) + p <- add myCost (Param p) + sp <- add myCost (Uni Square p) + lsp <- add myCost (Uni Log sp) + d <- add myCost (Bin Sub root v) + sd <- add myCost (Uni Square d) + x <- add myCost (Bin Div sd sp) + add myCost (Bin Add x lsp) + +buildNLLEGraph HGaussian m egraph root = runIdentity $ addToEg `runStateT` egraph + where + addToEg :: EGraphST Identity EClassId + addToEg = do v1 <- add myCost (Var (-1)) + v2 <- add myCost (Var (-2)) + c1 <- add myCost (Const (2*pi)) + c2 <- add myCost (Const m) + x <- add myCost (Bin Sub root v1) + y <- add myCost (Uni Square x) + z <- add myCost (Bin Div y v2) + w <- add myCost (Bin Mul c1 v2) + lw <- add myCost (Uni Log w) + p <- add myCost (Bin Mul c2 lw) + add myCost (Bin Add z p) + + +buildNLLEGraph Poisson m egraph root = runIdentity $ addToEg `runStateT` egraph + where + addToEg :: EGraphST Identity EClassId + addToEg = do v1 <- add myCost (Var (-1)) + lv <- add myCost (Uni Log v1) + x <- add myCost (Bin Mul v1 lv) + y <- add myCost (Uni Exp root) + z <- add myCost (Bin Add x y) + vt <- add myCost (Bin Mul v1 root) + add myCost (Bin Sub z vt) + +buildNLLEGraph Bernoulli m egraph root = runIdentity $ addToEg `runStateT` egraph + where + addToEg :: EGraphST Identity EClassId + addToEg = do v <- add myCost (Var (-1)) + c1 <- add myCost (Const 1) + c2 <- add myCost (Const (-1)) + mr <- add myCost (Bin Mul c2 root) + er <- add myCost (Uni Exp mr) + er1 <- add myCost (Bin Add c1 er) + ler1 <- add myCost (Uni Log er1) + v1 <- add myCost (Bin Sub c1 v) + v1r <- add myCost (Bin Mul v1 root) + add myCost (Bin Add ler1 v1r) + +buildNLLEGraph ROXY m egraph root = error "ROXY not supported with cache" + -- | Prediction for different distributions predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector predict MSE tree theta xss = evalTree xss theta tree @@ -279,28 +355,7 @@ treeArr = IntMap.toAscList $ tree2arr tree' j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] - {- - -- EXAMPLE OF FINITE DIFFERENCE - -- Implement for debugging -gradNLL ROXY mXerr mYerr xss ys tree theta = - (f, delay grad) - where - (Sz p) = M.size theta - (Sz2 m n) = M.size xss - yhat = predict Gaussian tree theta xss - f = nll ROXY mXerr mYerr xss ys tree theta - grad = makeArray @S (getComp xss) (Sz p) finiteDiff - eps = 1e-8 - finiteDiff ix = unsafePerformIO $ do - theta' <- Mut.thaw theta - v <- Mut.readM theta' ix - Mut.writeM theta' ix (v + eps) - theta'' <- Mut.freezeS theta' - let f'= nll ROXY mXerr mYerr xss ys tree theta'' - g = (f' - f)/eps - pure $ if isNaN g then (1/0) else g - -} nanTo0 x = x -- if isNaN x || isInfinite x then 0 else x {-# INLINE nanTo0 #-} @@ -362,6 +417,35 @@ where (yhat, grad) = reverseModeGraph xss ys mYerr theta tree grad' = VS.map nanTo0 grad + +-- | e-graph support +gradNLLEGraph MSE xss ys mYerr egraph cache root theta = + (M.sum yhat, grad') + where + (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta + grad' = VS.map nanTo0 grad +gradNLLEGraph Gaussian xss ys mYerr egraph cache root theta = + (M.sum yhat, grad') + where + (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta + grad' = VS.map nanTo0 grad +gradNLLEGraph Bernoulli xss ys mYerr egraph cache root theta + | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1." + | otherwise = (M.sum yhat, grad') + where + (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta + grad' = VS.map nanTo0 grad +gradNLLEGraph Poisson xss ys mYerr egraph cache root theta + | M.any (<0) ys = error "For Poisson distribution the output must be non-negative." + | otherwise = (M.sum yhat, grad') + where + (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta + grad' = VS.map nanTo0 grad +gradNLLEGraph ROXY xss ys mYerr egraph cache root theta = + ((*0.5) $ M.sum yhat, VS.map (*(0.5)) $ grad') + where + (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta + grad' = VS.map nanTo0 grad -- | Fisher information of negative log-likelihood fisherNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector diff --git a/src/Algorithm/SRTree/ModelSelection.hs b/src/Algorithm/SRTree/ModelSelection.hs --- a/src/Algorithm/SRTree/ModelSelection.hs +++ b/src/Algorithm/SRTree/ModelSelection.hs @@ -18,7 +18,7 @@ import Algorithm.Massiv.Utils ( det ) import Algorithm.SRTree.Likelihoods - ( PVector, SRMatrix, fisherNLL, hessianNLL, nll, Distribution ) + ( PVector, SRMatrix, fisherNLL, hessianNLL, nll, Distribution(..) ) import Data.Massiv.Array (Ix2 (..), Sz (..), (!-!)) import qualified Data.Massiv.Array as A import Data.SRTree @@ -30,7 +30,7 @@ -- | Bayesian information criterion bic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -bic dist mYerr xss ys theta tree = (p + 1) * log n + 2 * nll dist mYerr xss ys tree theta +bic dist mYerr xss ys theta tree = p * log n + 2 * nll dist mYerr xss ys tree theta where (A.Sz (fromIntegral -> p)) = A.size theta (A.Sz (fromIntegral -> n)) = A.size ys @@ -38,7 +38,7 @@ -- | Akaike information criterion aic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -aic dist mYerr xss ys theta tree = 2 * (p + 1) + 2 * nll dist mYerr xss ys tree theta +aic dist mYerr xss ys theta tree = 2 * p + 2 * nll dist mYerr xss ys tree theta where (A.Sz (fromIntegral -> p)) = A.size theta (A.Sz (fromIntegral -> n)) = A.size ys @@ -53,12 +53,25 @@ b = 1 / sqrt n {-# INLINE evidence #-} +fractionalBayesFactor :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +fractionalBayesFactor dist mYerr xss ys theta tree = (1 - b) * nll' - p / 2 * log b + f_compl + p / 2 * log(2*pi*nup) + where + nll_val = nll dist mYerr xss ys tree theta + nll_gaus = nll Gaussian mYerr xss ys tree theta + nll' = if dist == MSE then nll_gaus else nll_val + (A.Sz (fromIntegral -> p)) = A.size theta + (A.Sz (fromIntegral -> n)) = A.size ys + b = 1 / sqrt n + nup = exp(1 - log 3) + f_compl = countNodes tree * log (countUniqueTokens tree) +{-# INLINE fractionalBayesFactor #-} + -- | MDL as described in -- Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023). mdl :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double mdl dist mYerr xss ys theta tree = nll' dist mYerr xss ys theta tree + logFunctional tree - -- + logParameters dist mYerr xss ys theta tree + + logParameters dist mYerr xss ys theta tree where fisher = fisherNLL dist mYerr xss ys tree theta theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher @@ -87,7 +100,7 @@ -- log of the functional complexity logFunctional :: Fix SRTree -> Double -logFunctional tree = countNodes tree * log (countUniqueTokens tree') +logFunctional tree = countNodes tree * log (countUniqueTokens tree') + foldr (\c acc -> log (abs c) + acc) 0 consts + log(2) * numberOfConsts where diff --git a/src/Algorithm/SRTree/Opt.hs b/src/Algorithm/SRTree/Opt.hs --- a/src/Algorithm/SRTree/Opt.hs +++ b/src/Algorithm/SRTree/Opt.hs @@ -23,9 +23,40 @@ import qualified Data.Vector.Storable as VS import qualified Data.IntMap.Strict as IntMap import Data.SRTree.Recursion +import Algorithm.EqSat.Egraph hiding ( size ) +import Algorithm.EqSat.Build +import Control.Monad.State.Strict +import Control.Monad.Identity +import Algorithm.SRTree.AD (evalCache) import Debug.Trace +-- | minimizes the negative log-likelihood of the expression +minimizeNLLEGraph :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> EGraph -> EClassId -> ECache -> PVector -> (PVector, Double, Int, ECache) +minimizeNLLEGraph alg dist mYerr niter xss ys egraph root cache t0 + | niter == 0 = (t0, f, 0, cache') + | n == 0 = (t0, f, 0, cache') + | otherwise = (t_opt', fst aa, nEvs, cache') -- (t_opt', nll dist mYerr xss ys tree t_opt', nEvs, cache') + where + (rt, eg) = buildNLLEGraph dist (fromIntegral m) egraph root -- convertProtectedOps + t0' = toStorableVector t0 + (Sz n) = size t0 + (Sz m) = size ys + tree = runIdentity $ getBestExpr root `evalStateT` egraph + aa = gradNLLEGraph dist xss ys mYerr eg cache' rt t_opt + + funAndGrad = gradNLLEGraph dist xss ys mYerr eg cache' rt + (f, _) = gradNLLEGraph dist xss ys mYerr eg cache' rt t0' -- if there's no parameter or no iterations + cache' = evalCache xss egraph cache root t0' + + + algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) + stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + (t_opt, nEvs) = case minimizeLocal problem t0' of + Right sol -> (solutionParams sol, nEvals sol) + Left e -> (t0', 0) + t_opt' = fromStorableVector compMode t_opt -- | minimizes the negative log-likelihood of the expression diff --git a/src/Data/SRTree/Print.hs b/src/Data/SRTree/Print.hs --- a/src/Data/SRTree/Print.hs +++ b/src/Data/SRTree/Print.hs @@ -22,6 +22,7 @@ , showPython , printPython , showLatex + , showLatexWithVars , printLatex , showOp ) @@ -143,14 +144,34 @@ showLatex = cata alg . removeProtection where alg = \case - Var ix -> concat ["x_{, ", show ix, "}"] - Param ix -> concat ["\\theta_{, ", show ix, "}"] + Var ix -> concat ["x_{", show ix, "}"] + Param ix -> concat ["\\theta_{", show ix, "}"] Const c -> show c - Bin Power l r -> concat [l, "^{", r, "}"] + Bin Power l r -> concat ["{", l, "^{", r, "}}"] + Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"] + Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"] + Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"] Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"] Uni Abs t -> concat ["\\left |", t, "\\right |"] + Uni Recip t -> concat ["\\frac{1}{", t, "}"] Uni f t -> concat [showLatexFun f, "(", t, ")"] - + +showLatexWithVars :: [String] -> Fix SRTree -> String +showLatexWithVars varnames = cata alg . removeProtection + where + alg = \case + Var ix -> concat ["\\operatorname{", varnames !! ix, "}"] + Param ix -> concat ["\\theta_{", show ix, "}"] + Const c -> show c + Bin Power l r -> concat ["{", l, "^{", r, "}}"] + Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"] + Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"] + Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"] + Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"] + Uni Abs t -> concat ["\\left |", t, "\\right |"] + Uni Recip t -> concat ["\\frac{1}{", t, "}"] + Uni f t -> concat [showLatexFun f, "(", t, ")"] + showLatexFun :: Function -> String showLatexFun f = mconcat ["\\operatorname{", map toLower $ show f, "}"] {-# INLINE showLatexFun #-} diff --git a/src/Data/SRTree/Random.hs b/src/Data/SRTree/Random.hs --- a/src/Data/SRTree/Random.hs +++ b/src/Data/SRTree/Random.hs @@ -77,28 +77,28 @@ instance HasFuns FullParams where _funs (P _ _ _ fs) = fs -type Rng a = StateT StdGen IO a +type Rng m a = StateT StdGen m a -- auxiliary function to sample between False and True -toss :: StateT StdGen IO Bool +toss :: Monad m => Rng m Bool toss = state random {-# INLINE toss #-} -tossBiased :: Double -> Rng Bool +tossBiased :: Monad m => Double -> Rng m Bool tossBiased p = do r <- state random pure (r < p) -randomVal :: Rng Double +randomVal :: Monad m => Rng m Double randomVal = state random -- returns a random element of a list -randomFrom :: [a] -> StateT StdGen IO a +randomFrom :: Monad m => [a] -> Rng m a randomFrom funs = do n <- randomRange (0, length funs - 1) pure $ funs !! n {-# INLINE randomFrom #-} -- returns a random element within a range -randomRange :: (Ord val, Random val) => (val, val) -> StateT StdGen IO val +randomRange :: (Ord val, Random val, Monad m) => (val, val) -> Rng m val randomRange rng = state (randomR rng) {-# INLINE randomRange #-} @@ -116,31 +116,31 @@ -- | RndTree is a Monad Transformer to generate random trees of type `SRTree ix val` -- given the parameters `p ix val` using the random number generator `StdGen`. -type RndTree p = ReaderT p (StateT StdGen IO) (Fix SRTree) +type RndTree m p = ReaderT p (StateT StdGen m) (Fix SRTree) -- | Returns a random variable, the parameter `p` must have the `HasVars` property -randomVar :: HasVars p => RndTree p +randomVar :: Monad m => HasVars p => RndTree m p randomVar = do vars <- asks _vars lift $ Fix . Var <$> randomFrom vars -- | Returns a random constant, the parameter `p` must have the `HasConst` property -randomConst :: HasVals p => RndTree p +randomConst :: (HasVals p, Monad m) => RndTree m p randomConst = do rng <- asks _range lift $ Fix . Const <$> randomRange rng -- | Returns a random integer power node, the parameter `p` must have the `HasExps` property -randomPow :: HasExps p => RndTree p +randomPow :: (HasExps p, Monad m) => RndTree m p randomPow = do rng <- asks _exponents lift $ Fix . Bin Power 0 . Fix . Const . fromIntegral <$> randomRange rng -- | Returns a random function, the parameter `p` must have the `HasFuns` property -randomFunction :: HasFuns p => RndTree p +randomFunction :: (HasFuns p, Monad m) => RndTree m p randomFunction = do funs <- asks _funs f <- lift $ randomFrom funs lift $ pure $ Fix (Uni f 0) -- | Returns a random node, the parameter `p` must have every property. -randomNode :: HasEverything p => RndTree p +randomNode :: (HasEverything p, Monad m) => RndTree m p randomNode = do choice <- lift $ randomRange (0, 8 :: Int) case choice of @@ -155,7 +155,7 @@ 8 -> pure . Fix $ Bin Power 0 0 -- | Returns a random non-terminal node, the parameter `p` must have every property. -randomNonTerminal :: HasEverything p => RndTree p +randomNonTerminal :: (HasEverything p, Monad m) => RndTree m p randomNonTerminal = do choice <- lift $ randomRange (0, 6 :: Int) case choice of @@ -173,7 +173,7 @@ -- >>> tree <- evalStateT treeGen (mkStdGen 52) -- >>> showExpr tree -- "(-2.7631152121655838 / Exp((x0 / ((x0 * -7.681722660704317) - Log(3.378309080134594)))))" -randomTreeTemplate :: HasEverything p => Int -> RndTree p +randomTreeTemplate :: (HasEverything p, Monad m) => Int -> RndTree m p randomTreeTemplate 0 = do coin <- lift toss if coin @@ -192,7 +192,7 @@ -- >>> tree <- evalStateT treeGen (mkStdGen 42) -- >>> showExpr tree -- "Exp(Log((((7.784360517385774 * x0) - (3.6412224491658223 ^ x1)) ^ ((x0 ^ -4.09764995657091) + Log(-7.710216839988497)))))" -randomTreeBalanced :: HasEverything p => Int -> RndTree p +randomTreeBalanced :: (HasEverything p, Monad m) => Int -> RndTree m p randomTreeBalanced n | n <= 1 = do coin <- lift toss if coin @@ -205,10 +205,10 @@ 2 -> replaceFixChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) -randomVec :: Int -> Rng PVector +randomVec :: Monad m => Int -> Rng m 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 :: Monad m => Int -> Int -> Int -> Rng m (Fix SRTree) -> Rng m (SRTree ()) -> Bool -> Rng m (Fix SRTree) randomTree minDepth maxDepth maxSize genTerm genNonTerm grow | noSpaceLeft = genTerm | needNonTerm = genRecursion diff --git a/src/Text/ParseSR.hs b/src/Text/ParseSR.hs --- a/src/Text/ParseSR.hs +++ b/src/Text/ParseSR.hs @@ -26,7 +26,7 @@ import qualified Data.Map.Strict as Map import Data.List.Split ( splitOn ) -import Debug.Trace (trace) +import Debug.Trace (trace, traceShow) -- * Data types @@ -235,7 +235,7 @@ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft] , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] ] - var = do char 'X' + var = do char 'X' <|> char 'x' ix <- decimal pure $ Fix $ Var (ix - 1) -- Operon is not 0-based "var" diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -1,11 +1,11 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.37.0. +-- This file has been generated from package.yaml by hpack version 0.38.1. -- -- see: https://github.com/sol/hpack name: srtree -version: 2.0.1.5 +version: 2.0.1.6 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 @@ -34,6 +34,7 @@ Algorithm.EqSat.Info Algorithm.EqSat.Queries Algorithm.EqSat.SearchSR + Algorithm.EqSat.SearchSRCache Algorithm.EqSat.Simplify Algorithm.Massiv.Utils Algorithm.SRTree.AD