diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,9 @@ # Changelog for srtree +## 2.0.1.7 + +- Added log10 MSE fitness function + ## 2.0.1.6 - Added Fractional Bayes model selection 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 @@ -508,6 +508,35 @@ else do ts <- go n'' (d-1) ns pure (tt:ts) +getNEclassFrom :: Monad m => Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom n eid = getNEclassFrom' n 15 eid + +getNEclassFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom' _ 0 _ = pure [] +getNEclassFrom' n d eId' = do + eId <- canonical eId' + nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass) + (Prelude.map (eId:) <$> go n d nodes) + where + --go :: Int -> Int -> [ENode] -> EGraphST m [[EClassId]] + go n' _ [] = pure [] + go n' 0 ts = pure [] + go n' d (node:ns) = do + tt <- case node of + Bin op l r -> do l' <- getNEclassFrom' n' (d-1) l + r' <- getNEclassFrom' n' (d-1) r + pure $ Prelude.take n [li <> ri | li <- l', ri <- r'] + Uni f t -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]] + Var ix -> pure [[]] + Const x -> pure [[]] + Param ix -> pure [[]] + pure tt + --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 eId <- canonical eId' @@ -550,6 +579,21 @@ then pure [n] else do eids' <- mapM go eids pure ((n : eids) <> concat eids') + +getAllChildBestEClassesRep :: Monad m => EClassId -> EGraphST m [EClassId] +getAllChildBestEClassesRep eId' = do + eId <- canonical eId' + 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 : concat eids') -- | returns a random expression rooted at e-class `eId` getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree) diff --git a/src/Algorithm/EqSat/Info.hs b/src/Algorithm/EqSat/Info.hs --- a/src/Algorithm/EqSat/Info.hs +++ b/src/Algorithm/EqSat/Info.hs @@ -30,6 +30,7 @@ import Algorithm.EqSat.Egraph import Data.AEq (AEq ((~==))) import Algorithm.EqSat.Queries + import Data.Maybe import qualified Data.Set as TrueSet import Data.Sequence (Seq(..), (><)) @@ -170,6 +171,9 @@ insertFitness :: Monad m => EClassId -> Double -> [PVector] -> EGraphST m () insertFitness eId' fit params = do eId <- canonical eId' + tree <- getBestExpr' eId + let p = fromIntegral (length params) + let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0 ec <- gets ((IntMap.! eId) . _eClass) let oldFit = _fitness . _info $ ec --when (oldFit < Just fit) $ do @@ -181,6 +185,7 @@ then modify' $ over (eDB . unevaluated) (IntSet.delete eId) . over (eDB . fitRangeDB) (insertRange eId fit) . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty) + . over (eDB . dlRangeDB) (insertRange eId f_compl) else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit)) insertDL :: Monad m => EClassId -> Double -> EGraphST m () @@ -193,3 +198,11 @@ modify' $ over eClass (IntMap.insert eId newEc) modify' $ over (eDB . dlRangeDB) (insertRange eId fit) . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty) + +-- | TODO: remove from here gets the best expression given the default cost function +getBestExpr' :: Monad m => EClassId -> EGraphST m (Fix SRTree) +getBestExpr' eid = do eid' <- canonical eid + best <- gets (_best . _info . (IntMap.! eid') . _eClass) + childs <- mapM getBestExpr' $ childrenOf best + pure . Fix $ replaceChildren childs best + diff --git a/src/Algorithm/EqSat/Queries.hs b/src/Algorithm/EqSat/Queries.hs --- a/src/Algorithm/EqSat/Queries.hs +++ b/src/Algorithm/EqSat/Queries.hs @@ -80,6 +80,34 @@ then go (m-1) (ecId:bests) t else go m bests t +getTopEClassInRange :: Monad m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId] +getTopEClassInRange b n p range = do + let f = if b then _fitRangeDB else _dlRangeDB + gets (f . _eDB) + >>= go n [] range + where + inRange v (x, y) + | v >= x && v <= y = 0 + | v < x = -1 + | v > y = 1 + | otherwise = 1 + + go :: Monad m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId] + go _ bests [] _ = pure bests + go 0 bests (r:rs) rt = go n bests rs rt + go m bests (r:rs) rt = case rt of + Empty -> pure bests + t :|> y -> do let x = snd y + ecId <- canonical x + ec <- gets ((IntMap.! ecId) . _eClass) + if (isInfinite . fromJust . _fitness . _info $ ec) + then go m bests (r:rs) t + else do let v = p ec + case (v `inRange` r) of + 0 -> go (m-1) (ecId:bests) (r:rs) t -- it is in range, go to the next range + -1 -> go n bests rs (t :|> y) -- it is smaller than the range, get the first n of the next range + 1 -> go m bests (r:rs) t -- y is still greater than the range, keep looking in the same range + getTopECLassIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopECLassIn b n p ecs' = do let f = if b then _fitRangeDB else _dlRangeDB 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 @@ -63,7 +63,7 @@ -- | Supported distributions for negative log-likelihood -- MSE refers to mean squared error -- HGaussian is Gaussian with heteroscedasticity, where the error should be provided -data Distribution = MSE | Gaussian | HGaussian | Bernoulli | Poisson | ROXY +data Distribution = MSE | Gaussian | HGaussian | Bernoulli | Poisson | ROXY | LOG10 deriving (Show, Read, Enum, Bounded, Eq) -- | Sum-of-square errors or Sum-of-square residues @@ -128,6 +128,14 @@ -- | Mean Squared error (not a distribution) nll MSE _ xss ys t theta = mse xss ys t theta +nll LOG10 _ xss ys t theta = M.sum $ (M.map (logBase 10) $ (f (delay ys) / f yhat)) ^ (2 :: Int) + where + yhat = evalTree xss theta t + (Sz m) = M.size ys + f :: Array D Ix1 Double -> Array D Ix1 Double + f z = (z + M.map (\zi -> sqrt (zi^2 + 1e-10)) z) + -- log ys - log y = log (ys/y) + -- | Gaussian distribution, theta must contain an additional parameter corresponding -- to variance. nll Gaussian mYerr xss ys t theta @@ -230,6 +238,11 @@ -- WARNING: pass tree with parameters -- TODO: handle error similar to ROXY buildNLL MSE m tree = ((tree - var (-1)) ** 2) / constv m +buildNLL LOG10 m tree = (((log (y / tree')) / log 10) ** 2) / constv m + where + tree' = (tree + sqrt(tree^2 + 1e-10)) + y = (var (-1) + sqrt(var (-1) ^ 2 + 1e-10)) + buildNLL Gaussian m tree = (square(tree - var (-1)) / square (param p)) + log ((square (param p))) where square = Fix . Uni Square @@ -268,8 +281,36 @@ x <- add myCost (Bin Sub root v) y <- add myCost (Bin Power x c1) add myCost (Bin Div y c2) +buildNLLEGraph LOG10 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) + c3 <- add myCost (Const 10) + c4 <- add myCost (Const 1e-10) + -- log (x + sqrt (x^2 + 1)) / log 10 + log10 <- add myCost (Uni Log c3) + t2 <- add myCost (Uni Square root) + t2p1 <- add myCost (Bin Add t2 c4) + sqt <- add myCost (Uni Sqrt t2p1) + tpt <- add myCost (Bin Add root sqt) + -- same with y + y2 <- add myCost (Uni Square v) + y2p1 <- add myCost (Bin Add y2 c4) + sqy <- add myCost (Uni Sqrt y2p1) + ypy <- add myCost (Bin Add v sqy) + tptypy <- add myCost (Bin Div ypy tpt) + + logy <- add myCost (Uni Log tptypy) + log10y <- add myCost (Bin Div logy log10) + + --x <- add myCost (Bin Sub log10t v) + y <- add myCost (Bin Power tptypy c1) + add myCost (Bin Div y c2) + buildNLLEGraph Gaussian m egraph root = runIdentity (addToEg `runStateT` egraph) where p = countParamsUniqEg egraph root @@ -329,6 +370,7 @@ -- | Prediction for different distributions predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector predict MSE tree theta xss = evalTree xss theta tree +predict LOG10 tree theta xss = evalTree xss theta tree predict Gaussian tree theta xss = evalTree xss theta tree predict Bernoulli tree theta xss = logistic $ evalTree xss theta tree predict Poisson tree theta xss = exp $ evalTree xss theta tree @@ -348,13 +390,15 @@ eps = 1e-8 f = (/ fromIntegral m) . M.sum . M.map (^2) $ (predict MSE tree theta xss) - delay ys finitediff ix = let t1 = disturb ix - f' = (/ fromIntegral m) . M.sum . M.map (^2) $ (predict MSE tree t1 xss) - delay ys + f' = (/ fromIntegral m) . M.sum . M.map (^2) $ (predict MSE tree t1 xss) - ys' in (f' - f)/eps (Sz2 m _) = M.size xss tree' = buildNLL dist (fromIntegral m) tree treeArr = IntMap.toAscList $ tree2arr tree' j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] - + flog :: Array D Ix1 Double -> Array D Ix1 Double + flog z = M.map (logBase 10) (z + M.map sqrt (z^2 + 1e-10)) + ys' = (if dist==LOG10 then id else id) (delay ys) nanTo0 x = x -- if isNaN x || isInfinite x then 0 else x @@ -366,6 +410,11 @@ where (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix grad' = M.map nanTo0 grad +gradNLLArr LOG10 xss ys mYerr tree j2ix theta = + (M.sum yhat, delay grad') + where + (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix + grad' = M.map nanTo0 grad gradNLLArr Gaussian xss ys mYerr tree j2ix theta = (M.sum yhat, delay grad') where @@ -395,6 +444,11 @@ where (yhat, grad) = reverseModeGraph xss ys mYerr theta tree grad' = VS.map nanTo0 grad +gradNLLGraph LOG10 xss ys mYerr tree theta = + (M.sum yhat, grad') + where + (yhat, grad) = reverseModeGraph xss ys mYerr theta tree + grad' = VS.map nanTo0 grad gradNLLGraph Gaussian xss ys mYerr tree theta = (M.sum yhat, grad') where @@ -424,6 +478,13 @@ where (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta grad' = VS.map nanTo0 grad +gradNLLEGraph LOG10 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 + ys' :: PVector + ys' = M.computeAs M.S $ M.map (logBase 10) (delay ys + M.map sqrt (delay ys^2 + 1e-10)) gradNLLEGraph Gaussian xss ys mYerr egraph cache root theta = (M.sum yhat, grad') where @@ -543,6 +604,7 @@ (phi, phi') = case dist of MSE -> (yhat, M.replicate cmp (Sz m) 1) + LOG10 -> (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) diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack name: srtree -version: 2.0.1.6 +version: 2.0.1.8 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 @@ -67,7 +67,7 @@ , base >=4.19 && <5 , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.8 + , containers >=0.6.7 && <0.9 , dlist ==1.0.* , exceptions >=0.10.7 && <0.11 , filepath >=1.4.0.0 && <1.6 @@ -102,7 +102,7 @@ , base >=4.19 && <5 , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.8 + , containers >=0.6.7 && <0.9 , dlist ==1.0.* , exceptions >=0.10.7 && <0.11 , filepath >=1.4.0.0 && <1.6 @@ -142,7 +142,7 @@ , base >=4.19 && <5 , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.8 + , containers >=0.6.7 && <0.9 , dlist ==1.0.* , exceptions >=0.10.7 && <0.11 , filepath >=1.4.0.0 && <1.6 @@ -182,7 +182,7 @@ , base >=4.19 && <5 , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.8 + , containers >=0.6.7 && <0.9 , dlist ==1.0.* , exceptions >=0.10.7 && <0.11 , filepath >=1.4.0.0 && <1.6 @@ -222,7 +222,7 @@ , base >=4.19 && <5 , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.8 + , containers >=0.6.7 && <0.9 , dlist ==1.0.* , exceptions >=0.10.7 && <0.11 , filepath >=1.4.0.0 && <1.6