srtree 2.0.1.5 → 2.0.1.6
raw patch · 18 files changed
+754/−65 lines, 18 files
Files
- ChangeLog.md +4/−0
- apps/srtools/Args.hs +7/−0
- apps/srtools/IO.hs +9/−3
- apps/srtools/Main.hs +2/−1
- apps/srtools/Report.hs +8/−2
- src/Algorithm/EqSat.hs +1/−1
- src/Algorithm/EqSat/Build.hs +49/−1
- src/Algorithm/EqSat/Egraph.hs +1/−0
- src/Algorithm/EqSat/SearchSRCache.hs +244/−0
- src/Algorithm/EqSat/Simplify.hs +6/−1
- src/Algorithm/SRTree/AD.hs +219/−2
- src/Algorithm/SRTree/Likelihoods.hs +108/−24
- src/Algorithm/SRTree/ModelSelection.hs +18/−5
- src/Algorithm/SRTree/Opt.hs +31/−0
- src/Data/SRTree/Print.hs +25/−4
- src/Data/SRTree/Random.hs +17/−17
- src/Text/ParseSR.hs +2/−2
- srtree.cabal +3/−2
ChangeLog.md view
@@ -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
apps/srtools/Args.hs view
@@ -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"
apps/srtools/IO.hs view
@@ -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
apps/srtools/Main.hs view
@@ -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
apps/srtools/Report.hs view
@@ -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
src/Algorithm/EqSat.hs view
@@ -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'
src/Algorithm/EqSat/Build.hs view
@@ -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
src/Algorithm/EqSat/Egraph.hs view
@@ -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)
+ src/Algorithm/EqSat/SearchSRCache.hs view
@@ -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''
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, 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
src/Algorithm/SRTree/AD.hs view
@@ -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 <! ix, False)+ evalKey (Const v) = pure $ (M.replicate comp m v, False)+ evalKey (Param ix) = pure $ (M.replicate comp m (theta VS.! ix), True)+ evalKey (Uni f t) = do (v, b) <- getVal t+ pure $ (M.computeAs S . M.map (evalFun f) $ v, b)+ evalKey (Bin op l r) = do (vl, bl) <- getVal l+ (vr, br) <- getVal r+ pure $ (M.computeAs S $ M.zipWith (evalOp op) vl vr, bl || br)++ insKey (Var _) _ _ s = s+ insKey (Const _) _ _ s = s+ insKey (Param _) _ _ s = s+ insKey node v toLocal ((global,local), s) =+ let k = getId node+ in if toLocal+ then ((global, IntMap.insert k v local), s)+ else ((IntMap.insert k v global, local), s)++ insertLocal k v = do (c1, c2) <- get+ put (c1, IntMap.insert k v c2)+ insertGlobal k v = do (c1, c2) <- get+ put (IntMap.insert k v c1, c2)+ getVal rt' = do let rt = canon rt'+ n = getNode rt+ case n of+ Var ix -> 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 <! ix, False)+ evalKey (Const v) = pure $ (M.replicate comp m v, False)+ evalKey (Param ix) = pure $ (M.replicate comp m (theta VS.! ix), True)+ evalKey (Uni f t) = do (v, b) <- getVal t+ pure $ (M.computeAs S . M.map (evalFun f) $ v, b)+ evalKey (Bin op l r) = do (vl, bl) <- getVal l+ (vr, br) <- getVal r+ pure $ (M.computeAs S $ M.zipWith (evalOp op) vl vr, bl || br)++ insKey (Var _) _ _ s = s+ insKey (Const _) _ _ s = s+ insKey (Param _) _ _ s = s+ insKey node v toLocal ((global,local), s) =+ let k = getId node+ in if toLocal+ then ((global, IntMap.insert k v local), s)+ else ((IntMap.insert k v global, local), s)++ insertLocal k v = do (c1, c2) <- get+ put (c1, IntMap.insert k v c2)+ insertGlobal k v = do (c1, c2) <- get+ put (IntMap.insert k v c1, c2)+ getVal rt' = do let rt = canon rt'+ n = getNode rt+ case n of+ Var ix -> 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)
src/Algorithm/SRTree/Likelihoods.hs view
@@ -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
src/Algorithm/SRTree/ModelSelection.hs view
@@ -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
src/Algorithm/SRTree/Opt.hs view
@@ -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
src/Data/SRTree/Print.hs view
@@ -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 #-}
src/Data/SRTree/Random.hs view
@@ -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
src/Text/ParseSR.hs view
@@ -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"
srtree.cabal view
@@ -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