srtree 2.0.0.2 → 2.0.0.3
raw patch · 42 files changed
+3996/−1545 lines, 42 filesdep +ansi-terminaldep +binarydep +replinedep −normaldistributiondep ~basedep ~randomdep ~splitnew-component:exe:egraphSearchnew-component:exe:reggression
Dependencies added: ansi-terminal, binary, repline, scheduler, table-layout, unliftio, unliftio-core
Dependencies removed: normaldistribution
Dependency ranges changed: base, random, split
Files
- ChangeLog.md +8/−0
- README.md +1/−1
- apps/egraphGP/Main.hs +378/−494
- apps/egraphGP/Random.hs +2/−2
- apps/egraphGP/Util.hs +258/−0
- apps/egraphSearch/Main.hs +325/−0
- apps/egraphSearch/Random.hs +54/−0
- apps/egraphSearch/Util.hs +254/−0
- apps/ieeexplore/Main.hs +0/−101
- apps/rEGGression/Commands.hs +393/−0
- apps/rEGGression/Main.hs +336/−0
- apps/rEGGression/Random.hs +33/−0
- apps/rEGGression/Util.hs +272/−0
- apps/srtools/Args.hs +13/−14
- apps/srtools/IO.hs +53/−10
- apps/srtools/Main.hs +3/−1
- apps/srtools/Report.hs +47/−44
- apps/tinygp/GP.hs +97/−73
- apps/tinygp/Main.hs +43/−14
- apps/tinygp/Util.hs +62/−0
- src/Algorithm/EqSat.hs +19/−6
- src/Algorithm/EqSat/Build.hs +71/−15
- src/Algorithm/EqSat/DB.hs +26/−4
- src/Algorithm/EqSat/Egraph.hs +78/−11
- src/Algorithm/EqSat/Info.hs +36/−9
- src/Algorithm/EqSat/Queries.hs +113/−10
- src/Algorithm/EqSat/Simplify.hs +85/−24
- src/Algorithm/SRTree/AD.hs +224/−381
- src/Algorithm/SRTree/ConfidenceIntervals.hs +35/−35
- src/Algorithm/SRTree/Likelihoods.hs +308/−114
- src/Algorithm/SRTree/ModelSelection.hs +28/−27
- src/Algorithm/SRTree/NonlinearOpt.hs +4/−2
- src/Algorithm/SRTree/Opt.hs +45/−81
- src/Data/SRTree.hs +4/−0
- src/Data/SRTree/Datasets.hs +19/−10
- src/Data/SRTree/Derivative.hs +3/−1
- src/Data/SRTree/Eval.hs +5/−0
- src/Data/SRTree/Internal.hs +28/−0
- src/Data/SRTree/Print.hs +3/−2
- src/Numeric/Optimization/NLOPT/Bindings.hs +12/−2
- src/Text/ParseSR.hs +112/−29
- srtree.cabal +106/−28
ChangeLog.md view
@@ -1,5 +1,13 @@ # Changelog for srtree +## 2.0.0.3++- Fixed compatibility with random-1.3.0 and GHC-9.12.1 +- Fixed bug in Bernoulli distribution +- Removed `log(sqrt(x))` rule in parametric rules due to generating longer expressions +- Fixed DL calculation without the correct number of parameters +- Fixed memory issue when querying pattern distribution + ## 2.0.0.0 - Complete refactoring of the library
README.md view
@@ -274,4 +274,4 @@ - support more advanced functions - support conditional branching (`IF-THEN-ELSE`)-+- document egraph-search and ieeexplore
apps/egraphGP/Main.hs view
@@ -3,6 +3,7 @@ {-# LANGUAGE MultiWayIf #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} module Main where @@ -13,11 +14,13 @@ import Algorithm.EqSat.Info import Algorithm.EqSat.DB import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection import Algorithm.SRTree.Opt import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.)) import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_) import Control.Monad.State.Strict import qualified Data.IntMap.Strict as IM+import qualified Data.Map.Strict as Map import Data.Massiv.Array as MA hiding (forM_, forM) import Data.Maybe (fromJust, isNothing, isJust) import Data.SRTree@@ -29,7 +32,7 @@ import Random import System.Random import qualified Data.HashSet as Set-import Data.List ( sort, maximumBy, intercalate, sortOn )+import Data.List ( sort, maximumBy, intercalate, sortOn, intersperse, nub ) import Data.IntSet (IntSet) import qualified Data.IntSet as IntSet import qualified Data.Sequence as FingerTree@@ -37,477 +40,336 @@ import qualified Data.Foldable as Foldable import qualified Data.IntMap as IntMap import List.Shuffle ( shuffle )+import Algorithm.SRTree.NonlinearOpt+import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS import Debug.Trace-import Algorithm.EqSat (runEqSat)---- Insert random expression --- Evaluate random subtree --- Insert new random parent eNode --type RndEGraph a = EGraphST (StateT StdGen IO) a--io = lift . lift-{-# INLINE io #-}-rnd = lift-{-# INLINE rnd #-}--myCost :: SRTree Int -> Int-myCost (Var _) = 1-myCost (Const _) = 1-myCost (Param _) = 1-myCost (Bin op l r) = 2 + l + r-myCost (Uni _ t) = 3 + t--data Alg = OnlyRandom | BestFirst deriving (Show, Read, Eq)---- experiment 1 80/30-fitnessFun :: SRMatrix -> PVector -> SRMatrix -> PVector -> Fix SRTree -> RndEGraph (Double, PVector)-fitnessFun x y x_val y_val _tree = do- let tree = relabelParams _tree- nParams = countParams tree- thetaOrig <- rnd $ randomVec nParams -- = MA.replicate Seq nParams 1.0- let (theta, fit) = minimizeNLL Gaussian Nothing 50 x y tree thetaOrig- tr = negate . mse x y tree $ if nParams == 0 then thetaOrig else theta- val = negate . mse x_val y_val tree $ if nParams == 0 then thetaOrig else theta- -- val = r2 x y tree $ if nParams == 0 then thetaOrig else theta- pure $ if isNaN val || isNaN tr- then (-1/0, theta) -- infinity- else (min tr val, theta)-{-# INLINE fitnessFun #-}--fitnessFunRep :: SRMatrix -> PVector -> SRMatrix -> PVector -> Fix SRTree -> RndEGraph (Double, PVector)-fitnessFunRep x y x_val y_val _tree = do- fits <- replicateM 1 (fitnessFun x y x_val y_val _tree)- pure (maximumBy (compare `on` fst) fits)-{-# INLINE fitnessFunRep #-}---- helper query functions-fitnessIs p = p . _fitness . _info-{-# INLINE fitnessIs #-}--getFitness :: EClassId -> RndEGraph (Maybe Double)-getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)-{-# INLINE getFitness #-}-getSize :: EClassId -> RndEGraph Int-getSize c = gets (_size . _info . (IM.! c) . _eClass)-{-# INLINE getSize #-}-getSizeOf :: (Int -> Bool) -> [EClassId] -> RndEGraph [EClassId]-getSizeOf p = filterM (getSize >=> (pure . p))-{-# INLINE getSizeOf #-}--(&&&) p1 p2 x = p1 x && p2 x-{-# INLINE (&&&) #-}--isValidFitness = fitnessIs (isJust &&& (not . isNaN . fromJust) &&& (not . isInfinite . fromJust))-{-# INLINE isValidFitness #-}--evaluated = fitnessIs isJust-{-# INLINE evaluated #-}-unevaluated' = fitnessIs isNothing-{-# INLINE unevaluated' #-}--isSizeOf p = p . _size . _info-{-# INLINE isSizeOf #-}+import Algorithm.EqSat (runEqSat,applySingleMergeOnlyEqSat) -funDoesNotExistWith node = Prelude.any (not . (`sameFunAs` node) . snd) . _parents- where sameFunAs (Uni f _) (Uni g _) = f == g- sameFunAs _ _ = False-{-# INLINE funDoesNotExistWith #-}+import GHC.IO (unsafePerformIO)+import Control.Scheduler +import Control.Monad.IO.Unlift+import Data.SRTree (convertProtectedOps) -opDoesNotExistWith :: (SRTree ()) -> EClassId -> EClass -> Bool-opDoesNotExistWith node ecId = Prelude.any (not . (`sameOpAs` node) . snd) . _parents- where sameOpAs (Bin op1 l _) (Bin op2 _ _) = op1 == op2 && ecId == l- sameOpAs _ _ = False-{-# INLINE opDoesNotExistWith #-}+import Util -rewriteBasic2 :: [Rule]-rewriteBasic2 =- [- "x" * "y" :=> "y" * "x"- , "x" + "y" :=> "y" + "x"- , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"- , ("x" + "y") + "z" :=> "x" + ("y" + "z")- , ("x" * "y") * "z" :=> "x" * ("y" * "z")- , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")- , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"- ]+egraphGP :: DataSet -> DataSet -> DataSet -> Args -> StateT EGraph (StateT StdGen IO) ()+egraphGP dataTrain dataVal dataTest args = do+ when ((not.null) (_loadFrom args)) $ (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg) -egraphSearch alg x y x_val y_val x_te y_te terms nEvals maxSize = do- ec <- insertRndExpr maxSize- updateIfNothing ec+ pop <- replicateM (_nPop args) $ do ec <- insertRndExpr (_maxSize args) rndTerm rndNonTerm >>= canonical+ updateIfNothing fitFun ec+ pure ec insertTerms- evaluateUnevaluated- runEqSat myCost rewriteBasic2 1+ --runEqSat myCost rewritesParams 1+ --cleanDB+ evaluateUnevaluated fitFun+ pop' <- Prelude.mapM canonical pop+ when (_trace args) $ printPop pop'+ let m = (_nPop args) `div` (_maxSize args) - while ((<nEvals) . snd) (1,1) $- \(radius, nEvs) ->- do- --nEvs <- gets (FingerTree.size . _fitRangeDB . _eDB)- nCls <- gets (IM.size . _eClass)- nUnev <- gets (IntSet.size . _unevaluated . _eDB)- let nEvs = nCls - nUnev- --io . print $ (nCls, nEvs)- bestF <- getBestFitness+ finalPop <- iterateFor (_gens args) pop' $ \it ps' -> do+ -- ps <- Prelude.mapM canonical ps'+ --let chunks = 1+ -- (parts, r) = _nPop args `divMod` chunks -- chunks of 50+ -- genChunk = Prelude.mapM canonical ps' >>= \ps -> replicateM chunks (tournament ps) >>= Prelude.mapM combine+ -- genPart = genChunk >>= \chunk -> (runEqSat myCost rewritesSimple 1 >> cleanDB) >> pure chunk+ --newPop' <- (<>) <$> (concat <$> (replicateM parts genPart)) <*> (concat <$> replicateM r genPart)+ --parents <- replicateM (_nPop args - if (_moo args) then (_maxSize args) else 0) (tournament ps)+ --newPop' <- Prelude.mapM combine parents+ --runEqSat myCost rewritesSimple 1 --applySingleMergeOnlyEqSat myCost rewritesSimple+ --cleanDB+ newPop' <- replicateM (_nPop args) (evolve ps')+ --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB+ --newPop' <- Prelude.mapM (\eId -> canonical eId >>= \eId' -> (updateIfNothing fitFun eId' >> pure eId')) newPop''+ --Prelude.mapM_ (updateIfNothing fitFun) newPop' - (ecN, b) <- case alg of- OnlyRandom -> do let ratio = fromIntegral nEvs / fromIntegral nCls- b <- rnd (tossBiased ratio)- ec <- if b && ratio > 0.99 then insertRndExpr maxSize >>= canonical else evaluateRndUnevaluated >>= canonical- pure (ec, False)- BestFirst -> do- ecsPareto <- getParetoEcsUpTo radius- ecsBest <- getBestEcs (isSizeOf (<=maxSize)) radius+ totSz <- gets (IntMap.size . _eClass)+ let full = False -- totSz > max maxMem (_nPop args)+ when full cleanEGraph - ecPareto <- combineFrom ecsPareto- curFitPareto <- getFitness ecPareto+ newPop <- if (_moo args)+ then do+ let n_paretos = (_nPop args) `div` (_maxSize args)+ pareto <- concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n 2)+ let remainder = _nPop args - length pareto+ lft <- if full+ then getTopFitEClassThat remainder (const True)+ else pure $ Prelude.take remainder newPop'+ Prelude.mapM canonical (pareto <> lft)+ else if full+ then getTopFitEClassThat (_nPop args) (const True)+ else Prelude.mapM canonical newPop'+ when (_trace args) $ printPop newPop - if isNothing curFitPareto- then pure (ecPareto, False)- else do ecBest <- combineFrom ecsBest- curFitBest <- getFitness ecBest- if isNothing curFitBest- then pure (ecBest, False)- else do ee <- evalRndSubTree- case ee of- Nothing -> do ec <- insertRndExpr maxSize >>= canonical- pure (ec, True)- Just c -> pure (c, False) - upd <- updateIfNothing ecN- when (upd)- do runEqSat myCost rewriteBasic2 1- cleanDB- pure ()- let radius' = if b then (max 1 $ min (200 `div` maxSize) (radius+1)) else (max 1 $ radius-1)- nEvs' = nEvs + if upd then 1 else 0- pure (radius', nEvs')- eclasses <- gets (IntMap.toList . _eClass)- -- forM_ eclasses $ \(_, v) -> (io.print) (Set.size (_eNodes v), Set.size (_parents v))- paretoFront- --ft <- gets (_fitRangeDB . _eDB)- --io . print $ Foldable.toList ft+ pure newPop + io $ putStrLn "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"+ if (_printPareto args)+ then paretoFront fitFun (_maxSize args) printExpr+ else printBest fitFun printExpr+ when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode ) where- numberOfEvalClasses :: Monad m => Int -> EGraphST m Bool- numberOfEvalClasses nEvs =- (subtract <$> gets (IntSet.size . _unevaluated . _eDB) <*> gets (IM.size . _eClass))- >>= \n -> pure (n<nEvs)-- updateIfNothing ec = do- mf <- getFitness ec- case mf of- Nothing -> do- t <- getBest ec- (f, p) <- fitnessFunRep x y x_val y_val t- insertFitness ec f p- pure True- Just _ -> pure False-- getBestFitness = do- bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)- gets (_fitness . _info . (IM.! bec) . _eClass)-- evalRndSubTree :: RndEGraph (Maybe EClassId)- evalRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)- if not (null ecIds)- then do rndId <- rnd $ randomFrom ecIds- Just <$> canonical rndId- else pure Nothing--- combineFrom ecs = do- nt <- rnd rndNonTerm- p1 <- rnd (randomFrom ecs)- p2 <- rnd (randomFrom ecs)- l1 <- rnd (randomFrom [1..maxSize-2])- l2 <- rnd (randomFrom [1..(maxSize - l1 - 1)])- e1 <- randomChildFrom p1 l1- ml <- gets (_size . _info . (IM.! e1) . _eClass)- e2 <- randomChildFrom p2 l2- case nt of- Uni Id () -> canonical e1- Uni f () -> add myCost (Uni f e1) >>= canonical- Bin op () () -> do b <- rnd toss- if b- then add myCost (Bin op e1 e2) >>= canonical- else add myCost (Bin op e2 e1) >>= canonical-- getParetoEcsUpTo n = concat <$> (forM [1..maxSize] $ \i -> getBestEcsOfSize i n)-- getBestEcsOfSize i n = getTopECLassWithSize i n- --do- --ecs <- getTopECLassWithSize i n- --Prelude.mapM canonical ecs -- (Prelude.take n ecs)-- getBestEcs p n = getTopECLassThat n p- --do- --ecs <- getTopECLassThat n p- --fits <- Prelude.mapM getFitness ecs- --let sorted = sort $ Prelude.zip (Prelude.map (fmap negate) fits) ecs- --Prelude.mapM canonical (Prelude.take n ecs)-- randomChildFrom ec maxL = do- p <- rnd toss -- whether to go deeper or return this level- l <- gets (_size . _info . (IM.! ec) . _eClass )+ maxSize = (_maxSize args)+ maxMem = 10000 -- running 1 iter of eqsat for each new individual will consume ~3GB+ fitFun = fitnessFunRep (_optRepeat args) (_optIter args) (_distribution args) dataTrain dataVal+ nonTerms = parseNonTerms (_nonterminals args)+ (Sz2 _ nFeats) = MA.size (getX dataTrain)+ terms = if _distribution args == ROXY+ then [var 0, param 0]+ else [var ix | ix <- [0 .. nFeats-1]] <> [param 0]+ uniNonTerms = [t | t <- nonTerms, isUni t]+ binNonTerms = [t | t <- nonTerms, isBin t]+ isUni (Uni _ _) = True+ isUni _ = False+ isBin (Bin _ _ _) = True+ isBin _ = False - if p || l >= maxL- then do enodes <- gets (_eNodes . (IM.! ec) . _eClass)- enode <- gets (_best . _info . (IM.! ec) . _eClass) -- we should return the best otherwise we may build larger- case enode of- Uni _ eci -> randomChildFrom eci maxL- Bin _ ecl ecr -> do coin <- rnd toss- if coin- then randomChildFrom ecl maxL- else randomChildFrom ecr maxL- _ -> pure ec- else pure ec+ -- TODO: merge two or more egraphs+ cleanEGraph = do let nParetos = (maxMem `div` 5) `div` _maxSize args+ pareto <- (concat <$> (forM [1 .. _maxSize args] $ \n -> getTopFitEClassWithSize n nParetos))+ >>= Prelude.mapM canonical+ infos <- forM pareto (\c -> gets (_info . (IntMap.! c) . _eClass))+ exprs <- forM pareto getBestExpr+ put emptyGraph+ newIds <- fromTrees myCost exprs+ forM_ (Prelude.zip newIds (Prelude.reverse infos)) $ \(eId, info) ->+ insertFitness eId (fromJust $ _fitness info) (fromJust $ _theta info) - nonTerms = [ Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () ()- , Bin PowerAbs () (), Uni Recip (), Uni LogAbs (), Uni Exp (), Uni Sin (), Uni SqrtAbs ()] rndTerm = Random.randomFrom terms- rndNonTerm = Random.randomFrom $ (Uni Id ()) : nonTerms- rndNonTerm2 = Random.randomFrom nonTerms-- insertTerms =- forM terms $ \t -> do fromTree myCost t >>= canonical-- insertRndExpr :: Int -> RndEGraph EClassId- insertRndExpr maxSize =- do grow <- rnd toss- n <- rnd (randomFrom [3 .. maxSize])- t <- rnd $ Random.randomTree 2 8 n rndTerm rndNonTerm2 grow- fromTree myCost t >>= canonical-- insertBestExpr :: RndEGraph EClassId- insertBestExpr = do --let t = "t0" / (recip ("t1" - "x0") + powabs "t2" "x0")- let t = ((("t0" + (powabs "t0" "x0")) / "t0") * "x0")- ecId <- fromTree myCost t >>= canonical- (f, p) <- fitnessFunRep x y x_val y_val t- insertFitness ecId f p- io . putStrLn $ "Best fit global: " <> show f- pure ecId- where powabs l r = Fix (Bin PowerAbs l r)-- getBestEclassThat p =- do ecIds <- getTopECLassThat 1 p -- isValidFitness- --bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds- --ecIds' <- getEClassesThat (fitnessIs (== Just bestFit))- Prelude.mapM canonical $ Prelude.take 1 ecIds-- getBestExprWithSize n =- do ec <- getTopECLassWithSize n 1- if (not (null ec))- then do- bestFit <- getFitness $ head ec- bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)- (:[]) . (,bestP) . (,bestFit) <$> getBest (head ec)- else pure []-- getBestExprThat p =- do ec <- getBestEclassThat p- if (not (null ec))- then do- bestFit <- getFitness $ head ec- (:[]) . (,bestFit) <$> getBest (head ec)- else pure []-- printAll = do- ecs <- gets (IM.keys . _eClass)- forM_ ecs $ \ec ->- do t <- getBest ec- f <- gets (_fitness . _info . (IM.! ec) . _eClass)- io . putStrLn $ showExpr t <> " " <> show f-- paretoFront = go 1 (-1.0/0.0)- where- go n f- | n > maxSize = pure ()- | otherwise = do- ecList <- getBestExprWithSize n- if (not (null ecList))- then do let ((best, mf), mtheta) = head ecList- best' = relabelParams best- x_tot <- MA.computeAs MA.S <$> (MA.concatOuterM $ Prelude.map MA.toLoadArray [x, x_val])- y_tot <- MA.computeAs MA.S <$> (MA.concatOuterM $ Prelude.map MA.toLoadArray [y, y_val])-- --(fit_tr, theta) <- fitnessFunRep x_tot y_tot x_tot y_tot best'- let fit = fromJust mf- fit_tr = fit- theta = fromJust mtheta- fit_te = mse x_te y_te best' theta- str_th = intercalate ";" $ Prelude.map show $ MA.toList theta-- when (fit > f) do- io . putStrLn $ showExpr best <> "," <> str_th <> "," <> show (negate fit) <> "," <> show (negate fit_tr) <> "," <> show fit_te- go (n+1) (max fit f)- else go (n+1) f-- evaluateUnevaluated = do- ec <- gets (IntSet.toList . _unevaluated . _eDB)- forM_ ec $ \c -> do- t <- getBest c- (f, p) <- fitnessFun x y x_val y_val t- insertFitness c f p-- evaluateRndUnevaluated = do- ec <- gets (IntSet.toList . _unevaluated . _eDB)- c <- rnd . randomFrom $ ec - t <- getBest c- (f, p) <- fitnessFun x y x_val y_val t- insertFitness c f p- pure c--while p arg prog = do when (p arg) do arg' <- prog arg- while p arg' prog-- {--egraphGP :: SRMatrix -> PVector -> [Fix SRTree] -> Int -> RndEGraph (Fix SRTree, Double)-egraphGP x y terms nEvals = do- replicateM_ 200 insertRndExpr- getBestExpr- runEqSat myCost rewrites 50- evaluateUnevaluated- paretoFront- getBestExpr- where- paretoFront = do - forM_ [1..10] $ \i ->- do (best, fit) <- getBestExprThat (evaluated &&& isSizeOf (==i))- io . putStrLn $ showExpr best <> " " <> show fit -- evaluateUnevaluated = do - ec <- getEClassesThat unevaluated- forM_ ec $ \c -> do - t <- getBest c - f <- fitnessFun x y t- updateFitness f c ---+ rndNonTerm = Random.randomFrom nonTerms + refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList >>= pure . nub+ modify' $ over (eDB . refits) (const Set.empty)+ forM_ ids $ \ec -> do t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p + iterateFor 0 xs f = pure xs+ iterateFor n xs f = do xs' <- f n xs+ iterateFor (n-1) xs' f + evolve xs' = do xs <- Prelude.mapM canonical xs'+ parents <- tournament xs+ offspring <- combine parents+ --applySingleMergeOnlyEqSat myCost rewritesParams >> cleanDB+ runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged+ canonical offspring >>= updateIfNothing fitFun+ canonical offspring+ --pure offspring + tournament xs = do p1 <- applyTournament xs >>= canonical+ p2 <- applyTournament xs >>= canonical+ pure (p1, p2) + applyTournament :: [EClassId] -> RndEGraph EClassId+ applyTournament xs = do challengers <- replicateM (_nTournament args) (rnd $ randomFrom xs) >>= traverse canonical+ fits <- Prelude.map fromJust <$> Prelude.mapM getFitness challengers+ pure . snd . maximumBy (compare `on` fst) $ Prelude.zip fits challengers + combine (p1, p2) = (crossover p1 p2 >>= mutate) >>= canonical + crossover p1 p2 = do sz <- getSize p1+ coin <- rnd $ tossBiased (_pc args)+ if sz == 1 || not coin+ then rnd (randomFrom [p1, p2])+ else do pos <- rnd $ randomRange (1, sz-1)+ cands <- getAllSubClasses p2+ tree <- getSubtree pos 0 Nothing [] cands p1+ fromTree myCost tree >>= canonical + getSubtree :: Int -> Int -> Maybe (EClassId -> ENode) -> [Maybe (EClassId -> ENode)] -> [EClassId] -> EClassId -> RndEGraph (Fix SRTree)+ getSubtree 0 sz (Just parent) mGrandParents cands p' = do+ p <- canonical p'+ candidates' <- filterM (\c -> (<maxSize-sz) <$> getSize c) cands+ candidates <- filterM (\c -> doesNotExistGens mGrandParents (parent c)) candidates'+ >>= traverse canonical+ if null candidates+ then getBestExpr p+ else do subtree <- rnd (randomFrom candidates)+ getBestExpr subtree+ getSubtree pos sz parent mGrandParents cands p' = do+ p <- canonical p'+ root <- getBestENode p >>= canonize+ case root of+ Param ix -> pure . Fix $ Param ix+ Const x -> pure . Fix $ Const x+ Var ix -> pure . Fix $ Var ix+ Uni f t' -> do t <- canonical t'+ (Fix . Uni f) <$> getSubtree (pos-1) (sz+1) (Just $ Uni f) (parent:mGrandParents) cands t+ Bin op l'' r'' ->+ do l <- canonical l''+ r <- canonical r''+ szLft <- getSize l+ szRgt <- getSize r+ if szLft < pos+ then do l' <- getBestExpr l+ r' <- getSubtree (pos-szLft-1) (sz+szLft+1) (Just $ Bin op l) (parent:mGrandParents) cands r+ pure . Fix $ Bin op l' r'+ else do l' <- getSubtree (pos-1) (sz+szRgt+1) (Just (\t -> Bin op t r)) (parent:mGrandParents) cands l+ r' <- getBestExpr r+ pure . Fix $ Bin op l' r' + getAllSubClasses p' = do+ p <- canonical p'+ en <- getBestENode p+ case en of+ Bin _ l r -> do ls <- getAllSubClasses l+ rs <- getAllSubClasses r+ pure (p : (ls <> rs))+ Uni _ t -> (p:) <$> getAllSubClasses t+ _ -> pure [p] + mutate p = do sz <- getSize p+ coin <- rnd $ tossBiased (_pm args)+ if coin+ then do pos <- rnd $ randomRange (0, sz-1)+ tree <- mutAt pos maxSize Nothing p+ fromTree myCost tree >>= canonical+ else pure p + peel :: Fix SRTree -> SRTree ()+ peel (Fix (Bin op l r)) = Bin op () ()+ peel (Fix (Uni f t)) = Uni f ()+ peel (Fix (Param ix)) = Param ix+ peel (Fix (Var ix)) = Var ix+ peel (Fix (Const x)) = Const x + mutAt :: Int -> Int -> Maybe (EClassId -> ENode) -> EClassId -> RndEGraph (Fix SRTree)+ mutAt 0 sizeLeft Nothing _ = (insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical) >>= getBestExpr -- we chose to mutate the root+ mutAt 0 1 _ _ = rnd $ randomFrom terms -- we don't have size left+ mutAt 0 sizeLeft (Just parent) _ = do -- we reached the mutation place+ ec <- insertRndExpr sizeLeft rndTerm rndNonTerm >>= canonical -- create a random expression with the size limit+ (Fix tree) <- getBestExpr ec --+ root <- getBestENode ec+ exist <- canonize (parent ec) >>= doesExist+ if exist+ -- the expression `parent ec` already exists, try to fix+ then do let children = childrenOf root+ candidates <- case length children of+ 0 -> filterM (checkToken parent . (replaceChildren children)) (Prelude.map peel terms)+ 1 -> filterM (checkToken parent . (replaceChildren children)) uniNonTerms+ 2 -> filterM (checkToken parent . (replaceChildren children)) binNonTerms+ if null candidates+ then pure $ Fix tree -- there's no candidate, so we failed and admit defeat+ else do newToken <- rnd (randomFrom candidates)+ pure . Fix $ replaceChildren (childrenOf tree) newToken -------------------- GARBAGE CODE -------------------- go i = do n <- getEClassesThat isValidFitness- unless (length n >= nEvals)- do gpStep- when (i `mod` 1000 == 0) (getBestExpr >>= (io . print . snd))- when (i `mod` 1000000 == 0) $ do- n <- gets (IM.size . _eClass)- --io $ putStrLn ("before: " <> show n)- applyMergeOnlyDftl myCost- n1 <- gets (IM.size . _eClass)- when (n1 < n) $ io $ print (n,n1)- --io $ putStrLn ("after: " <> show n1)- go (i+1)+ else pure . Fix $ tree - rndTerm = Random.randomFrom terms- rndNonTerm = Random.randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () ()- , Bin PowerAbs () (), Uni Recip ()]+ mutAt pos sizeLeft parent p' = do+ p <- canonical p'+ root <- getBestENode p >>= canonize+ case root of+ Param ix -> pure . Fix $ Param ix+ Const x -> pure . Fix $ Const x+ Var ix -> pure . Fix $ Var ix+ Uni f t' -> canonical t' >>= \t -> (Fix . Uni f) <$> mutAt (pos-1) (sizeLeft-1) (Just $ Uni f) t+ Bin op ln rn -> do l <- canonical ln+ r <- canonical rn+ szLft <- getSize l+ szRgt <- getSize r+ if szLft < pos+ then do l' <- getBestExpr l+ r' <- mutAt (pos-szLft-1) (sizeLeft-szLft-1) (Just $ Bin op l) r+ pure . Fix $ Bin op l' r'+ else do l' <- mutAt (pos-1) (sizeLeft-szRgt-1) (Just (\t -> Bin op t r)) l+ r' <- getBestExpr r+ pure . Fix $ Bin op l' r' - getBestExpr :: RndEGraph (Fix SRTree, Double) - getBestExpr = do ecIds <- getEClassesThat evaluated -- isValidFitness- nc <- gets (IM.size . _eClass)- io . putStrLn $ "Evaluated expressions: " <> show (length ecIds) <> " / " <> show nc- bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds- ecIds' <- getEClassesThat (fitnessIs (== Just bestFit))- (,bestFit) <$> getBest (head ecIds')+ checkToken parent en' = do en <- canonize en'+ mEc <- gets ((Map.!? en) . _eNodeToEClass)+ case mEc of+ Nothing -> pure True+ Just ec -> do ec' <- canonical ec+ ec'' <- canonize (parent ec')+ not <$> doesExist ec''+ doesExist, doesNotExist :: ENode -> RndEGraph Bool+ doesExist en = gets ((Map.member en) . _eNodeToEClass)+ doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass) - getBestExprThat p = - do ecIds <- getEClassesThat p -- isValidFitness- nc <- gets (IM.size . _eClass)- bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds- ecIds' <- getEClassesThat (fitnessIs (== Just bestFit))- (,bestFit) <$> getBest (head ecIds')+ 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' - insertRndExpr :: RndEGraph () - insertRndExpr = do grow <- rnd toss- t <- rnd $ Random.randomTree 2 6 10 rndTerm rndNonTerm grow- f <- fitnessFun x y t- ecId <- fromTree myCost t >>= canonical- -- io $ print ('i', showExpr t, f)- updateFitness f ecId+ printExpr :: Int -> EClassId -> RndEGraph ()+ printExpr ix ec = do+ theta' <- gets (fromJust . _theta . _info . (IM.! ec) . _eClass)+ bestExpr <- getBestExpr ec+ let nParams = countParams bestExpr+ (MA.Sz nTheta) = MA.size theta'+ (_, theta) <- if (nParams /= nTheta)+ then fitFun bestExpr+ else pure (1.0, theta') - evalRndSubTree :: RndEGraph ()- evalRndSubTree = do ecIds <- getEClassesThat unevaluated- unless (null ecIds) do- rndId <- rnd $ randomFrom ecIds- rndId' <- canonical rndId - t <- getBest rndId'- f <- fitnessFun x y t- -- io $ print ('e', showExpr t, f)- updateFitness f rndId'+ let (x, y, mYErr) = dataTrain+ (x_val, y_val, mYErr_val) = dataVal+ (x_te, y_te, mYErr_te) = dataTest+ distribution = _distribution args+ best' = relabelParams bestExpr+ expr = paramsToConst (MA.toList theta) best'+ mse_train = mse x y best' theta+ mse_val = mse x_val y_val best' theta+ mse_te = mse x_te y_te best' theta+ r2_train = r2 x y best' theta+ r2_val = r2 x_val y_val best' theta+ r2_te = r2 x_te y_te best' theta+ nll_train = nll distribution mYErr x y best' theta+ nll_val = nll distribution mYErr_val x_val y_val best' theta+ nll_te = nll distribution mYErr_te x_te y_te best' theta+ mdl_train = mdl distribution mYErr x y theta best'+ mdl_val = mdl distribution mYErr_val x_val y_val theta best'+ mdl_te = mdl distribution mYErr_te x_te y_te theta best'+ vals = intercalate ","+ $ Prelude.map show [mse_train, mse_val, mse_te+ , r2_train, r2_val, r2_te+ , nll_train, nll_val, nll_te+ , mdl_train, mdl_val, mdl_te]+ thetaStr = intercalate ";" $ Prelude.map show (MA.toList theta)+ io . putStrLn $ show ix <> "," <> showExpr expr <> ","+ <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)+ <> "," <> vals - tournament :: Int -> [EClassId] -> RndEGraph EClassId- tournament n ecIds = do - (c0:cs) <- replicateM n (rnd (randomFrom ecIds))- f0 <- gets (_fitness . _info . (IM.! c0) . _eClass)- snd <$> foldM (\(facc, acc) c -> gets (_fitness . _info . (IM.! c) . _eClass)- >>= \f -> if f > facc- then pure (f, c)- else pure (facc, acc)) (f0, c0) cs+ printPop pop = forM_ pop $ \ecN'-> do+ ecN'' <- canonical ecN'+ _tree <- getBestExpr ecN''+ fi <- fromJust <$> getFitness ecN''+ theta <- fromJust <$> getTheta ecN''+ let thetaStr = intercalate ";" $ Prelude.map show (MA.toList theta)+ io . putStrLn $ showExpr _tree <> "," <> thetaStr <> "," <> show fi+ pure () - insertRndParent :: RndEGraph ()- insertRndParent = do nt <- rnd rndNonTerm- meId <- case nt of- Uni f _ -> do ecIds <- getEClassesThat (isSizeOf (<10) &&& isValidFitness &&& funDoesNotExistWith nt)- if null ecIds- then pure Nothing - else do rndId <- tournament 5 ecIds- sz <- getSize rndId- if sz < 10 - then do node <- canonize (Uni f rndId)- Just <$> add myCost node- else pure Nothing- Bin op _ _ -> do ecIds <- getEClassesThat (isSizeOf (<9) &&& isValidFitness)- if null ecIds- then pure Nothing- else do rndIdLeft <- rnd $ randomFrom ecIds- sz1 <- getSize rndIdLeft- ecIds' <- getEClassesThat (isSizeOf (< (10 - sz1)) &&& isValidFitness &&& opDoesNotExistWith nt rndIdLeft)- if null ecIds'- then pure Nothing- else do rndIdRight <- rnd $ randomFrom ecIds'- rndIdRight <- tournament 5 ecIds'- sz2 <- getSize rndIdRight- if sz1 + sz2 < 10- then Just <$> (canonize (Bin op rndIdLeft rndIdRight)- >>= add myCost)- else pure Nothing- when (isJust meId) do- let eId = fromJust meId- eId' <- canonical eId- curFit <- gets (_fitness . _info . (IM.! eId') . _eClass)- when (isNothing curFit) do- t <- getBest eId'- f <- fitnessFun x y t- updateFitness f eId'- -- io $ print ('p', showExpr t, f)+ insertTerms =+ forM terms $ \t -> do fromTree myCost t >>= canonical - gpStep :: RndEGraph () - gpStep = do choice <- rnd $ randomFrom [2,2,3,3,3]- if | choice == 1 -> insertRndExpr- | choice == 2 -> insertRndParent- | otherwise -> evalRndSubTree- rebuild myCost- -} data Args = Args- { dataset :: String,- gens :: Int,- _alg :: Alg,- _maxSize :: Int,- _split :: Int+ { _dataset :: String,+ _testData :: String,+ _gens :: Int,+ _maxSize :: Int,+ _split :: Int,+ _printPareto :: Bool,+ _trace :: Bool,+ _distribution :: Distribution,+ _optIter :: Int,+ _optRepeat :: Int,+ _nPop :: Int,+ _nTournament :: Int,+ _pc :: Double,+ _pm :: Double,+ _nonterminals :: String,+ _dumpTo :: String,+ _loadFrom :: String,+ _moo :: Bool } deriving (Show) @@ -519,6 +381,12 @@ <> short 'd' <> metavar "INPUT-FILE" <> help "CSV dataset." )+ <*> strOption+ ( long "test"+ <> short 't'+ <> value ""+ <> showDefault+ <> help "test data") <*> option auto ( long "generations" <> short 'g'@@ -526,11 +394,6 @@ <> showDefault <> value 100 <> help "Number of generations." )- <*> option auto- ( long "algorithm"- <> short 'a'- <> metavar "ALG"- <> help "Algorithm." ) <*> option auto ( long "maxSize" <> short 's'@@ -538,65 +401,86 @@ <*> option auto ( long "split" <> short 'k'- <> help "k-split ratio training-test")--chunksOf :: Int -> [e] -> [[e]]-chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))- where- splitter :: [e] -> ([e] -> a -> a) -> a -> a- splitter [] _ n = n- splitter l c n = l `c` splitter (Prelude.drop i l) c n- build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]- build g = g (:) []--splitData :: SRMatrix -> PVector -> Int -> State StdGen (SRMatrix, SRMatrix, PVector, PVector)-splitData x y k = do if k == 1- then pure (x, x, y, y)- else do- ixs' <- (state . shuffle) [0 .. sz-1]- let ixs = chunksOf k ixs' -- $ sortOn (\ix -> y MA.! ix) [0 .. sz-1]- --ixs <- forM sortedIxs $ \is -> state (shuffle is)- let xl = MA.toLists x :: [MA.ListItem MA.Ix2 Double]- x_tr = MA.fromLists' comp_x [xl !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]- x_te = MA.fromLists' comp_x [xl !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]- y_tr = MA.fromList comp_y [y MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]- y_te = MA.fromList comp_y [y MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]-- {-- ixs <- state (shuffle [0 .. sz-1])- let ixs_tr = sort $ Prelude.take qty_tr ixs- ixs_te = sort $ Prelude.drop qty_tr ixs-- x_tr = MA.fromLists' comp_x [xl !! ix | ix <- ixs_tr]- x_te = MA.fromLists' comp_x [xl !! ix | ix <- ixs_te]- y_tr = MA.fromList comp_y [y MA.! ix | ix <- ixs_tr]- y_te = MA.fromList comp_y [y MA.! ix | ix <- ixs_te]- -}- pure (x_tr, x_te, y_tr, y_te)- where- (MA.Sz sz) = MA.size y- --qty_tr = round (thr * fromIntegral sz)- --qty_te = sz - qty_tr- comp_x = MA.getComp x- comp_y = MA.getComp y+ <> value 1 + <> showDefault+ <> help "k-split ratio training-validation")+ <*> switch+ ( long "print-pareto"+ <> help "print Pareto front instead of best found expression")+ <*> switch+ ( long "trace"+ <> help "print all evaluated expressions.")+ <*> option auto+ ( long "distribution"+ <> value Gaussian+ <> showDefault+ <> help "distribution of the data.")+ <*> option auto+ ( long "opt-iter"+ <> value 30+ <> showDefault+ <> help "number of iterations in parameter optimization.")+ <*> option auto+ ( long "opt-retries"+ <> value 1+ <> showDefault+ <> help "number of retries of parameter fitting.")+ <*> option auto+ ( long "nPop"+ <> value 100+ <> showDefault+ <> help "population size (Default: 100).")+ <*> option auto+ ( long "tournament-size"+ <> value 2+ <> showDefault+ <> help "tournament size.")+ <*> option auto+ ( long "pc"+ <> value 1.0+ <> showDefault+ <> help "probability of crossover.")+ <*> option auto+ ( long "pm"+ <> value 0.3+ <> showDefault+ <> help "probability of mutation.")+ <*> strOption+ ( long "non-terminals"+ <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+ <> showDefault+ <> help "set of non-terminals to use in the search."+ )+ <*> strOption+ ( long "dump-to"+ <> value ""+ <> showDefault+ <> help "dump final e-graph to a file."+ )+ <*> strOption+ ( long "load-from"+ <> value ""+ <> showDefault+ <> help "load initial e-graph from a file."+ )+ <*> switch+ ( long "moo"+ <> help "replace the current population with the pareto front instead of replacing it with the generated children."+ ) main :: IO () main = do- --args <- pure (Args "nikuradse_2.csv" 100) -- execParser opts args <- execParser opts- g <- getStdGen- ((x, y, _, _), _, _) <- loadDataset (dataset args) True- let ((x', x_te, y', y_te),g') = runState (splitData x y $ _split args) g- ((x_tr, x_val, y_tr, y_val),g'') = runState (splitData x' y' 2) g'- let (Sz2 _ nFeats) = MA.size x- terms = [var ix | ix <- [0 .. nFeats-1]] <> [param 0] -- [param ix | ix <- [0 .. 5]]- alg = evalStateT (egraphSearch (_alg args) x_tr y_tr x_val y_val x_te y_te terms (gens args) (_maxSize args)) emptyGraph- --(bestExpr, fit) <- evalStateT alg g- --printExpr bestExpr- --print fit- evalStateT alg g''-+ g <- getStdGen+ dataTrain' <- loadTrainingOnly (_dataset args) True+ dataTest <- if null (_testData args)+ then pure dataTrain'+ else loadTrainingOnly (_testData args) True+ let ((dataTrain, dataVal), g') = runState (splitData dataTrain' $ _split args) g+ alg = evalStateT (egraphGP dataTrain dataVal dataTest args) emptyGraph+ evalStateT alg g' where opts = Opt.info (opt <**> helper)- ( fullDesc <> progDesc "Very simple example of GP using SRTree."- <> header "tinyGP - a very simple example of GP using SRTRee." )+ ( fullDesc <> progDesc "An implementation of GP with modified crossover and mutation\+ \ operators designed to exploit equality saturation and e-graphs."+ <> header "GPEgg - Genetic Programming for Symbolic Regression using e-graphs." )
apps/egraphGP/Random.hs view
@@ -30,7 +30,7 @@ {-# INLINE randomFrom #-} randomVec :: Int -> Rng PVector-randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-3.0, 3.0))+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1)) randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree) randomTree minDepth maxDepth maxSize genTerm genNonTerm grow @@ -48,7 +48,7 @@ node <- genNonTerm case node of Uni f _ -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow - Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 2) genTerm genNonTerm grow+ Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow pure . Fix $ Bin op l r {-# INLINE randomTree #-}
+ apps/egraphGP/Util.hs view
@@ -0,0 +1,258 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Random+import Algorithm.SRTree.Likelihoods+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries++type RndEGraph a = EGraphST (StateT StdGen IO) a+type DataSet = (SRMatrix, PVector, Maybe PVector)++csvHeader :: String+csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"++io :: IO a -> RndEGraph a+io = lift . lift+{-# INLINE io #-}+rnd :: StateT StdGen IO a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _) = 1+myCost (Const _) = 1+myCost (Param _) = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t) = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()+while p arg prog = do when (p arg) do arg' <- prog arg+ while p arg' prog++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =+ if isNaN val -- || isNaN tr+ then (-(1/0), theta) -- infinity+ else (val, theta)+ where+ --tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+ evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+ --tr = evalF x y mYErr+ val = evalF x_val y_val mYErr_val++--{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nRep nIter distribution dataTrain dataVal _tree = do+ let tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+ let fits = maximumBy (compare `on` fst) $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal tree) thetaOrigs+ pure fits+--{-# INLINE fitnessFunRep #-}+++-- helper query functions+-- TODO: move to egraph lib+getFitness :: EClassId -> RndEGraph (Maybe Double)+getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: EClassId -> RndEGraph (Maybe PVector)+getTheta c = gets (_theta . _info . (IM.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: EClassId -> RndEGraph Int+getSize c = gets (_size . _info . (IM.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}++getBestFitness :: RndEGraph (Maybe Double)+getBestFitness = do+ bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+ gets (_fitness . _info . (IM.! bec) . _eClass)++-- TODO: move to dataset lib+chunksOf :: Int -> [e] -> [[e]]+chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))+ where+ splitter :: [e] -> ([e] -> a -> a) -> a -> a+ splitter [] _ n = n+ splitter l c n = l `c` splitter (Prelude.drop i l) c n+ build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]+ build g = g (:) []++splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)+splitData (x, y, mYErr) k = do+ if k == 1+ then pure ((x, y, mYErr), (x, y, mYErr))+ else do+ ixs' <- (state . shuffle) [0 .. sz-1]+ let ixs = chunksOf k ixs'++ let (x_tr, x_te) = getX ixs x+ (y_tr, y_te) = getY ixs y+ mY = fmap (getY ixs) mYErr+ (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)+ pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))+ where+ (MA.Sz sz) = MA.size y+ comp_x = MA.getComp x+ comp_y = MA.getComp y++ getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)+ getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]+ in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )+ getY :: [[Int]] -> PVector -> (PVector, PVector)+ getY ixs ys = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b++parseNonTerms :: String -> [SRTree ()]+parseNonTerms = Prelude.map toNonTerm . splitOn ","+ where+ binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+ uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+ toNonTerm xs' = let xs = Prelude.map toLower xs'+ in case binTerms Map.!? xs of+ Just op -> Bin op () ()+ Nothing -> case uniTerms Map.!? xs of+ Just f -> Uni f ()+ Nothing -> error $ "invalid non-terminal " <> show xs++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, PVector)) -> RndEGraph EClassId+insertExpr t fitFun = do+ ecId <- fromTree myCost t >>= canonical+ (f, p) <- fitFun t+ insertFitness ecId f p+ io . putStrLn $ "Best fit global: " <> show f+ pure ecId+ where powabs l r = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+ mf <- getFitness ec+ case mf of+ Nothing -> do+ t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p+ pure True+ Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+ if not (null ecIds)+ then do rndId' <- rnd $ randomFrom ecIds+ rndId <- canonical rndId'+ constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+ case constType of+ NotConst -> pure $ Just rndId+ _ -> pure Nothing+ else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)++getBestExprWithSize n =+ do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+ if (not (null ec))+ then do+ bestFit <- getFitness $ head ec+ bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+ pure [(head ec, bestFit)]+ else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+ do grow <- rnd toss+ n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+ t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+ fromTree myCost t >>= canonical++refit fitFun ec = do+ t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p++--printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest fitFun printExprFun = do+ bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+ bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)+ --refit fitFun bec+ --io.print $ "should be " <> show bestFit+ printExprFun 0 bec++--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))+ where+ go :: Int -> Int -> Double -> RndEGraph ()+ go n ix f+ | n > maxSize = pure ()+ | otherwise = do+ ecList <- getBestExprWithSize n+ if not (null ecList)+ then do let (ec, mf) = head ecList+ improved = fromJust mf > f+ ec' <- canonical ec+ when improved $ refit fitFun ec' >> printExprFun ix ec'+ go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))+ else go (n+1) ix f++evaluateUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ forM_ ec $ \c -> do+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p++evaluateRndUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ c <- rnd . randomFrom $ ec+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p+ pure c
+ apps/egraphSearch/Main.hs view
@@ -0,0 +1,325 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}++module Main where ++import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Simplify+import Algorithm.EqSat.Build+import Algorithm.EqSat.Queries+import Algorithm.EqSat.Info+import Algorithm.EqSat.DB+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection+import Control.Monad ( forM_, forM, when )+import Control.Monad.State.Strict+import qualified Data.IntMap.Strict as IM+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.Maybe ( fromJust, isNothing )+import Data.SRTree+import Data.SRTree.Print ( showExpr )+import Options.Applicative as Opt hiding (Const)+import Random+import System.Random+import Data.List ( intercalate )+import qualified Data.IntSet as IntSet+import Algorithm.EqSat (runEqSat)+import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS+import Debug.Trace+import qualified Data.HashSet as Set+import Control.Lens (over)++import Util++data Alg = OnlyRandom | BestFirst deriving (Show, Read, Eq)++egraphSearch :: DataSet -> DataSet -> DataSet -> Args -> StateT EGraph (StateT StdGen IO) ()+-- terms nEvals maxSize printPareto printTrace slowIter slowRep =+egraphSearch dataTrain dataVal dataTest args = do+ if null (_loadFrom args)+ then do ecFst <- insertRndExpr (_maxSize args) rndTerm rndNonTerm2+ --ecFst <- insertBestExpr -- use only to debug+ updateIfNothing fitFun ecFst+ insertTerms+ evaluateUnevaluated fitFun+ runEqSat myCost rewritesParams 1+ cleanDB+ else (io $ BS.readFile (_loadFrom args)) >>= \eg -> put (decode eg)+ nCls <- gets (IM.size . _eClass)+ nUnev <- gets (IntSet.size . _unevaluated . _eDB)+ let nEvs = nCls - nUnev++ while ((<(_gens args)) . snd) (10, nEvs) $+ \(radius, nEvs) ->+ do+ nCls <- gets (IM.size . _eClass)+ --nUnev <- gets (IntSet.size . _unevaluated . _eDB)+ --let nEvs = nCls - nUnev -- WARNING: see if it affects results++ ecN <- case (_alg args) of+ OnlyRandom -> do let ratio = fromIntegral nEvs / fromIntegral nCls+ b <- rnd (tossBiased ratio)+ ec <- if b && ratio > 0.99+ then insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+ else do mEc <- pickRndSubTree -- evaluateRndUnevaluated fitFun >>= canonical+ case mEc of+ Nothing ->insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+ Just ec' -> pure ec'+ pure ec+ BestFirst -> do+ ecsPareto <- getParetoEcsUpTo 50 (_maxSize args)+ ecPareto <- combineFrom ecsPareto >>= canonical+ curFitPareto <- getFitness ecPareto++ if isNothing curFitPareto+ then pure ecPareto+ else do ecsBest <- getTopFitEClassThat 100 (isSizeOf (<(_maxSize args)))+ ecBest <- combineFrom ecsBest >>= canonical+ curFitBest <- getFitness ecBest+ if isNothing curFitBest+ then pure ecBest+ else do ee <- pickRndSubTree+ case ee of+ Nothing -> insertRndExpr (_maxSize args) rndTerm rndNonTerm2 >>= canonical+ Just c -> pure c++ -- when upd $+ ecN' <- canonical ecN+ upd <- updateIfNothing fitFun ecN'+ when upd $ runEqSat myCost rewritesParams 1 >> cleanDB >> refitChanged+++ when (upd && (_trace args))+ do+ ecN'' <- canonical ecN'+ _tree <- getBestExpr ecN''+ fi <- negate . fromJust <$> getFitness ecN''+ theta <- fromJust <$> getTheta ecN''+ let thetaStr = intercalate ";" $ Prelude.map show (MA.toList theta)+ io . putStrLn $ showExpr _tree <> "," <> thetaStr <> "," <> show fi+ pure ()+ let radius' = if (not upd) then (max 10 $ min (500 `div` (_maxSize args)) (radius+1)) else (max 10 $ radius-1)+ nEvs' = nEvs + if upd then 1 else 0+ pure (radius', nEvs')+ io $ putStrLn csvHeader+ if (_printPareto args)+ then paretoFront (_maxSize args) printExpr+ else printBest printExpr+ when ((not.null) (_dumpTo args)) $ get >>= (io . BS.writeFile (_dumpTo args) . encode )+ where+ fitFun = fitnessFunRep (_optRepeat args) (_optIter args) (_distribution args) dataTrain dataVal++ refitChanged = do ids <- gets (_refits . _eDB) >>= Prelude.mapM canonical . Set.toList+ modify' $ over (eDB . refits) (const Set.empty)+ forM_ ids $ \ec -> do t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p++ combineFrom [] = pure 0 -- this is the first terminal and it will always be already evaluated+ combineFrom ecs = do+ nt <- rnd rndNonTerm+ p1 <- rnd (randomFrom ecs)+ p2 <- rnd (randomFrom ecs)+ l1 <- rnd (randomFrom [2..(_maxSize args)-2]) -- sz 10: [2..8]++ e1 <- randomChildFrom p1 l1 >>= canonical+ ml <- gets (_size . _info . (IM.! e1) . _eClass)+ l2 <- rnd (randomFrom [1..((_maxSize args) - ml - 1)]) -- maxSize - maxSize + 2 - 2= 0 -- sz 10: [1..7] (2) / [1..1] (8)+ e2 <- randomChildFrom p2 l2 >>= canonical+ case nt of+ Uni Id () -> canonical e1+ Uni f () -> add myCost (Uni f e1) >>= canonical+ Bin op () () -> do b <- rnd toss+ if b+ then add myCost (Bin op e1 e2) >>= canonical+ else add myCost (Bin op e2 e1) >>= canonical+ _ -> canonical e1 -- it is a terminal, should it happen?++ randomChildFrom ec' maxL = do+ p <- rnd toss -- whether to go deeper or return this level+ ec <- canonical ec'+ l <- gets (_size . _info . (IM.! ec) . _eClass )++ if p || l > maxL+ then do --enodes <- gets (_eNodes . (IM.! ec) . _eClass)+ enode <- gets (_best . _info . (IM.! ec) . _eClass) -- we should return the best otherwise we may build larger exprs+ case enode of+ Uni _ eci -> randomChildFrom eci maxL+ Bin _ ecl ecr -> do coin <- rnd toss+ if coin+ then randomChildFrom ecl maxL+ else randomChildFrom ecr maxL+ _ -> pure ec -- this shouldn't happen unless maxL==0+ else pure ec++ nonTerms = parseNonTerms (_nonterminals args)+ --[ Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () (), Bin PowerAbs () (), Uni Recip ()]+ (Sz2 _ nFeats) = MA.size (getX dataTrain)+ terms = if _distribution args == ROXY+ then [var 0, param 0]+ else [var ix | ix <- [0 .. nFeats-1]] <> [param 0]++ rndTerm = Random.randomFrom terms+ rndNonTerm = Random.randomFrom $ (Uni Id ()) : nonTerms+ rndNonTerm2 = Random.randomFrom nonTerms++ insertTerms =+ forM terms $ \t -> do fromTree myCost t >>= canonical++ printExpr :: Int -> EClassId -> RndEGraph ()+ printExpr ix ec = do + theta' <- gets (fromJust . _theta . _info . (IM.! ec) . _eClass)+ bestExpr <- getBestExpr ec+ let nParams = countParams bestExpr+ (MA.Sz nTheta) = MA.size theta'+ (_, theta) <- if (nParams /= nTheta)+ then fitFun bestExpr+ else pure (1.0, theta')++ let (x, y, mYErr) = dataTrain+ (x_val, y_val, mYErr_val) = dataVal+ (x_te, y_te, mYErr_te) = dataTest+ distribution = _distribution args+ best' = relabelParams bestExpr+ expr = paramsToConst (MA.toList theta) best'+ mse_train = mse x y best' theta+ mse_val = mse x_val y_val best' theta+ mse_te = mse x_te y_te best' theta+ r2_train = r2 x y best' theta+ r2_val = r2 x_val y_val best' theta+ r2_te = r2 x_te y_te best' theta+ nll_train = nll distribution mYErr x y best' theta+ nll_val = nll distribution mYErr_val x_val y_val best' theta+ nll_te = nll distribution mYErr_te x_te y_te best' theta+ mdl_train = mdl distribution mYErr x y theta best'+ mdl_val = mdl distribution mYErr_val x_val y_val theta best'+ mdl_te = mdl distribution mYErr_te x_te y_te theta best'+ vals = intercalate ","+ $ Prelude.map show [mse_train, mse_val, mse_te+ , r2_train, r2_val, r2_te+ , nll_train, nll_val, nll_te+ , mdl_train, mdl_val, mdl_te]+ thetaStr = intercalate ";" $ Prelude.map show (MA.toList theta)+ io . putStrLn $ show ix <> "," <> showExpr expr <> ","+ <> thetaStr <> "," <> show (countNodes $ convertProtectedOps expr)+ <> "," <> vals+++data Args = Args+ { _dataset :: String,+ _testData :: String,+ _gens :: Int,+ _alg :: Alg,+ _maxSize :: Int,+ _split :: Int,+ _printPareto :: Bool,+ _trace :: Bool,+ _distribution :: Distribution,+ _optIter :: Int,+ _optRepeat :: Int,+ _nonterminals :: String,+ _dumpTo :: String,+ _loadFrom :: String+ }+ deriving (Show)++-- parser of command line arguments+opt :: Parser Args+opt = Args+ <$> strOption+ ( long "dataset"+ <> short 'd'+ <> metavar "INPUT-FILE"+ <> help "CSV dataset." )+ <*> strOption+ ( long "test"+ <> short 't'+ <> value ""+ <> showDefault+ <> help "test data")+ <*> option auto+ ( long "generations"+ <> short 'g'+ <> metavar "GENS"+ <> showDefault+ <> value 100+ <> help "Number of generations." )+ <*> option auto+ ( long "algorithm"+ <> short 'a'+ <> metavar "ALG"+ <> help "Algorithm." )+ <*> option auto+ ( long "maxSize"+ <> short 's'+ <> help "max-size." )+ <*> option auto+ ( long "split"+ <> short 'k'+ <> value 1 + <> showDefault+ <> help "k-split ratio training-validation")+ <*> switch+ ( long "print-pareto"+ <> help "print Pareto front instead of best found expression")+ <*> switch+ ( long "trace"+ <> help "print all evaluated expressions.")+ <*> option auto+ ( long "distribution"+ <> value Gaussian+ <> showDefault+ <> help "distribution of the data.")+ <*> option auto+ ( long "opt-iter"+ <> value 30+ <> showDefault+ <> help "number of iterations in parameter optimization.")+ <*> option auto+ ( long "opt-retries"+ <> value 1+ <> showDefault+ <> help "number of retries of parameter fitting.")+ <*> strOption+ ( long "non-terminals"+ <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+ <> showDefault+ <> help "set of non-terminals to use in the search."+ )+ <*> strOption+ ( long "dump-to"+ <> value ""+ <> showDefault+ <> help "dump final e-graph to a file."+ )+ <*> strOption+ ( long "load-from"+ <> value ""+ <> showDefault+ <> help "load initial e-graph from a file."+ )++main :: IO ()+main = do+ args <- execParser opts+ g <- getStdGen+ dataTrain' <- loadTrainingOnly (_dataset args) True+ dataTest <- if null (_testData args)+ then pure dataTrain'+ else loadTrainingOnly (_testData args) True+ let ((dataTrain, dataVal), g') = runState (splitData dataTrain' $ _split args) g+ alg = evalStateT (egraphSearch dataTrain dataVal dataTest args) emptyGraph+ evalStateT alg g'++ where+ opts = Opt.info (opt <**> helper)+ ( fullDesc <> progDesc "Symbolic Regression search algorithm\+ \ exploiting the potentials of equality saturation\+ \ and e-graphs."+ <> header "SymREgg - symbolic regression with e-graphs."+ )
+ apps/egraphSearch/Random.hs view
@@ -0,0 +1,54 @@+module Random where ++import System.Random +import Control.Monad.State.Strict+import Control.Monad+import Data.SRTree +import Data.SRTree.Eval+import Data.Massiv.Array as MA++type Rng a = StateT StdGen IO a ++toss :: Rng Bool+toss = state random+{-# INLINE toss #-}++tossBiased :: Double -> Rng Bool+tossBiased p = do r <- state random + pure (r < p)++randomVal :: Rng Double+randomVal = state random++randomRange :: (Ord val, Random val) => (val, val) -> Rng val+randomRange rng = state (randomR rng)+{-# INLINE randomRange #-}++randomFrom :: [a] -> Rng a+randomFrom funs = do n <- randomRange (0, length funs - 1)+ pure $ funs !! n+{-# INLINE randomFrom #-}++randomVec :: Int -> Rng PVector+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))++randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree)+randomTree minDepth maxDepth maxSize genTerm genNonTerm grow + | noSpaceLeft = genTerm+ | needNonTerm = genRecursion + | otherwise = do r <- toss+ if r + then genTerm + else genRecursion+ where + noSpaceLeft = maxDepth <= 1 || maxSize <= 2+ needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2++ genRecursion = do + node <- genNonTerm+ case node of + Uni f _ -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow + Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow+ r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow + pure . Fix $ Bin op l r+{-# INLINE randomTree #-}
+ apps/egraphSearch/Util.hs view
@@ -0,0 +1,254 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Random+import Algorithm.SRTree.Likelihoods+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries++type RndEGraph a = EGraphST (StateT StdGen IO) a+type DataSet = (SRMatrix, PVector, Maybe PVector)++csvHeader :: String+csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"++io :: IO a -> RndEGraph a+io = lift . lift+{-# INLINE io #-}+rnd :: StateT StdGen IO a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _) = 1+myCost (Const _) = 1+myCost (Param _) = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t) = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f ()+while p arg prog = do when (p arg) do arg' <- prog arg+ while p arg' prog++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) _tree thetaOrig =+ if isNaN val || isNaN tr+ then (-(1/0), theta) -- infinity+ else (val, theta) -- (min tr val, theta)+ where+ tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+ evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+ tr = evalF x y mYErr+ val = evalF x_val y_val mYErr_val++{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nRep nIter distribution dataTrain dataVal _tree = do+ let tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+ let fits = Prelude.map (fitnessFun nIter distribution dataTrain dataVal _tree) thetaOrigs+ pure (maximumBy (compare `on` fst) fits)+{-# INLINE fitnessFunRep #-}++--fitnessMV :: Int -> Int -> Distribution -> [DataSet] -> [DataSet] -> Fix SRTree -> RndEGraph (Double, [PVector])+--fitnessMV nRep nIter distribution dataTrains dataVals _tree = do+-- response <- forM (zip dataTrains dataVals) $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv _tree+-- pure (minimum (map fst response), map snd response)++-- helper query functions+-- TODO: move to egraph lib+getFitness :: EClassId -> RndEGraph (Maybe Double)+getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: EClassId -> RndEGraph (Maybe PVector)+getTheta c = gets (_theta . _info . (IM.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: EClassId -> RndEGraph Int+getSize c = gets (_size . _info . (IM.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}++getBestFitness :: RndEGraph (Maybe Double)+getBestFitness = do+ bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+ gets (_fitness . _info . (IM.! bec) . _eClass)++-- TODO: move to dataset lib+chunksOf :: Int -> [e] -> [[e]]+chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))+ where+ splitter :: [e] -> ([e] -> a -> a) -> a -> a+ splitter [] _ n = n+ splitter l c n = l `c` splitter (Prelude.drop i l) c n+ build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]+ build g = g (:) []++splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)+splitData (x, y, mYErr) k = do+ if k == 1+ then pure ((x, y, mYErr), (x, y, mYErr))+ else do+ ixs' <- (state . shuffle) [0 .. sz-1]+ let ixs = chunksOf k ixs'++ let (x_tr, x_te) = getX ixs x+ (y_tr, y_te) = getY ixs y+ mY = fmap (getY ixs) mYErr+ (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)+ pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))+ where+ (MA.Sz sz) = MA.size y+ comp_x = MA.getComp x+ comp_y = MA.getComp y++ getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)+ getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]+ in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )+ getY :: [[Int]] -> PVector -> (PVector, PVector)+ getY ixs ys = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b++parseNonTerms :: String -> [SRTree ()]+parseNonTerms = Prelude.map toNonTerm . splitOn ","+ where+ binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+ uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+ toNonTerm xs' = let xs = Prelude.map toLower xs'+ in case binTerms Map.!? xs of+ Just op -> Bin op () ()+ Nothing -> case uniTerms Map.!? xs of+ Just f -> Uni f ()+ Nothing -> error $ "invalid non-terminal " <> show xs++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, PVector)) -> RndEGraph EClassId+insertExpr t fitFun = do+ ecId <- fromTree myCost t >>= canonical+ (f, p) <- fitFun t+ insertFitness ecId f p+ io . putStrLn $ "Best fit global: " <> show f+ pure ecId+ where powabs l r = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+ mf <- getFitness ec+ case mf of+ Nothing -> do+ t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p+ pure True+ Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+ if not (null ecIds)+ then do rndId' <- rnd $ randomFrom ecIds+ rndId <- canonical rndId'+ constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+ case constType of+ NotConst -> pure $ Just rndId+ _ -> pure Nothing+ else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)++getBestExprWithSize n =+ do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+ if (not (null ec))+ then do+ bestFit <- getFitness $ head ec+ bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+ (:[]) . (,bestP) . (,bestFit) . (,ec) <$> getBestExpr (head ec)+ else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+ do grow <- rnd toss+ n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+ t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+ fromTree myCost t >>= canonical++printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest printExprFun = do+ bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+ printExprFun 0 bec++paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront maxSize printExprFun = go 1 0 (-(1.0/0.0))+ where+ go :: Int -> Int -> Double -> RndEGraph ()+ go n ix f+ | n > maxSize = pure ()+ | otherwise = do+ ecList <- getBestExprWithSize n+ if not (null ecList)+ then do let (((_, ec), mf), _) = head ecList+ improved = fromJust mf > f+ ec' <- traverse canonical ec+ when improved $ printExprFun ix (head ec')+ go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))+ else go (n+1) ix f++evaluateUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ forM_ ec $ \c -> do+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p++evaluateRndUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ c <- rnd . randomFrom $ ec+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p+ pure c
− apps/ieeexplore/Main.hs
@@ -1,101 +0,0 @@-{-# LANGUAGE BlockArguments #-}-{-# LANGUAGE TupleSections #-}-{-# LANGUAGE MultiWayIf #-}-{-# LANGUAGE OverloadedStrings #-}--module Main where --import Algorithm.EqSat.Egraph-import Algorithm.EqSat.Simplify-import Algorithm.SRTree.Likelihoods-import Algorithm.SRTree.Opt-import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.))-import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_)-import Control.Monad.State-import qualified Data.IntMap as IM-import Data.Massiv.Array as MA hiding (forM_, forM)-import Data.Maybe (fromJust, isNothing, isJust)-import Data.SRTree-import Data.SRTree.Datasets-import Data.SRTree.Eval-import Data.SRTree.Random (randomTree)-import Data.SRTree.Print-import Options.Applicative as Opt hiding (Const)-import System.Random-import qualified Data.Set as Set-import Data.List ( sort )--import Debug.Trace-import Algorithm.EqSat (runEqSat)---type RangedTree a = a --data QueryDB = QDB { _mseDB :: RangedTree Double - , _maeDB :: RangedTree Double - , _r2DB :: RangedTree Double - , _mdlDB :: RangedTree Double - , _bicDB :: RangedTree Double - , _aicDB :: RangedTree Double - , _lenDB :: RangedTree Int- , _ptDB :: RangedTree Int -- DB - } --data ExtraInfo = ExtraInfo { _thetaMap :: IM.IntMap PVector- , _qdb :: QueryDB- }--type InfoEGraph a = EGraphST (StateT ExtraInfo IO) a----io = lift . lift---{-# INLINE io #-}---ext = lift---{-# INLINE ext #-}--myCost :: SRTree Int -> Int-myCost (Var _) = 1-myCost (Const _) = 1-myCost (Param _) = 1-myCost (Bin op l r) = 2 + l + r-myCost (Uni _ t) = 3 + t----data Args = Args- { dataset :: String,- gens :: Int,- _maxSize :: Int- }- deriving (Show)---- parser of command line arguments-opt :: Parser Args-opt = Args- <$> strOption- ( long "dataset"- <> short 'd'- <> metavar "INPUT-FILE"- <> help "CSV dataset." )- <*> option auto- ( long "generations"- <> short 'g'- <> metavar "GENS"- <> showDefault- <> value 100- <> help "Number of generations." )- <*> option auto- ( long "maxSize"- <> short 's'- <> help "max-size." )--main :: IO ()-main = do- --args <- pure (Args "nikuradse_2.csv" 100) -- execParser opts- args <- execParser opts- ((x, y, _, _), _, _) <- loadDataset (dataset args) True- print "opa"-- where- opts = Opt.info (opt <**> helper)- ( fullDesc <> progDesc "Very simple example of GP using SRTree."- <> header "tinyGP - a very simple example of GP using SRTRee." )
+ apps/rEGGression/Commands.hs view
@@ -0,0 +1,393 @@+{-# language OverloadedStrings #-}+{-# language TupleSections #-}++module Commands where++import Control.Applicative ((<|>))+import Data.Attoparsec.ByteString.Char8 hiding ( match )+import qualified Data.ByteString.Char8 as B+import Data.Maybe+import Text.Read ( readMaybe )+import Data.Monoid (All(..))+import qualified Data.IntMap.Strict as IntMap+import qualified Data.IntSet as IntSet+import Control.Monad.State.Strict+import Control.Monad ( forM_ )+import Data.Char ( toUpper )+import qualified Data.Map as Map+import qualified Data.HashSet as Set+import qualified Data.Massiv.Array as MA+import Data.List ( nub, sortOn )+import Data.List.Split ( splitOn )++import Data.SRTree+import Data.SRTree.Datasets+import Data.SRTree.Recursion+import Data.SRTree.Eval+import Data.SRTree.Print hiding ( printExpr )+import Text.ParseSR (SRAlgs(..), parseSR, parsePat, Output(..), showOutput)++import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.Opt++import Algorithm.EqSat+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info+import Algorithm.EqSat.Queries+import Algorithm.EqSat.DB+import Algorithm.EqSat.Simplify++import Algorithm.SRTree.ModelSelection++import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS++import Util++-- * Parsing++-- top 5 by fitness|mdl [less than 5 params, less than 10 nodes]+data Command = Top Int Filter Criteria PatStr+ | Distribution FilterDist (Maybe Limit)+ -- below these will not be a parsable command+ | Report EClassId ArgOpt+ | Optimize EClassId Int ArgOpt+ | Insert String ArgOpt+ | Subtrees EClassId+ | Pareto Criteria+ | CountPat String+ | Save String+ | Load String+ | Import String Distribution String Bool++type Filter = EClass -> Bool -- pattern?+type FilterDist = Int -> Bool+data Criteria = ByFitness | ByDL deriving Eq+data Limit = Limit Int Bool deriving Show+data PatStr = PatStr String Bool | AntiPatStr String Bool | NoPat+type ArgOpt = (Distribution, DataSet, DataSet)++-- top 10 with <=10|=10 size with <=4 parameters by fitness|dl matching pat+-- report id+-- optimize id+-- insert eq+-- subtrees id+-- distribution with size <=10 limited at 10 asc|dsc+--+++parseCmd parser = eitherResult . (`feed` "") . parse parser . putEOL . B.strip++stripSp = many' (char ' ')++parseTop = do n <- decimal+ stripSp+ filters <- many' parseFilter+ stripSp+ criteria <- fromMaybe ByFitness . listToMaybe <$> many' parseCriteria+ stripSp+ pats' <- many' (parsePattern <|> parseAnti)+ pats <- case pats' of+ [] -> pure $ NoPat+ (x:_) -> pure $ x+ pure $ Top n (getAll . mconcat filters) criteria pats++parseDist = do filters' <- many' parseFilterDist+ let filters = if null filters'+ then [(\pat -> All $ pat <= 10)]+ else filters'+ stripSp+ limit <- listToMaybe <$> many' parseLimit+ pure $ Distribution (getAll . mconcat filters) limit++parseFilter = do stringCI "with"+ stripSp+ field <- parseSz <|> parseCost <|> parseParams+ stripSp+ cmp <- parseCmp+ stripSp+ pure (\ec -> All $ cmp (field ec))+parseFilterDist = do stringCI "with"+ stripSp+ stringCI "size"+ stripSp+ cmp <- parseCmp+ stripSp+ pure (\pat -> All $ cmp pat)++parseSz = stringCI "size" >> pure (_size . _info)+parseCost = stringCI "cost" >> pure (_cost . _info)+parseParams = stringCI "parameters" >> pure (mbLen . _theta . _info)+ where+ mbLen Nothing = 0+ mbLen (Just ps) = MA.unSz $ MA.size ps+parseCmp = do op <- parseLEQ <|> parseLT <|> parseEQ <|> parseGEQ <|> parseGT+ stripSp+ n <- decimal+ pure (`op` n)++parseLT = string "<" >> pure (<)+parseLEQ = string "<=" >> pure (<=)+parseEQ = string "=" >> pure (==)+parseGEQ = string ">=" >> pure (>=)+parseGT = string ">" >> pure (>)++parsePattern = do stringCI "matching"+ stripSp+ b <- option True parseRoot+ pat <- many' anyChar+ pure $ PatStr pat b+parseAnti = do stringCI "not matching"+ stripSp+ b <- option True parseRoot+ pat <- many' anyChar+ pure $ AntiPatStr pat b++parseLimit = do stringCI "limited at"+ stripSp+ n <- decimal+ stripSp+ ascOrdsc <- stringCI "asc" <|> stringCI "dsc"+ pure $ Limit n (ascOrdsc == "asc")+parseRoot = do stringCI "root"+ stripSp+ pure False+parseCriteria = parseByFit <|> parseByDL+parseByFit = do stringCI "by fitness"+ pure ByFitness+parseByDL = do stringCI "by dl"+ pure ByDL++putEOL :: B.ByteString -> B.ByteString+putEOL bs | B.last bs == '\n' = bs+ | otherwise = B.snoc bs '\n'++-- running+run (Top n filters criteria NoPat) = do+ let getFun = if criteria == ByFitness then getTopFitEClassThat else getTopDLEClassThat+ ids <- egraph $ getFun n filters+ printSimpleMultiExprs $ reverse ids++run (Top n filters criteria withPat) = do+ let (pat', getFun, isParents) =+ case withPat of+ PatStr p parent -> (p, if criteria == ByFitness then getTopFitEClassIn else getTopDLEClassIn, parent)+ AntiPatStr p parent -> (p, if criteria == ByFitness then getTopFitEClassNotIn else getTopDLEClassNotIn, parent)++ let etree = parsePat $ B.pack pat'+ case etree of+ Left _ -> io.putStrLn $ "no parse for " <> pat'+ Right pat -> do+ ecs' <- egraph $ (Prelude.map fromLeft . Prelude.filter isLeft . Prelude.map snd) <$> match pat+ ecs <- egraph $ Prelude.mapM canonical ecs'+ >>= getParents isParents+ ids <- egraph $ getFun n filters ecs+ printSimpleMultiExprs (reverse $ nub ids)++run (Distribution pSz mLimit) = do+ ee <- egraph $ IntSet.toList . IntSet.fromList <$> getAllEvaluatedEClasses+ allPats <- egraph $ getAllPatternsFrom pSz Map.empty ee+ let (n, isAsc) = case mLimit of+ Nothing -> (Map.size allPats, True)+ Just (Limit sz asc) -> (sz, asc)+ printMultiCounts (Prelude.take n+ $ sortOn (if isAsc then snd else negate . snd)+ $ Map.toList+ $ Map.filterWithKey (\k v -> k /= VarPat 'A' && pSz (lenPat k))+ allPats)++run (Report eid (dist, trainData, testData)) = egraph $ printExpr trainData testData dist eid++run (Optimize eid nIters (dist, trainData@(x, y, mYErr), testData)) = do -- dist trainData testData+ t <- egraph $ relabelParams <$> getBestExpr eid+ (f, theta) <- egraph $ fitnessFunRep nIters dist trainData t+ egraph $ insertFitness eid f theta+ let mdl_train = mdl dist mYErr x y theta t+ egraph $ insertDL eid mdl_train+ printSimpleMultiExprs [eid]++run (Insert expr argOpt) = do+ let etree = parseSR TIR "" False $ B.pack expr+ case etree of+ Left _ -> io.putStrLn $ "no parse for " <> expr+ Right tree -> do eid <- egraph $ fromTree myCost tree+ run (Optimize eid 100 argOpt)++run (Subtrees eid) = do+ isValid <- egraph $ gets ((IntMap.member eid) . _eClass)+ if isValid+ then do ids <- egraph $ getAllChildEClasses eid+ printSimpleMultiExprs ids+ else io.putStrLn $ "Invalid id."++run (Pareto crit) = do+ maxSize <- egraph $ gets (fst . IntMap.findMax . _sizeFitDB . _eDB)+ ecs <- egraph $ case crit of+ ByFitness -> getParetoEcsUpTo True 1 maxSize+ ByDL -> getParetoEcsUpTo False 1 maxSize+ printSimpleMultiExprs ecs++run (CountPat spat) = do+ let etree = parsePat $ B.pack spat+ case etree of+ Left _ -> io.putStrLn $ "no parse for " <> spat+ Right pat -> do (p, cnt) <- countPattern pat+ io . putStrLn $ spat <> " appears in " <> show cnt <> " equations."++run (Save fname) = do+ eg <- egraph get+ io $ BS.writeFile fname (encode eg)++run (Load fname) = do+ eg <- io $ BS.readFile fname+ egraph $ put (decode eg)++run (Import fname dist varnames params) = do+ egraph $ importCSV dist fname varnames params++-- * auxiliary functions+importCSV :: Distribution -> String -> String -> Bool -> RndEGraph ()+importCSV dist fname hdr convertParam = cleanDB >> parseEqs >> createDB >> rebuildAllRanges+ where+ alg = getFormat fname++ toTuple :: [String] -> (String, [Double], Double)+ toTuple [eq, t, f] = (eq, Prelude.map Prelude.read $ Prelude.filter (not.null) $ splitOn ";" t, fromMaybe (-1.0/0.0) $ readMaybe f)+ toTuple xss = error $ show xss++ parseEqs :: RndEGraph ()+ parseEqs = do content <- Prelude.map (toTuple . splitOn ",") . lines <$> (liftIO $ readFile fname)+ forM_ content $ \(eq, params, f) -> do+ case parseSR alg (B.pack hdr) False (B.pack eq) of+ Left _ -> liftIO $ putStrLn $ "Skippping " <> eq+ Right tree' -> do+ let (tree, ps) = if convertParam then floatConstsToParam tree' else (tree', theta)+ theta = if convertParam then if dist==MSE then ps <> params else ps else params+ eid <- fromTree myCost tree >>= canonical+ insertFitness eid f $ MA.fromList MA.Seq theta+ runEqSat myCost rewritesParams 1+ cleanDB+++parseCSV :: Distribution -> String -> String -> Bool -> IO EGraph+parseCSV dist fname hdr convertParam = execStateT parseEqs emptyGraph+ where+ alg = getFormat fname++ toTuple :: [String] -> (String, [Double], Double)+ toTuple [eq, t, f] = (eq, Prelude.map Prelude.read $ Prelude.filter (not.null) $ splitOn ";" t, fromMaybe (-1.0/0.0) $ readMaybe f)+ toTuple xss = error $ show xss++ parseEqs :: RndEGraph ()+ parseEqs = do content <- Prelude.map (toTuple . splitOn ",") . lines <$> (liftIO $ readFile fname)+ forM_ content $ \(eq, params, f) -> do+ case parseSR alg (B.pack hdr) False (B.pack eq) of+ Left _ -> liftIO $ putStrLn $ "Skippping " <> eq+ Right tree' -> do+ let (tree, ps) = if convertParam then floatConstsToParam tree' else (tree', theta)+ theta = if convertParam then if dist==MSE then ps <> params else ps else params+ eid <- fromTree myCost tree >>= canonical+ insertFitness eid f $ MA.fromList MA.Seq theta+ runEqSat myCost rewritesParams 1+ cleanDB+getFormat :: String -> SRAlgs+getFormat = Prelude.read . Prelude.map toUpper . Prelude.last . splitOn "."++++convert :: String -> Output -> String -> IO ()+convert fname out hdr = do+ let alg = getFormat fname+ content <- Prelude.map (toTuple . splitOn ",") . lines <$> readFile fname+ forM_ content $ \(eq, params, f) -> do+ case parseSR alg (B.pack hdr) False (B.pack eq) of+ Left _ -> pure ()+ Right tree -> do+ putStr (showOutput out tree)+ putChar ','+ putStr params+ putChar ','+ putStrLn f+ where+ toTuple :: [String] -> (String, String, String)+ toTuple [eq, t, f] = (eq, t, f)+ toTuple xss = error $ show xss++getParents False ecs = pure ecs+getParents True ecs = IntSet.toList <$> getParentsOf (IntSet.fromList ecs) 500000 (IntSet.fromList ecs)++isBest (e', en') = do e <- canonical e'+ best <- gets (_best . _info . (IntMap.! e) . _eClass) >>= canonize+ en <- canonize en'+ pure (en == best)++getParentsOf :: IntSet.IntSet -> Int -> IntSet.IntSet -> RndEGraph IntSet.IntSet+getParentsOf visited n queue | IntSet.size queue >= n = pure queue+getParentsOf visited n queue =+ do ecs <- Prelude.mapM canonical (IntSet.toList queue)+ parents' <- IntSet.unions <$> Prelude.mapM canonizeParents ecs+ uneval <- gets (_unevaluated . _eDB)+ grandParents <- getParentsOf (IntSet.union visited parents') (n-1) parents'+ pure (IntSet.filter (not . (`IntSet.member` uneval)) $ (queue <> grandParents))+ where+ canonizeParents ec = do parents <- gets (_parents . (IntMap.! ec) . _eClass)+ >>= Prelude.mapM (\(e, en) -> isBest (e, en) >>= \b -> pure (e, en, b)) . Set.toList+ >>= pure . Set.fromList++ pure $ IntSet.fromList . Prelude.map (\(e, _, _) -> e) . Set.toList $+ Set.filter (\(e, en, b) -> b && ec `Prelude.elem` (childrenOf en)+ && not (e `IntSet.member` visited)+ ) parents+++isLeft (Left _) = True+isLeft _ = False+fromLeft (Left x) = x+fromLeft _ = undefined++getAllPatternsFrom :: Monad m => (Int -> Bool) -> Map.Map Pattern Int -> [EClassId] -> EGraphST m (Map.Map Pattern Int)+getAllPatternsFrom pSz counts [] = pure counts+getAllPatternsFrom pSz counts (x:xs) = do pats <- Map.fromListWith (+) . Prelude.map (,1) . Prelude.filter (pSz . lenPat) <$> getAllPatterns (<=max_pat) x+ getAllPatternsFrom pSz (Map.unionWith (+) pats counts) xs+ where max_pat = 10++relabelVarPat :: Pattern -> Pattern+relabelVarPat t = alg t `evalState` 65+ where+ alg :: Pattern -> State Int Pattern+ alg (VarPat _) = do ix <- Control.Monad.State.Strict.get; Control.Monad.State.Strict.modify (+1); pure (VarPat $ toEnum ix)+ alg (Fixed (Uni f t')) = do t <- alg t'; pure $ Fixed (Uni f t)+ alg (Fixed (Bin op l' r')) = do l <- alg l'; r <- alg r'; pure $ Fixed (Bin op l r)+ alg pt = pure pt++lenPat :: Pattern -> Int+lenPat (Fixed (Uni _ t)) = 1 + lenPat t+lenPat (Fixed (Bin _ l r)) = 1 + lenPat l + lenPat r+lenPat _ = 1++countPattern pat = do+ ecs' <- egraph $ (Prelude.map fromLeft . Prelude.filter isLeft . Prelude.map snd) <$> match pat+ ecs <- egraph $ Prelude.mapM canonical ecs'+ >>= getEvaluated+ pure (pat, IntSet.size ecs)++getEvaluated ecs = getParentsOf (IntSet.fromList ecs) 500000 (IntSet.fromList ecs)++getAllPatterns :: Monad m => (Int -> Bool) -> EClassId -> EGraphST m [Pattern]+getAllPatterns pSz eid = do+ eid' <- canonical eid+ best <- gets (_best . _info . (IntMap.! eid') . _eClass)+ case best of+ Var ix -> pure [VarPat 'A', Fixed (Var ix)]+ Param ix -> pure [VarPat 'A', Fixed (Param ix)]+ Const x -> pure [VarPat 'A', Fixed (Const x)]+ Uni f t -> do pats <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz t+ pure (VarPat 'A' : [relabelVarPat $ Fixed (Uni f t') | t' <- pats])+ Bin op l r | l==r -> do pats <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz l+ pure (VarPat 'A' : [relabelVarPat $ Fixed (Bin op l' l') | l' <- pats])+ | otherwise -> do patsL <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz l+ patsR <- Prelude.filter (pSz . lenPat) <$> getAllPatterns pSz r+ pure (VarPat 'A' : [relabelVarPat $ Fixed (Bin op l' r') | l' <- patsL, r' <- patsR])++
+ apps/rEGGression/Main.hs view
@@ -0,0 +1,336 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TypeSynonymInstances #-}+{-# LANGUAGE TypeApplications #-}++module Main where ++import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.Opt+import Control.Monad.State.Strict++import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info+import Algorithm.EqSat.Queries+import Algorithm.EqSat.DB+import Algorithm.EqSat.Simplify++import qualified Data.IntMap as IM+import Data.Massiv.Array as MA hiding (forM_, forM, Continue, convert)+import Data.Maybe (fromJust, isNothing, isJust)+import Data.SRTree+import Data.SRTree.Recursion+import Data.SRTree.Datasets+import Data.SRTree.Eval+import Data.SRTree.Random (randomTree)+import Data.SRTree.Print hiding ( printExpr )+import Options.Applicative as Opt hiding (Const, columns)+import System.Random+import qualified Data.HashSet as Set+import Data.List ( sort, sortOn )+import Data.List.Split ( splitOn )+import qualified Data.Map as Map+import Data.Map ( Map )+import qualified Data.IntMap.Strict as IntMap+import Data.Char ( toLower, toUpper )+import Debug.Trace+import Algorithm.EqSat (runEqSat)++import System.Console.Repline hiding (Repl)+import Util+import Commands+import Data.List ( isPrefixOf, intercalate, nub )+import Text.Read+import Control.Monad ( forM, when, forM_ )+import Data.Binary ( encode, decode )+import qualified Data.ByteString.Lazy as BS+import Data.Maybe ( fromMaybe )+import Text.ParseSR (SRAlgs(..), parseSR, parsePat, Output(..), showOutput)+import qualified Data.ByteString.Char8 as B+import qualified Data.IntSet as IntSet+import qualified Data.Set as SSet++data Args = Args+ { _dataset :: String,+ _testData :: String,+ _distribution :: Distribution,+ _dumpTo :: String,+ _loadFrom :: String,+ _parseCSV :: String,+ _convertFromTo :: String,+ _parseParams :: Bool,+ _to :: Output,+ _calcDL :: Bool+ }+ deriving (Show)++-- parser of command line arguments+opt :: Parser Args+opt = Args+ <$> strOption+ ( long "dataset"+ <> short 'd'+ <> metavar "INPUT-FILE"+ <> help "CSV dataset." )+ <*> strOption+ ( long "test"+ <> short 't'+ <> value ""+ <> showDefault+ <> help "test data")+ <*> option auto+ ( long "distribution"+ <> value Gaussian+ <> showDefault+ <> help "distribution of the data.")+ <*> strOption+ ( long "dump-to"+ <> value ""+ <> showDefault+ <> help "dump final e-graph to a file."+ )+ <*> strOption+ ( long "load-from"+ <> value ""+ <> showDefault+ <> help "load initial e-graph from a file."+ )+ <*> strOption+ ( long "parse-csv"+ <> value ""+ <> showDefault+ <> help "parse-csv CSV file with the format expression,parameters,fitness. The fitness value should be maximization and the parameters a ; separated list (there must be an additional parameter for sigma in the Gaussian distribution). The format of the equation is determined by the extension of the file, supported extensions are operon, pysr, tir, itea, hl (heuristiclab), gomea, feat, etc."+ )+ <*> strOption+ ( long "convert"+ <> value ""+ <> showDefault+ <> help "convert FROM TO, converts equation format from a given format (see 'parse-csv') to either 'math' or 'numpy'. The 'math' format is compatible with the tir format, so you can use this to standardize the equations from multiple sources into a single file. The output will be written to stdout."+ )+ <*> switch+ ( long "parse-parameters"+ <> help "Extract the numerical parameters from the expression. In this case the csv file should be formatted as \"equation,error,fitness, where 'error' is the error term used in Gaussia likelihood, it can be empty if using other distributions.\""+ )+ <*> option auto+ ( long "to"+ <> value MATH+ <> showDefault+ <> help "Format to convert to."+ )+ <*> switch+ ( long "calculate-dl"+ <> help "(re)calculate DL."+ )++runIfRight cmd = case cmd of+ Left err -> io.print $ "wrong command format."+ Right c -> run c++topCmd :: [String] -> Repl ()+topCmd [] = helpCmd ["top"]+topCmd args = do+ let cmd = parseCmd parseTop (B.pack $ unwords args)+ runIfRight cmd++distCmd :: [String] -> Repl ()+distCmd [] = helpCmd ["distribution"]+distCmd args = do+ let cmd = parseCmd parseDist (B.pack $ unwords args)+ runIfRight cmd++reportCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+reportCmd _ _ _ [] = helpCmd ["report"]+reportCmd dist trainData testData args =+ case readMaybe @Int (head args) of+ Nothing -> io.putStrLn $ "The id must be an integer."+ Just n -> run (Report n (dist, trainData, testData))++optimizeCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+optimizeCmd _ _ _ [] = helpCmd ["optimize"]+optimizeCmd dist trainData testData args =+ case readMaybe @Int (head args) of+ Nothing -> io.putStrLn $ "The id must be an integer."+ Just n -> do let nIters = if length args > 1 then fromMaybe 100 (readMaybe @Int (args !! 1)) else 100+ run (Optimize n nIters (dist, trainData, trainData))++subtreesCmd :: [String] -> Repl ()+subtreesCmd [] = helpCmd ["subtrees"]+subtreesCmd (arg:_) = case readMaybe @Int arg of+ Nothing -> io.putStrLn $ "The argument must be an integer."+ Just n -> run (Subtrees n)++insertCmd :: Distribution -> DataSet -> DataSet -> [String] -> Repl ()+insertCmd dist trainData testData [] = helpCmd ["insert"]+insertCmd dist trainData testData args = do+ let etree = parseSR TIR "" False $ B.pack (unwords args)+ case etree of+ Left _ -> io.putStrLn $ "no parse for " <> unwords args+ Right tree -> do ec <- egraph $ fromTree myCost tree+ run (Optimize ec 100 (dist, trainData, trainData))++paretoCmd :: [String] -> Repl ()+paretoCmd [] = run (Pareto ByFitness)+paretoCmd args = case (Prelude.map toLower $ unwords args) of+ "by fitness" -> run (Pareto ByFitness )+ "by dl" -> run (Pareto ByDL)+ _ -> helpCmd ["pareto"]++countPatCmd :: [String] -> Repl ()+countPatCmd [] = helpCmd ["count-pattern"]+countPatCmd args = run (CountPat (unwords args))++saveCmd :: [String] -> Repl ()+saveCmd [] = helpCmd ["save"]+saveCmd args = run (Save (unwords args))++loadCmd :: [String] -> Repl ()+loadCmd [] = helpCmd ["load"]+loadCmd args = run (Load (unwords args))++importCmd :: Distribution -> String -> [String] -> Repl ()+importCmd dist varnames (fname:params:_) = run (Import fname dist varnames (Prelude.read params))+importCmd dist varnames _ = helpCmd ["import"]+++commands = ["help", "top", "report", "optimize", "subtrees", "insert", "count-pattern", "distribution", "pareto", "save", "load", "import"]++topHlp = "top N [FILTER...] [CRITERIA] [[not] matching [root] PATTERN] \n \+ \ \n \+ \ FILTER: with [size|cost|parameters] [<|<=|=|>|>=] N \n \+ \ CRITERIA: [by fitness | by dl] \n \+ \ \n \+ \ where \"dl\" is the description length, \"cost\" is the default cost function \n \+ \ and \"parameters\" refer to the number of parameters. The cost function \n \+ \ assigns a cost of 1 to terminals, 2 to binary operators and 3 to \n \+ \ nonlinear functions. \n \+ \ \n \+ \ Example: \n \+ \ \n \+ \ top 10 with size <= 10 with parameters > 2 by fitness matching v0 * x0 + t0 \n \+ \ \n \+ \ This will return the 10 best expressions by fitness with size less than \n \+ \ or equal to 10 and more than 2 parameters containing any sub-expression \n \+ \ in the format f(x) * x0 + t0. \n \+ \ To create a pattern for matching you can use x0 .. xn to represent a variable \n \+ \ t0 .. tn to represent a numerical parameter, and v0 .. vn to represent wildcards. \n \+ \ Notice that v0 * x0 + v0 will pattern expressions such as (sin(t0) + x0) * x0 + (sin(t0) + x0) \n \+ \ but not (sin(t0) + x0) * x0 + t0, since both occurrences of v0 will match the same expression. \n \+ \ (see `help count-pattern` for more details) \+ \ The keyword \"root\" will matches only expressions starting with this pattern."++distHlp = "distribution [FILTER] [LIMIT] \n\n \+ \ FILTER: with size [<|<=|=|>|>=] N \n \+ \ LIMIT: limited at N [asc|dsc] \n\n \+ \ Shows the distribution of all the patterns in the set of evaluated expressions.\n \+ \ The list can be filtered by the size of the pattern and limited by the top most frequent (dsc) \n \+ \ or least frequent (asc) patterns. \n\n \+ \ See `help count-pattern` for details on the syntax of pattern."++countHlp = "count-pattern PAT \n\n \+ \ Count the number of occurrence of the pattern PAT in the e-graph. \n\n \+ \ A pattern follows the same syntax of an expression: \n\n\+ \ EXPR := FUN(EXPR) | EXPR OP EXPR | TERM \n\+ \ FUN := abs | sin | cos | tan | sinh | cosh | tanh | asin | acos | atan | asinh | acosh | atanh | sqrt | sqrtabs | cbrt | square | log | logabs | exp | recip | cube \n\+ \ OP := + | - | * | / | aq | ^ | |^| \n\+ \ TERM := xN | tN | vN \n\n\+ \ where: \n \+ \ - aq is the analytical quotient (x aq y = x / sqrt(1 + y^2)) \n \+ \ - x |^| y = abs(x) ^ y \n \+ \ - xN is the N-th input variable \n \+ \ - tN is the N-th numerical parameter \n \+ \ - vN is the N-th pattern variable (see below) \n\n \+ \ The pattern variable works as a wildcard matching any expression. \n \+ \ If we use the same pattern variable multiple times in the expression, \n \+ \ the pattern must be the same in every occurrence. \n\n \+ \ Examples: \n\n \+ \ v0 + x0 will match anything added to x0\n \+ \ v0 + v1 * x0 will match anything added to any expression multiplied by x0. \+ \ For example: t0 ^ 2 + exp(t1 + x1) * x0. \n \+ \ v0 + v0 * x0 will match any expression added with this same expression multiplied by x0. \+ \ For example: t0 ^ 2 + (t0 ^ 2) * x0."++hlpMap = Map.fromList $ Prelude.zip commands+ [ "help <cmd>: shows a brief explanation for the command."+ , topHlp+ , "report N: displays a detailed report for the expression with id N."+ , "optimize N: (re)optimize expression with id N."+ , "subtrees N: shows the subtrees for the tree rotted with id N."+ , "insert EXPR: inserts a new expression EXPR and evaluates."+ , countHlp+ , distHlp+ , "pareto [by fitness| by dl]: shows the pareto front where the first objective is the criteria (default: fitness) and the second objective is model size."+ , "save FILE: save current e-graph to a file named FILE."+ , "load FILE: load current e-graph from a file named FILE."+ ]++-- Evaluation+cmd :: Map String ([String] -> Repl ()) -> String -> Repl ()+cmd cmdMap input = do let (cmd':args) = words input+ case cmdMap Map.!? cmd' of+ Nothing -> io.putStrLn $ "Command not found!!!"+ Just f -> f args++helpCmd :: [String] -> Repl ()+helpCmd [] = io . putStrLn $ "Usage: help <command-name>."+helpCmd (x:_) = case hlpMap Map.!? x of+ Nothing -> io.putStrLn $ "Command not found!!!"+ Just hlp -> io.putStrLn $ hlp+-- Completion+comp :: (Monad m, MonadState EGraph m) => WordCompleter m+comp n = pure $ filter (isPrefixOf n) commands++ini :: Repl ()+ini = do (io . putStrLn) "Welcome to r🥚ression.\nPress Ctrl-D to exit.\nPress <TAB> to see the commands."+ pure ()++final :: Repl ExitDecision+final = do io.print $ "good-bye!"+ return Exit+-- TODO: DL is sorted by min not max as the fitness+main :: IO ()+main = do+ args <- execParser opts+ g <- getStdGen+ (dataTrain, varnames) <- loadTrainingOnly (_dataset args) True+ (dataTest, _) <- if null (_testData args)+ then pure (dataTrain, "")+ else loadTrainingOnly (_testData args) True+ eg <- if (not.null) (_loadFrom args)+ then decode <$> BS.readFile (_loadFrom args)+ else if (not. null) (_parseCSV args)+ then parseCSV (_distribution args) (_parseCSV args) varnames (_parseParams args)+ else pure emptyGraph+ let --alg = evalStateT (repl dataTrain dataVal dataTest args) emptyGraph+ dist = _distribution args+ funs = [ helpCmd+ , topCmd+ , reportCmd dist dataTrain dataTest+ , optimizeCmd dist dataTrain dataTest+ , subtreesCmd+ , insertCmd dist dataTrain dataTest+ , countPatCmd+ , distCmd+ , paretoCmd+ , saveCmd+ , loadCmd+ , importCmd dist varnames+ ]+ cmdMap = Map.fromList $ Prelude.zip commands funs++ repl = evalRepl (const $ pure ">>> ") (cmd cmdMap) [] Nothing Nothing (Word comp) ini final+ crDB = if _calcDL args then (createDB >> fillDL dist dataTrain >> rebuildAllRanges) else (createDB >> rebuildAllRanges)+ if (not.null) (_convertFromTo args)+ then convert (_convertFromTo args) (_to args) varnames+ else do when (_calcDL args) $ putStrLn "Calculating DL..."+ eg' <- execStateT crDB eg+ evalStateT repl eg'+ where+ opts = Opt.info (opt <**> helper)+ ( fullDesc <> progDesc "Exploration and query system for a database of regression models using e-graphs."+ <> header "r🥚ression - Nonlinear regression models exploration and query system with e-graphs (egg)." )+
+ apps/rEGGression/Random.hs view
@@ -0,0 +1,33 @@+module Random where ++import System.Random +import Control.Monad.State.Strict+import Control.Monad+import Data.SRTree +import Data.SRTree.Eval+import Data.Massiv.Array as MA++type Rng a = IO a++toss :: Rng Bool+toss = randomIO+{-# INLINE toss #-}++tossBiased :: Double -> Rng Bool+tossBiased p = do r <- randomIO+ pure (r < p)++randomVal :: Rng Double+randomVal = randomIO++randomRange :: (Ord val, Random val) => (val, val) -> Rng val+randomRange rng = (randomRIO rng)+{-# INLINE randomRange #-}++randomFrom :: [a] -> Rng a+randomFrom funs = do n <- randomRange (0, length funs - 1)+ pure $ funs !! n+{-# INLINE randomFrom #-}++randomVec :: Int -> Rng PVector+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))
+ apps/rEGGression/Util.hs view
@@ -0,0 +1,272 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info+import Algorithm.EqSat ( recalculateBest )++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Random+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection+import Data.SRTree.Print+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy, intercalate )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries+import Algorithm.EqSat.DB+import Data.List (nub)+import System.Console.Repline hiding (Repl)+import Text.Printf+import Text.Layout.Table hiding (top)+import Text.Layout.Table.Cell.Formatted+import Text.Layout.Table.Cell+import System.Console.ANSI.Codes+++type RndEGraph = StateT EGraph IO+type DataSet = (SRMatrix, PVector, Maybe PVector)+data Info = Info {_training :: DataSet, _test :: DataSet, _dist :: Distribution}+type Repl = HaskelineT (StateT EGraph IO)++csvHeader :: String+csvHeader = "id,Expression,theta,size,MSE_train,MSE_val,MSE_test,R2_train,R2_val,R2_test,nll_train,nll_val,nll_test,mdl_train,mdl_val,mdl_test"++io :: IO a -> Repl a+io = lift . lift+{-# INLINE io #-}+egraph :: RndEGraph a -> Repl a+egraph = lift++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++fitnessFun :: Int -> Distribution -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) _tree thetaOrig =+ if isNaN tr+ then (-(1/0), theta) -- infinity+ else (tr, theta)+ where+ tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+ evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+ tr = evalF x y mYErr++{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Distribution -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nIter distribution dataTrain _tree = do+ let tree = relabelParams _tree+ nParams = countParams tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0++ thetaOrigs <- lift (randomVec nParams)+ --lift $ print thetaOrigs+ pure (fitnessFun nIter distribution dataTrain tree thetaOrigs)+{-# INLINE fitnessFunRep #-}+++-- helper query functions+-- TODO: move to egraph lib+getFitness :: EClassId -> RndEGraph (Maybe Double)+getFitness c = gets (_fitness . _info . (IM.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: EClassId -> RndEGraph (Maybe PVector)+getTheta c = gets (_theta . _info . (IM.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: EClassId -> RndEGraph Int+getSize c = gets (_size . _info . (IM.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}+getDL :: EClassId -> RndEGraph (Maybe Double)+getDL c = gets (_dl . _info . (IM.! c) . _eClass)+{-# INLINE getDL #-}++getBestFitness :: RndEGraph (Maybe Double)+getBestFitness = do+ bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+ gets (_fitness . _info . (IM.! bec) . _eClass)++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> ((a, b1, c2), c3)+getTrain ((a, b, _, _), (c, _), varnames, _) = ((a,b,c), varnames)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b++printExpr :: DataSet -> DataSet -> Distribution -> EClassId -> RndEGraph ()+printExpr dataTrain dataTest distribution ec = do+ theta <- fromJust <$> getTheta ec++ bestExpr <- getBestExpr ec+ let (x, y, mYErr) = dataTrain+ (x_te, y_te, mYErr_te) = dataTest+ best' = relabelParams bestExpr+ expr = paramsToConst (MA.toList theta) best'+ mse_train = mse x y best' theta+ mse_te = mse x_te y_te best' theta+ r2_train = r2 x y best' theta+ r2_te = r2 x_te y_te best' theta+ nll_train = nll distribution mYErr x y best' theta+ nll_te = nll distribution mYErr_te x_te y_te best' theta+ mdl_train = mdl distribution mYErr x y theta best'+ mdl_te = mdl distribution mYErr_te x_te y_te theta best'+ thetaStr = intercalate ", " $ Prelude.map show (MA.toList theta)+ lift . putStr $ "Evaluation metrics for expression (" <> (show ec) <> "): "+ lift . putStr $ setSGRCode [SetConsoleIntensity BoldIntensity]+ lift . putStrLn $ showExpr best'+ lift . putStr $ setSGRCode [Reset]+ lift . putStrLn $ "# of nodes\t" <> show (countNodes $ convertProtectedOps expr)+ lift . putStrLn $ "params:\t[" <> thetaStr <> "]"++ let rows = [ rowG ["MSE", printf "%.4f" mse_train, printf "%.4f" mse_te]+ , rowG ["R^2", printf "%.4f" r2_train, printf "%.4f" r2_te]+ , rowG ["nll", printf "%.4f" nll_train, printf "%.4f" nll_te]+ , rowG ["DL", printf "%.4f" mdl_train, printf "%.4f" mdl_te]+ ]+ columnsReport = [def, numCol, numCol]+ headerReport = titlesH $ Prelude.map bold ["Metric", "Training", "Test"]+ lift . putStrLn $ tableString (columnHeaderTableS columnsReport unicodeS headerReport rows)++printsimpleExpr eid = do t <- egraph $ relabelParams <$> getBestExpr eid+ fit <- egraph $ getFitness eid+ sz <- egraph $ getSize eid+ p <- egraph $ getTheta eid+ dl <- egraph $ getDL eid+ let fit' = case fit of+ Nothing -> "--"+ Just f -> printf "%.4f" f+ p' = case p of+ Nothing -> "--"+ Just ps -> "[" <> intercalate ", " (Prelude.map (printf "%.4f") (MA.toList ps)) <> "]"+ dl' = case dl of+ Nothing -> "--"+ Just d -> printf "%.4f" d++ pure $ colsAllG center [[show eid], justifyText 50 $ showExpr t, [fit'], justifyText 50 p', [show sz], [dl']]++printCounts (pat, cnt) = do+ let spat = showPat pat+ pure $ colsAllG center [justifyText 50 spat, [show cnt]]+ where+ showPat (Fixed (Var ix)) = 'x' : show ix+ showPat (Fixed (Param ix)) = 't' : show ix+ showPat (Fixed (Const x)) = show x+ showPat (Fixed (Bin op l r)) = concat ["(", showPat l, " ", showOp op, " ", showPat r, ")"]+ showPat (Fixed (Uni f t)) = concat [show f, "(", showPat t, ")"]+ showPat (VarPat ix) = 'v' : show (fromEnum ix-65)++printSimpleMultiExprs eids = do rows <- forM (nub eids) printsimpleExpr+ io.putStrLn $ tableString (columnHeaderTableS columns unicodeS headerSimple rows)++printMultiCounts cnts = do rows <- forM cnts printCounts+ io.putStrLn $ tableString (columnHeaderTableS [fixedLeftCol 50, numCol] unicodeS headerCount rows)++bold s = formatted (setSGRCode [SetConsoleIntensity BoldIntensity]) (plain s) (setSGRCode [Reset])++headerSimple :: HeaderSpec LineStyle (Formatted String)+headerSimple = titlesH $ Prelude.map (bold) ["Id", "Expression", "Fitness", "Parameters", "Size", "DL"]+columns = [numCol, fixedLeftCol 50, numCol, fixedLeftCol 50, numCol, numCol]+headerCount :: HeaderSpec LineStyle (Formatted String)+headerCount = titlesH $ Prelude.map bold ["Pattern", "Count"]++-- 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 t+ insertFitness ec f p+ pure True+ Just _ -> pure False++getParetoEcsUpTo b n maxSize = concat <$> forM [1..maxSize] (\i -> getTopEClassWithSize b i n)++getBestExprWithSize n =+ do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+ if (not (null ec))+ then do+ bestFit <- getFitness $ head ec+ bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+ (:[]) . (,bestP) . (,bestFit) . (,ec) <$> getBestExpr (head ec)+ else pure []++printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest printExprFun = do+ bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+ printExprFun 0 bec++paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront maxSize printExprFun = go 1 0 (-(1.0/0.0))+ where+ go :: Int -> Int -> Double -> RndEGraph ()+ go n ix f+ | n > maxSize = pure ()+ | otherwise = do+ ecList <- getBestExprWithSize n+ if not (null ecList)+ then do let (((_, ec), mf), _) = head ecList+ improved = fromJust mf > f+ ec' <- traverse canonical ec+ when improved $ printExprFun ix (head ec')+ go (n+1) (ix + if improved then 1 else 0) (max f (fromJust mf))+ else go (n+1) ix f++evaluateUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ forM_ ec $ \c -> do+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p++fillDL dist (x, y, mYErr) = do+ ecs <- getAllEvaluatedEClasses+ forM_ ecs $ \ec -> do+ theta <- fromJust <$> getTheta ec+ bestExpr <- relabelParams <$> getBestExpr ec+ if MA.size theta /= countParams bestExpr+ then (lift . putStrLn) $ "Wrong number of parameters in " <> showExpr bestExpr <> ": " <> show theta <> " " <> show ec+ else do let mdl_train = mdl dist mYErr x y theta bestExpr+ insertDL ec mdl_train
apps/srtools/Args.hs view
@@ -19,11 +19,11 @@ , hasHeader :: Bool , simpl :: Bool , dist :: Distribution- , msErr :: Maybe Double , restart :: Bool , rseed :: Int , toScreen :: Bool , useProfile :: Bool+ , simple :: Bool , alpha :: Double , ptype :: PType } deriving Show@@ -61,11 +61,14 @@ \ It will auto-detect and handle gzipped file based on gz extension. \ \ It will also auto-detect the delimiter.\n\ \ The filename can include extra information: \- \ filename.csv:start:end:target:vars where start and end \+ \ filename.csv:start:end:target:cols:yerr:xerr where start and end \ \ corresponds to the range of rows that should be used for fitting,\ \ target is the column index (or name) of the target variable and cols\- \ is a comma separated list of column indeces or names of the variables\- \ in the same order as used by the symbolic model." )+ \ is a comma separated list of column indices or names of the variables\+ \ in the same order as used by the symbolic model.\+ \ The yerr field corresponds to the column with the error of the target,\+ \ while xerr a comma separated indices of the columns with the error of the\+ \ variables. If nothing passed, it will ignore measurement errors." ) <*> strOption ( long "test" <> metavar "TEST"@@ -96,13 +99,6 @@ \ the avaliable distributions.\ \ The default is Gaussian." )- <*> option s2Reader- ( long "sErr"- <> metavar "Serr"- <> showDefault- <> value Nothing- <> help "Estimated standard error of the data.\- \ If not passed, it uses the model MSE.") <*> switch ( long "restart" <> help "If set, it samples the initial values of\@@ -123,6 +119,9 @@ <*> switch ( long "profile" <> help "If set, it will use profile likelihood to calculate the CIs." )+ <*> switch+ ( long "simple"+ <> help "If set, calculates only SSE.") <*> option auto ( long "alpha" <> metavar "ALPHA"@@ -166,9 +165,9 @@ where errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp -s2Reader :: ReadM (Maybe Double)-s2Reader =- str >>= \s -> mkReader ("wrong format " <> s) Just s+--s2Reader :: ReadM (Maybe Double)+--s2Reader =+-- str >>= \s -> mkReader ("wrong format " <> s) Just s distRead :: ReadM Distribution distRead =
apps/srtools/IO.hs view
@@ -9,13 +9,12 @@ import System.Random ( StdGen ) import Data.SRTree ( SRTree (..), Fix (..), var, floatConstsToParam, relabelVars )-import Algorithm.SRTree.Opt ( estimateSErr ) import Algorithm.SRTree.Likelihoods ( Distribution (..) ) import Algorithm.SRTree.ConfidenceIntervals ( printCI, BasicStats(_stdErr, _corr), CI ) import qualified Data.SRTree.Print as P import Data.SRTree.Eval ( compMode ) -import Args ( Args(outfile, alpha,msErr,dist,niter) )+import Args ( Args(outfile, alpha,dist,niter) ) import Report import Data.SRTree.Recursion ( cata ) @@ -26,6 +25,10 @@ csvHeader = intercalate "," (basicFields <> optFields <> modelFields) {-# inline csvHeader #-} +csvHeaderSimple :: String+csvHeaderSimple = intercalate "," (basicFields <> optFields)+{-# inline csvHeaderSimple #-}+ -- Open file if filename is not empty openWriteWithDefault :: Handle -> String -> IO Handle openWriteWithDefault dflt "" = pure dflt@@ -42,20 +45,35 @@ processTree args seed dset t ix = (basic, sseOrig, sseOpt, info, cis) where (tree, theta0) = floatConstsToParam t- mSErr' = case dist args of- Gaussian -> estimateSErr Gaussian (msErr args) (_xTr dset) (_yTr dset) (A.fromList compMode theta0) tree (niter args)- _ -> Nothing- args' = args{ msErr = mSErr' }- basic = getBasicStats args' seed dset tree theta0 ix++ basic = getBasicStats args seed dset tree theta0 ix treeVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> _expr basic (_, Nothing) -> _expr basic- (Just xV, Just yV) -> _expr $ getBasicStats args' seed dset{_xTr = xV, _yTr = yV} tree theta0 ix+ (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix sseOrig = getSSE dset t sseOpt = getSSE dset (_expr basic)- info = getInfo args' dset (_expr basic) treeVal- cis = getCI args' dset basic (alpha args')+ info = getInfo args dset (_expr basic) treeVal+ cis = getCI args dset basic (alpha args) +processTreeSimple :: Args -- command line arguments+ -> StdGen -- random number generator+ -> Datasets -- datasets+ -> Fix SRTree -- expression in tree format+ -> Int -- index of the parsed expression+ -> (BasicInfo, SSE, SSE)+processTreeSimple args seed dset t ix = (basic, sseOrig, sseOpt)+ where+ (tree, theta0) = floatConstsToParam t++ basic = getBasicStats args seed dset tree theta0 ix+ treeVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> _expr basic+ (_, Nothing) -> _expr basic+ (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix+ sseOrig = getSSE dset t+ sseOpt = getSSE dset (_expr basic)+ -- print the results to a csv format (except CI) printResults :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO () printResults args seed dset varnames exprs = do@@ -69,6 +87,18 @@ in hPutStrLn hStat (toCsv treeData varnames) unless (null (outfile args)) (hClose hStat) +printResultsSimple :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO ()+printResultsSimple args seed dset varnames exprs = do+ hStat <- openWriteWithDefault stdout (outfile args)+ hPutStrLn hStat csvHeaderSimple+ forM_ (zip [0..] exprs)+ \(ix, tree) ->+ case tree of+ Left err -> hPutStrLn stderr ("invalid expression: " <> err)+ Right t -> let treeData = processTreeSimple args seed dset t ix+ in hPutStrLn hStat (toCsvSimple treeData varnames)+ unless (null (outfile args)) (hClose hStat)+ -- change the stats into a string toCsv :: (BasicInfo, SSE, SSE, Info, e) -> [String] -> String toCsv (basic, sseOrig, sseOpt, info, _) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt <> sInfo)@@ -76,11 +106,24 @@ sBasic = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic) , show (_nNodes basic), show (_nParams basic) , intercalate ";" (map show (_params basic))+ , show (_nEvals basic) ] sSSEOrig = map (showF sseOrig) [_sseTr, _sseVal, _sseTe] sSSEOpt = map (showF sseOpt) [_sseTr, _sseVal, _sseTe] sInfo = map (showF info) [_bic, _bicVal, _aic, _aicVal, _evidence, _evidenceVal, _mdl, _mdlFreq, _mdlLatt, _mdlVal, _mdlFreqVal, _mdlLattVal, _nllTr, _nllVal, _nllTe, _cc, _cp] <> [intercalate ";" (map show (_fisher info))]+ showF p f = show (f p)++toCsvSimple :: (BasicInfo, SSE, SSE) -> [String] -> String+toCsvSimple (basic, sseOrig, sseOpt) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt)+ where+ sBasic = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic)+ , show (_nNodes basic), show (_nParams basic)+ , intercalate ";" (map show (_params basic))+ , show (_nEvals basic)+ ]+ sSSEOrig = map (showF sseOrig) [_sseTr, _sseVal, _sseTe]+ sSSEOpt = map (showF sseOpt) [_sseTr, _sseVal, _sseTe] showF p f = show (f p) -- get trees of transformed features
apps/srtools/Main.hs view
@@ -21,7 +21,9 @@ withInput (infile args) (from args) varnames False (simpl args) >>= if toScreen args then printResultsScreen args seed dset varnames' tgname -- full report on screne- else printResults args seed dset varnames' -- csv file+ else if simple args+ then printResultsSimple args seed dset varnames' -- csv file+ else printResults args seed dset varnames' -- csv file where opts = info (opt <**> helper) ( fullDesc <> progDesc "Optimize the parameters of\
apps/srtools/Report.hs view
@@ -7,16 +7,15 @@ import Statistics.Distribution.FDistribution ( fDistribution ) import Statistics.Distribution.ChiSquared ( chiSquared ) import Statistics.Distribution ( quantile )-import System.Random ( StdGen, split )-import Data.Random.Normal ( normals )+import System.Random ( StdGen, split, randomRs ) import Data.SRTree ( SRTree, Fix (..), floatConstsToParam, paramsToConst, countNodes ) import Data.SRTree.Eval-import Algorithm.SRTree.AD ( reverseModeUnique, forwardModeUniqueJac )+import Algorithm.SRTree.AD ( forwardModeUniqueJac ) import Algorithm.SRTree.Likelihoods import Algorithm.SRTree.ModelSelection ( aic, bic, evidence, logFunctional, logParameters, mdl, mdlFreq, mdlLatt ) import Algorithm.SRTree.ConfidenceIntervals-import Algorithm.SRTree.Opt (minimizeNLLWithFixedParam, minimizeNLL, minimizeNLLNonUnique)+import Algorithm.SRTree.Opt (minimizeNLLWithFixedParam, minimizeNLL) import Data.SRTree.Datasets ( loadDataset ) import Data.SRTree.Print ( showExpr ) import Debug.Trace ( trace, traceShow )@@ -24,12 +23,15 @@ import Args -- store the datasets split into training, validation and test-data Datasets = DS { _xTr :: SRMatrix- , _yTr :: PVector- , _xVal :: Maybe SRMatrix- , _yVal :: Maybe PVector- , _xTe :: Maybe SRMatrix- , _yTe :: Maybe PVector+data Datasets = DS { _xTr :: SRMatrix+ , _yTr :: PVector+ , _xVal :: Maybe SRMatrix+ , _yVal :: Maybe PVector+ , _xTe :: Maybe SRMatrix+ , _yTe :: Maybe PVector+ , _yErrTr :: Maybe PVector+ , _yErrVal :: Maybe PVector+ , _yErrTe :: Maybe PVector } -- basic fields name@@ -40,6 +42,7 @@ , "Number_of_nodes" , "Number_of_parameters" , "Parameters"+ , "Number_of_evaluations" ] -- basic information about the tree@@ -49,6 +52,7 @@ , _nNodes :: Int , _nParams :: Int , _params :: [Double]+ , _nEvals :: Int } -- optimization fields@@ -113,27 +117,27 @@ -- load the datasets getDataset :: Args -> IO (Datasets, String, String) getDataset args = do- ((xTr, yTr, xVal, yVal), varnames, tgname) <- loadDataset (dataset args) (hasHeader args)+ ((xTr, yTr, xVal, yVal), (yErrTr, yErrVal), varnames, tgname) <- loadDataset (dataset args) (hasHeader args) let (A.Sz m) = A.size yVal let (mXVal, mYVal) = if m == 0 then (Nothing, Nothing) else (Just xVal, Just yVal)- (mXTe, mYTe) <- if null (test args)- then pure (Nothing, Nothing)- else do ((xTe, yTe, _, _), _, _) <- loadDataset (test args) (hasHeader args)- pure (Just xTe, Just yTe)- pure (DS xTr yTr mXVal mYVal mXTe mYTe, varnames, tgname)+ (mXTe, mYTe, mYErrTe) <- if null (test args)+ then pure (Nothing, Nothing, Nothing)+ else do ((xTe, yTe, _, _), (yErrTe, _), _, _) <- loadDataset (test args) (hasHeader args)+ pure (Just xTe, Just yTe, yErrTe)+ pure (DS xTr yTr mXVal mYVal mXTe mYTe yErrTr yErrVal mYErrTe, varnames, tgname) getBasicStats :: Args -> StdGen -> Datasets -> Fix SRTree -> [Double] -> Int -> BasicInfo getBasicStats args seed dset tree theta0 ix | anyNaN = getBasicStats args (snd $ split seed) dset tree theta0 ix- | otherwise = Basic ix (infile args) tOpt nNodes nParams params+ | otherwise = Basic ix (infile args) tOpt nNodes nParams params nEvs where -- (tree', theta0) = floatConstsToParam tree thetas = if restart args- then A.fromList compMode $ take nParams (normals seed)+ then A.fromList compMode $ take nParams (randomRs (-1.0, 1.0) seed) else A.fromList compMode theta0- t = fst $ minimizeNLL (dist args) (msErr args) (niter args) (_xTr dset) (_yTr dset) tree thetas+ (t,_,nEvs) = minimizeNLL (dist args) (_yErrTr dset) (niter args) (_xTr dset) (_yTr dset) tree thetas tOpt = paramsToConst (A.toList t) tree nNodes = countNodes tOpt :: Int nParams = length theta0@@ -156,15 +160,15 @@ getInfo :: Args -> Datasets -> Fix SRTree -> Fix SRTree -> Info getInfo args dset tree treeVal =- Info { _bic = bic dist' msErr' xTr yTr thetaOpt' tOpt+ Info { _bic = bic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt , _bicVal = bicVal- , _aic = aic dist' msErr' xTr yTr thetaOpt' tOpt+ , _aic = aic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt , _aicVal = aicVal- , _evidence = evidence dist' msErr' xTr yTr thetaOpt' tOpt+ , _evidence = evidence dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt , _evidenceVal = evidenceVal- , _mdl = mdl dist' msErr' xTr yTr thetaOpt' tOpt- , _mdlFreq = mdlFreq dist' msErr' xTr yTr thetaOpt' tOpt- , _mdlLatt = mdlLatt dist' msErr' xTr yTr thetaOpt' tOpt+ , _mdl = mdl dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _mdlFreq = mdlFreq dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _mdlLatt = mdlLatt dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt , _mdlVal = mdlVal , _mdlFreqVal = mdlFreqVal , _mdlLattVal = mdlLattVal@@ -172,8 +176,8 @@ , _nllVal = nllVal , _nllTe = nllTe , _cc = logFunctional tOpt- , _cp = logParameters dist' msErr' xTr yTr thetaOpt' tOpt- , _fisher = A.toList $ fisherNLL dist' (msErr args) xTr yTr tOpt thetaOpt'+ , _cp = logParameters dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _fisher = A.toList $ fisherNLL dist' (_yErrTr dset) xTr yTr tOpt thetaOpt' } where (xTr, yTr) = (_xTr dset, _yTr dset)@@ -188,40 +192,40 @@ thetaOptVal' = A.fromList compMode thetaOptVal dist' = dist args- msErr' = msErr args- nllTr = nll dist' msErr' (_xTr dset) (_yTr dset) tOpt (A.fromList compMode thetaOpt)++ nllTr = nll dist' (_yErrTr dset) (_xTr dset) (_yTr dset) tOpt (A.fromList compMode thetaOpt) bicVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> bic dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> bic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal aicVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> aic dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> aic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal evidenceVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> evidence dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> evidence dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal nllVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> nll dist' msErr' xVal yVal tOptVal (A.fromList compMode thetaOptVal)+ _ -> nll dist' (_yErrVal dset) xVal yVal tOptVal (A.fromList compMode thetaOptVal) mdlVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> mdl dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> mdl dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal mdlFreqVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> mdlFreq dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> mdlFreq dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal mdlLattVal = case (_xVal dset, _yVal dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- _ -> mdlLatt dist' msErr' xVal yVal thetaOptVal' tOptVal+ _ -> mdlLatt dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal nllTe = case (_xTe dset, _yTe dset) of (Nothing, _) -> 0.0 (_, Nothing) -> 0.0- (Just xTe, Just yTe) -> nll dist' msErr' xTe yTe tOpt (A.fromList compMode thetaOpt)+ (Just xTe, Just yTe) -> nll dist' (_yErrTe dset) xTe yTe tOpt (A.fromList compMode thetaOpt) getCI :: Args -> Datasets -> BasicInfo -> Double -> (BasicStats, [CI], [CI], [CI], [CI]) getCI args dset basic alpha' = (stats', cis, pis_tr, pis_val, pis_te)@@ -233,23 +237,22 @@ tau_max' = sqrt $ quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - alpha') (xTr, yTr) = (_xTr dset, _yTr dset) dist' = dist args- msErr' = msErr args- stats' = getStatsFromModel dist' msErr' xTr yTr tree (A.fromList compMode theta)- profiles = getAllProfiles (ptype args) dist' msErr' xTr yTr tree (A.fromList compMode theta) (_stdErr stats') estCIs alpha'+ stats' = getStatsFromModel dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta)+ profiles = getAllProfiles (ptype args) dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta) (_stdErr stats') estCIs alpha' method = if useProfile args then Profile stats' profiles else Laplace stats' predFun = A.computeAs A.S . predict dist' tree (A.fromList compMode theta) prof estPi th t =- let (thOpt, _) = minimizeNLLNonUnique dist' (Just 1) 100 xTr yTr t th+ let (thOpt, _, _) = minimizeNLL dist' (_yErrTr dset) 100 xTr yTr t th ssr = sse xTr yTr t thOpt est = sqrt $ ssr / fromIntegral (n - _nParams basic) stdErr = _stdErr stats' A.! 0 fun = case ptype args of- Bates -> getProfile dist' (Just est) xTr yTr t thOpt stdErr tau_max 0- ODE -> getProfileODE dist' (Just est) xTr yTr t thOpt stdErr estPi tau_max 0- Constrained -> getProfileCnstr dist' (Just est) xTr yTr t thOpt stdErr tau_max' 0+ Bates -> getProfile dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max 0+ ODE -> getProfileODE dist' (_yErrTr dset) xTr yTr t thOpt stdErr estPi tau_max 0+ Constrained -> getProfileCnstr dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max' 0 in case fun of Left th' -> trace "found better optima" $ prof estPi th' t Right p -> (_tau2theta p, _opt p)
apps/tinygp/GP.hs view
@@ -10,19 +10,22 @@ import Data.SRTree.Eval import Data.SRTree.Recursion ( cata ) import System.Random-import Control.Monad.State+import Control.Monad.State.Strict import Control.Monad import Data.Vector qualified as V import Control.Monad (when) import Data.Massiv.Array qualified as M import Debug.Trace ( traceShow, trace )+import Util+import Data.List ( intercalate )+import qualified Data.Vector.Mutable as MV data Method = Grow | Full | BTC+type Rng a = StateT StdGen IO a -type Rng a = StateT StdGen IO a type GenUni = Fix SRTree -> Fix SRTree type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree-type FitFun = Individual -> Individual +type FitFun = Individual -> Rng Individual data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: PVector } @@ -82,7 +85,6 @@ if r then randomFrom term else genNonTerm -{-# INLINE randomTree #-} data HyperParams = HP { _minDepth :: Int @@ -104,19 +106,17 @@ if null selection then randomFromV pop else randomFromV champions-{-# INLINE tournament #-} randomIndividual :: HyperParams -> FitFun -> Bool -> Rng Individual randomIndividual hyperparams fitFun grow = do t <- randomTree hyperparams grow let p = countParams t theta' <- replicateM p (randomRange (-1,1))- let ind = fitFun $ Individual t 0.0 (M.fromList compMode theta' :: PVector)- pure ind+ fitFun $ Individual t 0.0 (M.fromList compMode theta' :: PVector)+ --pure ind --if isInfinite (_fit ind) -- then randomIndividual hyperparams fitFun grow -- else pure ind-{-# INLINE randomIndividual #-} initialPop :: HyperParams -> FitFun -> Rng (V.Vector Individual) initialPop hyperparams fitFun = do @@ -126,77 +126,101 @@ g = V.fromList . take m $ cycle [True, False] mapM (randomIndividual hyperparams{ _maxDepth = md} fitFun) g pure (V.concat pop)-{-# INLINE initialPop #-} -fitness :: SRMatrix -> PVector -> Individual -> Individual-fitness x y ind =- let - tree = relabelParams $ _tree ind- thetaOrig = _params ind- (theta, fit) = minimizeNLL Gaussian (Just 1) 10 x y tree thetaOrig- --theta = _params ind- fit' = negate $ mse x y tree thetaOrig -- nll Gaussian (Just 1) x y (relabelParams $ _tree ind) (_params ind)- -- (fit, g) = gradNLL Gaussian Nothing x y (_tree ind) (_params ind)- in if M.isNull (_params ind)- then ind{_fit=fit'}- else ind{_fit = negate (mse x y tree theta), _params = theta}- --in ind{_fit = fit, _params = theta}-{-# INLINE fitness #-}+fitness :: SRMatrix -> PVector -> Individual -> Rng Individual+fitness x y ind = do+ let tree = relabelParams $ _tree ind+ p = countParams tree+ theta1' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))+ theta2' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))+ let (theta1, f1, _) = minimizeNLL MSE Nothing 50 x y tree theta1'+ (theta2, f2, _) = minimizeNLL MSE Nothing 50 x y tree theta2'+ fit1 = negate f1+ fit2 = negate f2+ thetaOpt = if fit1 > fit2 then theta1 else theta2+ fitOpt = max fit1 fit2+ pure ind{_fit = fitOpt, _params = thetaOpt} -isAbs (Fix (Uni Abs _)) = True -isAbs _ = False -{-# INLINE isAbs #-} -isInv (Fix (Bin Div (Fix (Const 1.0)) _)) = True -isInv _ = False -{-# INLINE isInv #-}--mutate :: HyperParams -> Individual -> Rng Individual+mutate :: HyperParams -> Individual -> Rng (Maybe Individual) mutate hp ind = do let sz = countNodes' (_tree ind)- (t, b) <- go sz (_pm hp) (_tree ind)- if b - then pure $ Individual t 0.0 M.empty- else pure ind+ p <- state $ randomR (0, sz-1)+ b <- state random+ t <- go p (_maxSize hp) (_tree ind)+ --(t, b) <- go sz (_pm hp) (_tree ind)+ if b <= _pm hp && countNodes t <= _maxSize hp+ then pure . Just $ Individual t 0.0 M.empty+ else pure Nothing where- go s p t- | isAbs t = do let [x] = getChildren t- (t', b) <- go s p x- pure (Fix $ replaceChildren [t'] $ unfix t, b)- | otherwise = do- v <- state random - if v < p - then do let sz2 = countNodes' t - maxSz = _maxSize hp - s + sz2 - 2- (, True) <$> randomTree hp{_maxSize = maxSz} True - else case arity t of - 0 -> pure (t, False)- 1 -> do let [x] = getChildren t - (t', b) <- go s p x - pure (Fix $ replaceChildren [t'] $ unfix t, b)- 2 -> do let [l,r] = getChildren t - (l', b) <- if isInv t - then pure (l, False)- else go s p l- if b - then pure (Fix $ replaceChildren [l', r] $ unfix t, b)- else do (r', b') <- go s p r - pure (Fix $ replaceChildren [l, r'] $ unfix t, b')+ go 0 msz t = randomTree hp{_maxSize = msz-1} True+ go n msz (Fix (Uni f t)) = Fix . Uni f <$> go (n-1) (msz-1) t+ go n msz (Fix (Bin op l r)) = do+ let nl = countNodes l+ nr = countNodes r+ if nl <= n - 1+ then Fix . Bin op l <$> go (n-nl-1) (msz-nl-1) r+ else do l' <- go (n-1) (msz-nr-1) l+ pure $ Fix $ Bin op l' r -crossover :: HyperParams -> Individual -> Individual -> Rng Individual-crossover hp ind1 ind2 = pure ind1+crossover :: HyperParams -> Individual -> Individual -> Rng (Maybe Individual)+crossover hp ind1 ind2 = do+ b <- state random+ if b < (_pc hp)+ then do let n1 = countNodes $ _tree ind1+ n2 = countNodes $ _tree ind2+ p1 <- state $ randomR (0, n1-1)+ p2 <- state $ randomR (0, n2-1)+ let part1 = pickLeft p1 $ _tree ind1+ part2 = pickRight p2 $ _tree ind2+ t = part1 part2+ n = countNodes t+ if n <= _maxSize hp+ then pure . Just $ ind1{_tree = t}+ else pure Nothing+ else pure Nothing+ where+ pickRight :: Int -> Fix SRTree -> Fix SRTree+ pickRight 0 node = node+ pickRight n (Fix (Uni f t)) = pickRight (n-1) t+ pickRight n (Fix (Bin op l r)) = let nl = countNodes l+ in if nl <= n-1+ then pickRight (n-nl-1) r+ else pickRight (n-1) l+ pickLeft :: Int -> Fix SRTree -> (Fix SRTree -> Fix SRTree)+ pickLeft 0 node = \t -> t+ pickLeft n (Fix (Uni f t)) = let g = pickLeft (n-1) t in \t' -> Fix $ Uni f (g t')+ pickLeft n (Fix (Bin op l r)) = let nl = countNodes l+ in if nl <= n-1+ then let g = pickLeft (n-nl-1) r in \t -> Fix $ Bin op l (g t)+ else let g = pickLeft (n-1) l in \t -> Fix $ Bin op (g t) r + evolve :: HyperParams -> FitFun -> V.Vector Individual -> Rng Individual evolve hp fitFun pop = do parent1 <- tournament hp pop parent2 <- tournament hp pop - child <- crossover hp parent1 parent2- child' <- mutate hp child- let p = countParams (_tree child')- theta' <- M.fromList compMode <$> replicateM p (randomRange (-1,1))- pure $ fitFun child'{_params = theta'}-{-# INLINE evolve #-}+ mChild <- crossover hp parent1 parent2+ child' <- case mChild of+ Nothing -> mutate hp parent1+ Just child -> mutate hp child+ --let p = countParams (_tree child')+ --theta' <- M.fromList compMode <$> replicateM p (randomRange (-1,1))+ case child' of+ Nothing -> pure parent1+ Just c -> fitFun c +printFinal ind x y x_test y_test = do+ let tree = relabelParams $ _tree ind+ theta = _params ind+ mseTrain = mse x y tree theta+ mseTest = mse x_test y_test tree theta+ r2Train = r2 x y tree theta+ r2Test = r2 x_test y_test tree theta+ thetaStr = intercalate ";" $ map show $ M.toList theta+ putStrLn "id,Expression,theta,size,MSE_train,MSE_test,R2_train,R2_test"+ putStr $ "0," <> showExpr tree <> "," <> thetaStr <> "," <> show (countNodes tree) <> "," <> show mseTrain <> "," <> show mseTest <> "," <> show r2Train <> "," <> show r2Test+ report :: Int -> V.Vector Individual -> IO () report gen = mapM_ reportOne where reportOne ind = do putStr (show gen)@@ -208,15 +232,15 @@ print (M.toList $ _params ind) {-# INLINE report #-} -evolution :: Int -> HyperParams -> FitFun -> Rng (V.Vector Individual)+evolution :: Int -> HyperParams -> FitFun -> Rng (Individual) evolution gen hp fitFun = do pop <- initialPop hp fitFun- liftIO $ report (-1) pop- go gen pop + --liftIO $ report (-1) pop+ go gen pop where - go 0 pop = pure pop - go n pop = do + go 0 !pop = pure $ pop V.! 0+ go n !pop = do let best = V.maximumOn _fit $ V.filter (not.isNaN._fit) pop- pop' <- V.cons best <$> V.replicateM (_popSize hp - 1) (evolve hp fitFun pop)- liftIO $ report (gen-n) pop'+ pop' <- V.modify (\v -> MV.write v 0 best) <$> V.replicateM (_popSize hp) (evolve hp fitFun pop)+ --liftIO $ report (gen-n) pop' go (n-1) pop'
apps/tinygp/Main.hs view
@@ -1,23 +1,29 @@ module Main (main) where -import GP ( HyperParams(HP), fitness, evolution )-import Data.SRTree ( param, var )+import GP ( HyperParams(HP), fitness, evolution, printFinal )+import Data.SRTree import System.Random ( getStdGen )-import Control.Monad.State ( evalStateT )+import Control.Monad.State.Strict ( evalStateT ) import Data.SRTree.Datasets ( loadDataset ) import Options.Applicative import Data.Massiv.Array +import Util -- Data type to store command line arguments data Args = Args- { dataset :: String,- popSize :: Int,- gens :: Int,- pc :: Double,- pm :: Double+ { dataset :: String,+ _testData :: String,+ popSize :: Int,+ gens :: Int,+ _maxSize :: Int,+ pc :: Double,+ pm :: Double,+ _nonterminals :: String,+ _nTournament :: Int } deriving (Show) + -- parser of command line arguments opt :: Parser Args opt = Args@@ -26,6 +32,11 @@ <> short 'd' <> metavar "INPUT-FILE" <> help "CSV dataset." )+ <*> strOption+ ( long "test"+ <> value ""+ <> metavar "INPUT-FILE"+ <> help "CSV dataset." ) <*> option auto ( long "population" <> short 'p'@@ -41,6 +52,12 @@ <> value 100 <> help "Number of generations." ) <*> option auto+ ( long "max-size"+ <> metavar "SIZE"+ <> showDefault+ <> value 20+ <> help "maximum expression size." )+ <*> option auto ( long "probCx" <> metavar "PC" <> showDefault@@ -52,20 +69,32 @@ <> showDefault <> value 0.3 <> help "Mutation probability." )+ <*> strOption+ ( long "non-terminals"+ <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+ <> showDefault+ <> help "set of non-terminals to use in the search."+ )+ <*> option auto+ ( long "tournament-size"+ <> value 2+ <> showDefault+ <> help "tournament size."+ ) -nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> abs l ** r), Left (1/)]---nonterms = [Right (+), Right (-), Right (*)]+nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> Fix $ Bin PowerAbs l r), Left recip, Left log, Left exp, Left (\t -> Fix $ Uni SqrtAbs t)] main :: IO () main = do args <- execParser opts g <- getStdGen- ((x, y, _, _), _, _) <- loadDataset (dataset args) True- let hp = HP 2 4 25 (popSize args) 2 (pc args) (pm args) terms nonterms + (x, y, _) <- loadTrainingOnly (dataset args) True+ (x_test, y_test, _) <- loadTrainingOnly (_testData args) True+ let hp = HP 3 10 (_maxSize args) (popSize args) (_nTournament args) (pc args) (pm args) terms (parseNonTerms $ _nonterminals args) (Sz2 _ nFeats) = size x terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]]- pop <- evalStateT (evolution (gens args) hp (fitness x y)) g- putStrLn "Fin"+ best <- evalStateT (evolution (gens args) hp (fitness x y)) g+ printFinal best x y x_test y_test where opts = info (opt <**> helper) ( fullDesc <> progDesc "Very simple example of GP using SRTree."
+ apps/tinygp/Util.hs view
@@ -0,0 +1,62 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Algorithm.SRTree.Likelihoods+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries+++type DataSet = (SRMatrix, PVector, Maybe PVector)+++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b+++parseNonTerms = Prelude.map toNonTerm . splitOn ","+ where+ binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+ uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+ toNonTerm xs' = let xs = Prelude.map toLower xs'+ in case binTerms Map.!? xs of+ Just op -> Right $ \l r -> Fix $ Bin op l r+ Nothing -> case uniTerms Map.!? xs of+ Just f -> Left $ \t -> Fix $ Uni f t+ Nothing -> error $ "invalid non-terminal " <> show xs+
src/Algorithm/EqSat.hs view
@@ -52,12 +52,12 @@ eqSat expr rules costFun maxIt = do root <- fromTree costFun expr (end, it) <- runEqSat costFun rules maxIt- best <- getBest root+ best <- getBestExpr root --info <- gets ((IntMap.! root) . _eClass) --info2 <- gets ((IntMap.! 9) . _eClass) --traceShow (info, info2) $ if not end -- if had an early stop- then do eqSat best rules costFun it -- reapplies eqsat on the best so far + then do modify' (const emptyGraph) >> eqSat best rules costFun it -- reapplies eqsat on the best so far else pure best type CostMap = Map EClassId (Int, Fix SRTree)@@ -116,7 +116,7 @@ go it sch = do eNodes <- gets _eNodeToEClass eClasses <- gets _eClass- --createDB+ --createDB -- TODO: partial db is still incomplete --db <- gets (_patDB . _eDB) -- createDB -- creates the DB -- step 1: match the rules@@ -147,9 +147,10 @@ do db <- gets (_patDB . _eDB) -- createDB let matchSch = matchWithScheduler 10 matchAll = zipWithM matchSch [0..]- (rules, sch') = runState (matchAll rules') IntMap.empty- matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rules- mapM_ (uncurry (applyMergeOnlyMatch costFun)) $ concat matches+ (rls, sch') = runState (matchAll rules') IntMap.empty+ --matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rls+ --mapM_ (uncurry (applyMergeOnlyMatch costFun)) $ take 500 $ concat matches+ matches <- getNMatches 500 rls rebuild costFun -- recalculate heights --calculateHeights@@ -161,6 +162,18 @@ replaceEqRules (p1 :=> p2) = [p1 :=> p2] replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1] replaceEqRules (r :| cond) = map (:| cond) $ replaceEqRules r++ getNMatches n [] = pure []+ getNMatches 0 _ = pure []+ getNMatches n ([]:rss) = getNMatches n rss+ getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> match (source r)+ let (x, y) = splitAt n matches+ m = length x+ if m == n+ then pure matches+ else do matches' <- getNMatches (n - length x) (rs:rss)+ pure (matches <> matches')+ -- | matches the rules given a scheduler matchWithScheduler :: Int -> Int -> Rule -> Scheduler [Rule] -- [(Rule, (Map ClassOrVar ClassOrVar, ClassOrVar))]
src/Algorithm/EqSat/Build.hs view
@@ -35,7 +35,7 @@ import qualified Data.IntSet as IntSet import Data.Maybe import Data.Sequence (Seq(..), (><))-+import Data.List ( nub ) import Debug.Trace (trace, traceShow) -- | adds a new or existing e-node (merging if necessary)@@ -43,10 +43,20 @@ add costFun enode = do enode'' <- canonize enode -- canonize e-node constEnode <- calculateConsts enode''- let enode' = case constEnode of- ConstVal x -> Const x- ParamIx x -> Param x- _ -> enode''+ enode' <- case constEnode of+ ConstVal x -> pure $ Const x+ ParamIx x -> pure $ Param x+ _ -> case enode'' of+ Bin Sub c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)+ pure $ case constType of+ ParamIx x -> Bin Add c1 c2+ _ -> enode''+ Bin Div c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)+ pure $ case constType of+ ParamIx x -> Bin Mul c1 c2+ _ -> enode''+ _ -> pure $ enode''+ maybeEid <- gets ((Map.!? enode') . _eNodeToEClass) -- check if canonical e-node exists case maybeEid of Just eid -> pure eid@@ -86,6 +96,7 @@ . over (eDB . analysis) (const Set.empty) forM_ wl (uncurry (repair costFun)) forM_ al (uncurry (repairAnalysis costFun))+{-# INLINE rebuild #-} -- | repairs e-node by canonizing its children -- if the canonized e-node already exists in@@ -100,7 +111,7 @@ Just ecIdCanon -> do mergedId <- merge costFun ecIdCanon ecId' modify' $ over eNodeToEClass (Map.insert enode' mergedId) Nothing -> modify' $ over eNodeToEClass (Map.insert enode' ecId')-+{-# INLINE repair #-} -- | repair the analysis of the e-class -- considering the new added e-node@@ -115,8 +126,10 @@ when (_info eclass /= newData) $ do modify' $ over (eDB . analysis) (_parents eclass <>) . over eClass (IntMap.insert ecId' eclass')+ . over (eDB . refits) (Set.insert ecId') _ <- modifyEClass costFun ecId' pure ()+{-# INLINE repairAnalysis #-} -- | merge to equivalent e-classes merge :: Monad m => CostFun -> EClassId -> EClassId -> EGraphST m EClassId@@ -130,7 +143,7 @@ where mergeClasses :: Monad m => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId mergeClasses led ledC ledO sub subC subO =- do modify' $ over canonicalMap (IntMap.insert sub led) -- points sub e-class to leader to maintain consistency+ do modify' $ over canonicalMap (IntMap.insert sub led . IntMap.insert subO led) -- points sub e-class to leader to maintain consistency let -- create new e-class with same id as led newC = EClass led (_eNodes ledC `Set.union` _eNodes subC)@@ -142,10 +155,12 @@ . over (eDB . worklist) (_parents subC <>) -- insert parents of sub into worklist when (_info newC /= _info ledC) -- if there was change in data, $ modify' $ over (eDB . analysis) (_parents ledC <>) -- insert parents into analysis+ . over (eDB . refits) (Set.insert led) when (_info newC /= _info subC) $ modify' $ over (eDB . analysis) (_parents subC <>) updateDBs newC led ledC ledO sub subC subO modifyEClass costFun led+ --forM_ (_eNodes newC) $ \en -> addToDB (decodeEnode en) led pure led getLeaderSub c1 c1O c2 c2O =@@ -206,10 +221,11 @@ let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } maybeEid <- gets ((Map.!? en) . _eNodeToEClass) modify' $ over eClass (IntMap.insert ecId ec{_eNodes = Set.singleton (encodeEnode en) , _info = infoEc})+ when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId) case maybeEid of Nothing -> pure ecId Just eid' -> merge costFun eid' ecId- {-+ ParamIx x -> do let en = Param x c <- calculateCost costFun en@@ -217,11 +233,12 @@ let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } maybeEid <- gets ((Map.!? en) . _eNodeToEClass) modify' $ over eClass (IntMap.insert ecId ec{_eNodes = Set.insert (encodeEnode en) (_eNodes ec), _info = infoEc})+ when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId) -- TODO: what happen to the orphans? case maybeEid of Nothing -> pure ecId- Just eid' -> trace "merge" $ merge costFun eid' ecId- -}+ Just eid' -> merge costFun eid' ecId+ _ -> pure ecId where@@ -244,6 +261,7 @@ ecls <- gets (Map.toList . _eNodeToEClass) mapM_ (uncurry addToDB) ecls gets (_patDB . _eDB)+{-# INLINE createDB #-} -- | `addToDB` adds an e-node and e-class id to the database addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB ()@@ -254,6 +272,7 @@ case populate trie ids of -- populates the trie Nothing -> pure () Just t -> modify' $ over (eDB . patDB) (Map.insert op t) -- if something was created, insert back into the DB+{-# INLINE addToDB #-} -- | Populates an IntTrie with a sequence of e-class ids populate :: Maybe IntTrie -> [EClassId] -> Maybe IntTrie@@ -270,6 +289,7 @@ nextTrie = _trie tId IntMap.!? eid val = fromMaybe (trie eid IntMap.empty) $ populate nextTrie eids in Just $ IntTrie keys (IntMap.insert eid val (_trie tId))+{-# INLINE populate #-} canonizeMap :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m (Map ClassOrVar ClassOrVar, ClassOrVar) canonizeMap (subst, cv) = (,cv) <$> traverse g subst -- Map.fromList <$> traverse f (Map.toList subst)@@ -282,6 +302,7 @@ f (e1, Left e2) = do e2' <- canonical e2 pure (e1, Left e2') f (e1, e2) = pure (e1, e2)+{-# INLINE canonizeMap #-} applyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m () applyMatch costFun rule match' =@@ -293,6 +314,7 @@ do new_eclass <- reprPrat costFun (fst match) (target rule) merge costFun (getInt (snd match)) new_eclass pure ()+{-# INLINE applyMatch #-} applyMergeOnlyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m () applyMergeOnlyMatch costFun rule match' =@@ -306,6 +328,7 @@ Nothing -> pure () Just eid -> do merge costFun (getInt (snd match)) eid pure ()+{-# INLINE applyMergeOnlyMatch #-} -- | gets the e-node of the target of the rule -- TODO: add consts and modify@@ -326,12 +349,14 @@ rebuild costFun -- eid new_enode pure (Just eid) else gets ((Map.!? new_enode) . _eNodeToEClass)+{-# INLINE classOfENode #-} -- | adds the target of the rule into the e-graph reprPrat :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m EClassId reprPrat costFun subst (VarPat c) = canonical $ getInt $ subst Map.! Right (fromEnum c) reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target) add costFun (replaceChildren newChildren target)+{-# INLINE reprPrat #-} isValidHeight :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool isValidHeight match = do@@ -340,29 +365,38 @@ gets (_height . (IntMap.! ec') . _eClass) Right _ -> pure 0 pure $ h < 15+{-# INLINE isValidHeight #-} -- | returns `True` if the condition of a rule is valid for that match isValidConditions :: Monad m => Condition -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool isValidConditions cond match = gets $ cond (fst match)+{-# INLINE isValidConditions #-} -- * Tree to e-graph conversion and utility functions -- | Creates an e-graph from an expression tree fromTree :: Monad m => CostFun -> Fix SRTree -> EGraphST m EClassId fromTree costFun = cataM sequence (add costFun)+{-# INLINE fromTree #-} -- | Builds an e-graph from multiple independent trees fromTrees :: Monad m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId] fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) []+{-# INLINE fromTrees #-} -- | gets the best expression given the default cost function-getBest :: Monad m => EClassId -> EGraphST m (Fix SRTree)-getBest eid = do eid' <- canonical eid- best <- gets (_best . _info . (IntMap.! eid') . _eClass)- childs <- mapM getBest $ childrenOf best- pure . Fix $ replaceChildren childs best+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+{-# INLINE getBestExpr #-} +getBestENode eid = do eid' <- canonical eid+ gets (_best . _info . (IntMap.! eid') . _eClass)+{-# INLINE getBestENode #-}+ -- | returns one expression rooted at e-class `eId` -- TODO: avoid loopings getExpressionFrom :: Monad m => EClassId -> EGraphST m (Fix SRTree)@@ -383,6 +417,7 @@ isTerm (Const _) = True isTerm (Param _) = True isTerm _ = False+{-# INLINE getExpressionFrom #-} -- | returns all expressions rooted at e-class `eId` -- TODO: check for infinite list@@ -417,7 +452,25 @@ Param ix -> pure [Param ix] ts <- go ns pure (t:ts)+{-# INLINE getAllExpressionsFrom #-} +getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildEClasses eId' = do+ eId <- canonical eId'+ nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)+ nub <$> go eId++ where+ go :: Monad m => Int -> EGraphST m [Int]+ go n = do nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)+ let hasTerminal = any (null . childrenOf) nodes+ eids <- mapM canonical $ concatMap childrenOf nodes+ if hasTerminal+ then pure [n]+ else do eids' <- mapM go eids+ pure ((n : eids) <> concat eids')+{-# INLINE getAllChildEClasses #-}+ -- | returns a random expression rooted at e-class `eId` getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree) getRndExpressionFrom eId' = do@@ -434,6 +487,7 @@ randomRange rng = state (randomR rng) randomFrom xs = do n <- randomRange (0, length xs - 1) pure $ xs !! n+{-# INLINE getRndExpressionFrom #-} cleanMaps :: Monad m => EGraphST m () cleanMaps = do@@ -455,6 +509,8 @@ eDB' <- gets _eDB put $ EGraph canon enode2eclass' eclassMap' eDB' forceState+{-# INLINE cleanMaps #-} forceState :: Monad m => StateT s m () forceState = get >>= \ !_ -> return ()+{-# INLINE forceState #-}
src/Algorithm/EqSat/DB.hs view
@@ -31,12 +31,13 @@ import Data.HashSet (HashSet) import qualified Data.HashSet as Set import Data.String (IsString (..))+import Data.SRTree.Recursion (cata) import Debug.Trace -- A Pattern is either a fixed-point of a tree or an -- index to a pattern variable. The pattern variable matches anything. -data Pattern = Fixed (SRTree Pattern) | VarPat Char deriving Show -- Fixed structure of a pattern or a variable that matches anything+data Pattern = Fixed (SRTree Pattern) | VarPat Char deriving (Show, Eq, Ord) -- Fixed structure of a pattern or a variable that matches anything -- The instance for `IsString` for a `Pattern` is -- valid only for a single letter char from a-zA-Z. @@ -47,6 +48,14 @@ fromString [c] | n >= 65 && n <= 122 = VarPat c where n = fromEnum c fromString s = error $ "invalid string in VarPat: " <> s +tree2pat :: Fix SRTree -> Pattern+tree2pat = cata alg+ where+ alg (Param ix) = if ix >= 100 then VarPat (toEnum $ ix - 100 + 65) else Fixed $ Param ix+ alg (Var ix) = Fixed $ Var ix+ alg (Const x) = Fixed $ Const x+ alg (Bin op l r) = Fixed $ Bin op l r+ alg (Uni f t) = Fixed $ Uni f t -- A rule is either a directional rule where pat1 can be replaced by pat2, a bidirectional rule -- where pat1 can be replaced or replace pat2, or a pattern with a conditional function -- describing when to apply the rule @@ -149,19 +158,22 @@ target (r :| _) = target r target (_ :=> t) = t target (_ :==: t) = t+{-# INLINE target #-} source :: Rule -> Pattern source (r :| _) = source r source (s :=> _) = s source (s :==: _) = s+{-# INLINE source #-} getConditions :: Rule -> [Condition] getConditions (r :| c) = c : getConditions r getConditions _ = []-+{-# INLINE getConditions #-} cleanDB :: Monad m => EGraphST m () cleanDB = modify' $ over (eDB. patDB) (const Map.empty)+{-# INLINE cleanDB #-} -- | Returns the substitution rules -- for every match of the pattern `source` inside the e-graph.@@ -170,6 +182,7 @@ let (q, root) = compileToQuery src -- compile the source of the pattern into a query substs <- genericJoin q root -- find the substituion rules for this pattern pure [(s, s Map.! root) | s <- substs, Map.size s > 0]+{-# INLINE match #-} -- | Returns a Query (list of atoms) of a pattern compileToQuery :: Pattern -> (Query, ClassOrVar)@@ -195,17 +208,20 @@ atom = Atom root (replaceChildren roots pat) atoms' = atom:atoms pure (atoms', root)+{-# INLINE compileToQuery #-} -- get the value from the Either Int Int getInt :: ClassOrVar -> Int getInt (Left a) = a getInt (Right a) = a+{-# INLINE getInt #-} -- | returns the list of the children values getElems :: SRTree a -> [a] getElems (Bin _ l r) = [l,r] getElems (Uni _ t) = [t] getElems _ = []+{-# INLINE getElems #-} -- | Creates the substituion map for -- the pattern variables for each one of the@@ -225,7 +241,7 @@ maps <- forM cIds1 $ \classId -> do map (Map.insert x classId) <$> go (updateVar x classId atoms) vars pure (concat maps)-+{-# INLINE genericJoin #-} -- [Map.insert x classId y | classId <- domainX db x atoms -- , y <- go (updateVar x classId atoms) vars]@@ -237,7 +253,7 @@ domainX var atoms root = do let atoms' = filter (elemOfAtom var) atoms -- :: [ClassOrVar] -- look only in the atoms with this var map Left <$> intersectAtoms var atoms' root -- find the intersection of possible keys by each atom-+{-# INLINE domainX #-} --let ss = (map Left -- $ intersectAtoms var db -- $@@ -267,6 +283,7 @@ -- try to find an intersection of the tries that matches each atom of the pattern -- then -- else pure Set.empty+{-# INLINE intersectAtoms #-} -- | searches for the intersection of e-class ids that -- matches each part of the query.@@ -314,6 +331,7 @@ Nothing -> acc Just s -> Set.union acc s ) Set.empty (_trie trie)+{-# INLINE intersectTries #-} -- | updates all occurrence of var with the new id x updateVar :: ClassOrVar -> ClassOrVar -> Query -> Query@@ -322,6 +340,7 @@ replace (Atom r t) = let children = [if c == var then x else c | c <- getElems t] t' = replaceChildren children t in Atom (if r == var then x else r) t'+{-# INLINE updateVar #-} -- | checks whether two ClassOrVar are different -- only check if it is a pattern variable, else returns true@@ -329,6 +348,7 @@ isDiffFrom x y = case y of Left _ -> False Right z -> x /= z+{-# INLINE isDiffFrom #-} -- | checks if v is an element of an atom elemOfAtom :: ClassOrVar -> Atom -> Bool@@ -336,6 +356,7 @@ case root of Left _ -> v `elem` getElems tree Right x -> Right x == v || v `elem` getElems tree+{-# INLINE elemOfAtom #-} -- | sorts the variables in a query by the most frequently occurring orderedVars :: Query -> [ClassOrVar]@@ -349,3 +370,4 @@ varCost var = foldr (\a acc -> if elemOfAtom var a then acc - 100 + atomLen a else acc) 0 atoms atomLen (Atom _ t) = 1 + length (getElems t)+{-# INLINE orderedVars #-}
src/Algorithm/EqSat/Egraph.hs view
@@ -1,6 +1,7 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE StrictData #-}+{-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} -----------------------------------------------------------------------------@@ -22,7 +23,7 @@ import Control.Lens (element, makeLenses, view, over, (&), (+~), (-~), (.~), (^.)) --import Control.Monad (forM, forM_, when, foldM, void) import Data.List ( intercalate )-import Control.Monad.State.Strict+import Control.Monad.State.Strict hiding ( get, put ) import Data.IntMap.Strict (IntMap) import qualified Data.IntMap.Strict as IntMap import Data.Map.Strict (Map)@@ -37,7 +38,12 @@ import Data.SRTree import Data.SRTree.Eval import Data.Hashable+import Data.Binary+import qualified Data.Binary as Bin+import qualified Data.Massiv.Array as MA +import GHC.Generics+ import Debug.Trace type EClassId = Int -- NOTE: DO NOT CHANGE THIS, this will break the use of IntMap and IntSet@@ -125,6 +131,9 @@ n = FingerTree.length root +++ -- TODO: check this \/ getWithinRange :: Ord a => a -> a -> RangeTree a -> [EClassId] getWithinRange lb ub rt = map snd . toList $ go rt@@ -150,50 +159,105 @@ getSmallest rt = case rt of Empty -> error "empty finger" x :<| t -> x+{-# INLINE getSmallest #-}+ getGreatest :: Ord a => RangeTree a -> (a, EClassId) getGreatest rt = case rt of Empty -> error "empty finger" t :|> x -> x-+{-# INLINE getGreatest #-} data EGraph = EGraph { _canonicalMap :: ClassIdMap EClassId -- maps an e-class id to its canonical form , _eNodeToEClass :: Map ENode EClassId -- maps an e-node to its e-class id , _eClass :: ClassIdMap EClass -- maps an e-class id to its e-class data , _eDB :: EGraphDB- } deriving Show+ } deriving (Show, Generic) data EGraphDB = EDB { _worklist :: HashSet (EClassId, ENode) -- e-nodes and e-class schedule for analysis , _analysis :: HashSet (EClassId, ENode) -- e-nodes and e-class that changed data+ , _refits :: HashSet EClassId , _patDB :: DB -- database of patterns , _fitRangeDB :: RangeTree Double -- database of valid fitness+ , _dlRangeDB :: RangeTree Double , _sizeDB :: IntMap IntSet -- database of model sizes , _sizeFitDB :: IntMap (RangeTree Double) -- hacky! Size x Fitness DB+ , _sizeDLDB :: IntMap (RangeTree Double) , _unevaluated :: IntSet -- set of not-evaluated e-classes , _nextId :: Int -- next available id- } deriving Show+ } deriving (Show, Generic) data EClass = EClass { _eClassId :: Int -- e-class id (maybe we don't need that here) , _eNodes :: HashSet ENodeEnc -- set of e-nodes inside this e-class , _parents :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s , _height :: Int -- height , _info :: EClassData -- data- } deriving (Show, Eq)+ } deriving (Show, Eq, Generic) -data Consts = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq)-data Property = Positive | Negative | NonZero | Real deriving (Show, Eq) -- TODO: incorporate properties+data Consts = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq, Generic)+data Property = Positive | Negative | NonZero | Real deriving (Show, Eq, Generic) -- TODO: incorporate properties data EClassData = EData { _cost :: Cost , _best :: ENode , _consts :: Consts , _fitness :: Maybe Double -- NOTE: this cannot be NaN+ , _dl :: Maybe Double , _theta :: Maybe PVector , _size :: Int -- , _properties :: Property -- TODO: include evaluation of expression from this e-class- } deriving (Show)+ } deriving (Show, Generic) +-- * Serialization+instance Generic (EClassId, ENode)++instance Binary (SRTree EClassId) where+ put (Var ix) = put (0 :: Word8) >> put ix+ put (Param ix) = put (1 :: Word8) >> put ix+ put (Const x) = put (2 :: Word8) >> put x+ put (Uni f t) = put (3 :: Word8) >> put (fromEnum f) >> put t+ put (Bin op l r) = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r++ get = do t <- get :: Get Word8+ case t of+ 0 -> Var <$> get+ 1 -> Param <$> get+ 2 -> Const <$> get+ 3 -> Uni <$> (toEnum <$> get) <*> get+ 4 -> Bin <$> (toEnum <$> get) <*> get <*> get++instance Binary (SRTree ()) where+ put (Var ix) = put (0 :: Word8) >> put ix+ put (Param ix) = put (1 :: Word8) >> put ix+ put (Const x) = put (2 :: Word8) >> put x+ put (Uni f t) = put (3 :: Word8) >> put (fromEnum f)+ put (Bin op l r) = put (4 :: Word8) >> put (fromEnum op)++ get = do t <- get :: Get Word8+ case t of+ 0 -> Var <$> get+ 1 -> Param <$> get+ 2 -> Const <$> get+ 3 -> Uni <$> (toEnum <$> get) <*> pure ()+ 4 -> Bin <$> (toEnum <$> get) <*> pure () <*> pure ()++instance (Binary a, Hashable a) => Binary (HashSet a) where+ put hs = put (Set.toList hs)+ get = Set.fromList <$> get++instance Binary PVector where+ put xs = put (MA.toList xs)+ get = MA.fromList compMode <$> get++instance Binary IntTrie+instance Binary EClass+instance Binary Consts+instance Binary Property+instance Binary EClassData+instance Binary EGraphDB+instance Binary EGraph+ instance Eq EClassData where- EData c1 b1 cs1 ft1 _ s1 == EData c2 b2 cs2 ft2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && s1==s2+ EData c1 b1 cs1 ft1 dl1 _ s1 == EData c2 b2 cs2 ft2 dl2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && dl1==dl2 && s1==s2 -- The database maps a symbol to an IntTrie -- The IntTrie stores the possible paths from a certain e-class@@ -202,7 +266,7 @@ -- The IntTrie is composed of the set of available keys (for convenience) -- and an IntMap that maps one e-class id to the first child IntTrie, -- the first child IntTrie will point to the next child and so on-data IntTrie = IntTrie { _keys :: HashSet EClassId, _trie :: IntMap IntTrie } -- deriving Show+data IntTrie = IntTrie { _keys :: HashSet EClassId, _trie :: IntMap IntTrie } deriving (Generic) -- Shows the IntTrie as {keys} -> {show IntTries} instance Show IntTrie where@@ -220,10 +284,12 @@ -- | returns an empty e-graph emptyGraph :: EGraph emptyGraph = EGraph IntMap.empty Map.empty IntMap.empty emptyDB+{-# INLINE emptyGraph #-} -- | returns an empty e-graph DB emptyDB :: EGraphDB-emptyDB = EDB Set.empty Set.empty Map.empty FingerTree.empty IntMap.empty IntMap.empty IntSet.empty 0+emptyDB = EDB Set.empty Set.empty Set.empty Map.empty FingerTree.empty FingerTree.empty IntMap.empty IntMap.empty IntMap.empty IntSet.empty 0+{-# INLINE emptyDB #-} -- | Creates a new e-class from an e-class id, a new e-node, -- and the info of this e-class @@ -260,6 +326,7 @@ -- | Creates a singleton trie from an e-class id trie :: EClassId -> IntMap IntTrie -> IntTrie trie eid = IntTrie (Set.singleton eid)+{-# INLINE trie #-} -- | Check whether an e-class is a constant value isConst :: Monad m => EClassId -> EGraphST m Bool
src/Algorithm/EqSat/Info.hs view
@@ -42,16 +42,27 @@ -- TODO: instead of folding, just do not apply rules -- list of values instead of single value joinData :: EClassData -> EClassData -> EClassData-joinData (EData c1 b1 cn1 fit1 p1 sz1) (EData c2 b2 cn2 fit2 p2 sz2) =- EData (min c1 c2) b (combineConsts cn1 cn2) (minMaybe fit1 fit2) (bestParam p1 p2 fit1 fit2) (min sz1 sz2)+joinData (EData c1 b1 cn1 fit1 dl1 p1 sz1) (EData c2 b2 cn2 fit2 dl2 p2 sz2) =+ --EData (min c1 c2) b (combineConsts cn1 cn2) (minMaybe fit1 fit2) (bestParam p1 p2 fit1 fit2) (min sz1 sz2)+ EData (min c1 c2) (choose b1 b2) (choose cn1 cn2) (maxMaybe fit1 fit2) (choose dl1 dl2) (choose p1 p2) (choose sz1 sz2) where- minMaybe Nothing x = x- minMaybe x Nothing = x- minMaybe x y = min x y+ isFst = c1 <= c2+ choose x y = if isFst then x else y+ chooseF x y = if maxIsFst then x else y + maxIsFst = case (fit1, fit2) of+ (Nothing, Nothing) -> True+ (Nothing, Just f) -> False+ (Just f , Nothing) -> True+ (Just f1, Just f2) -> f1 >= f2++ maxMaybe Nothing x = x+ maxMaybe x Nothing = x+ maxMaybe x y = max x y+ bestParam Nothing x _ _ = x bestParam x Nothing _ _ = x- bestParam x y (Just f1) (Just f2) = if f1 < f2 then x else y+ bestParam x y (Just f1) (Just f2) = if f1 >= f2 then x else y b = if c1 <= c2 then b1 else b2 combineConsts (ConstVal x) (ConstVal y)@@ -68,6 +79,8 @@ combineConsts (ParamIx ix) (ParamIx iy) = ParamIx (min ix iy) combineConsts NotConst x = x combineConsts x NotConst = x+ combineConsts (ParamIx ix) (ConstVal x) = ConstVal x+ combineConsts (ConstVal x) (ParamIx ix) = ConstVal x -- p - p = 0 combineConsts x y = error (show x <> " " <> show y) -- | Calculate e-node data (constant values and cost)@@ -77,7 +90,7 @@ enode' <- canonize enode cost <- calculateCost costFun enode' sz <- sum <$> mapM (\ecId -> gets (_size . _info . (IntMap.! ecId) . _eClass)) (childrenOf enode')- pure $ EData cost enode' consts Nothing Nothing (sz+1)+ pure $ EData cost enode' consts Nothing Nothing Nothing (sz+1) getChildrenMinHeight :: Monad m => ENode -> EGraphST m Int getChildrenMinHeight enode = do@@ -146,6 +159,7 @@ combineConsts (Var _) = NotConst combineConsts (Uni f t) = case t of ConstVal x -> ConstVal $ evalFun f x+ --ParamIx x -> ParamIx x _ -> t combineConsts (Bin op l r) = evalOp' l r where@@ -154,10 +168,12 @@ evalOp' _ _ = NotConst insertFitness :: Monad m => EClassId -> Double -> PVector -> EGraphST m ()-insertFitness eId fit params = do+insertFitness eId' fit params = do+ eId <- canonical eId' ec <- gets ((IntMap.! eId) . _eClass) let oldFit = _fitness . _info $ ec- newInfo = (_info ec){_fitness = Just fit, _theta = Just params}+ --when (oldFit < Just fit) $ do+ let newInfo = (_info ec){_fitness = Just fit, _theta = Just params} newEc = ec{_info = newInfo} sz = _size newInfo modify' $ over eClass (IntMap.insert eId newEc)@@ -166,3 +182,14 @@ . over (eDB . fitRangeDB) (insertRange eId fit) . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty) else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit))++insertDL :: Monad m => EClassId -> Double -> EGraphST m ()+insertDL eId fit' = do+ let fit = negate fit'+ ec <- gets ((IntMap.! eId) . _eClass)+ let sz = _size . _info $ ec+ newInfo = (_info ec){_dl = Just fit'}+ newEc = ec{_info=newInfo}+ 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)
src/Algorithm/EqSat/Queries.hs view
@@ -1,5 +1,6 @@ {-# LANGUAGE ViewPatterns #-} {-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TupleSections #-} ----------------------------------------------------------------------------- -- | -- Module : Algorithm.EqSat.Queries@@ -25,6 +26,8 @@ import Control.Lens ( over ) import Data.Maybe import Data.Sequence ( Seq(..) )+import qualified Data.Sequence as FingerTree+import qualified Data.Foldable as Foldable import Debug.Trace @@ -57,9 +60,10 @@ -- | returns the e-class id with the best fitness that -- is true to a predicate-getTopECLassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]-getTopECLassThat n p = do- gets (_fitRangeDB . _eDB)+getTopECLassThat :: Monad m => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopECLassThat b n p = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB) >>= go n [] where go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]@@ -70,18 +74,117 @@ ecId <- canonical x ec <- gets ((IntMap.! ecId) . _eClass) if (isInfinite . fromJust . _fitness . _info $ ec)- then pure bests+ then go m bests t else if p ec- then go (m-1) (x:bests) t+ then go (m-1) (ecId:bests) t else go m bests t-getTopECLassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]-getTopECLassWithSize sz n = do- go n [] <$> gets ((IntMap.!? sz) . _sizeFitDB . _eDB)- >>= mapM canonical++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+ gets (f . _eDB)+ >>= go n [] where+ ecs = Set.fromList ecs'+ go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go 0 bests rt = pure bests+ go m bests 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 t -- pure bests+ else if ecId `Set.member` ecs && p ec+ then go (m-1) (ecId:bests) t+ else go m bests t+getTopECLassNotIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopECLassNotIn b n p ecs' = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB)+ >>= go n []+ where+ ecs = Set.fromList ecs'++ go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go 0 bests rt = pure bests+ go m bests 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 t+ else if not (ecId `Set.member` ecs) && p ec+ then go (m-1) (ecId:bests) t+ else go m bests t++getAllEvaluatedEClasses :: Monad m => EGraphST m [EClassId]+getAllEvaluatedEClasses = do+ gets (_fitRangeDB . _eDB)+ >>= go []+ where+ go :: Monad m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go bests 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 bests t+ else go (ecId:bests) t++getTopEClassWithSize :: Monad m => Bool -> Int -> Int -> EGraphST m [EClassId]+getTopEClassWithSize b sz n = do+ let fun = if b then _sizeFitDB else _sizeDLDB+ gets (go n [] . (IntMap.!? sz) . fun . _eDB)+ -- >>= mapM canonical+ where -- go :: Monad m => Int -> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId] go _ bests Nothing = [] go 0 bests (Just rt) = bests go m bests (Just rt) = case rt of Empty -> bests- t :|> (f, x) -> if isInfinite f then bests else go (m-1) (x:bests) (Just t)+ t :|> (f, x) -> if isInfinite f || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t)++getTopFitEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopFitEClassThat = getTopECLassThat True+getTopDLEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopDLEClassThat = getTopECLassThat False+getTopFitEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassIn = getTopECLassIn True+getTopDLEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassIn = getTopECLassIn False+getTopFitEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassNotIn = getTopECLassNotIn True+getTopDLEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassNotIn = getTopECLassNotIn True+getTopFitEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]+getTopFitEClassWithSize = getTopEClassWithSize True+getTopDLEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]+getTopDLEClassWithSize = getTopEClassWithSize False++rebuildAllRanges :: Monad m => EGraphST m ()+rebuildAllRanges = do szF <- gets (_sizeFitDB._eDB) >>= traverse rebuildRange+ dlF <- gets (_sizeDLDB._eDB) >>= traverse rebuildRange+ fR <- gets (_fitRangeDB._eDB) >>= rebuildRange+ dR <- gets (_dlRangeDB._eDB) >>= rebuildRange++ modify' $ over (eDB.fitRangeDB) (const fR)+ . over (eDB.dlRangeDB) (const dR)+ . over (eDB.sizeFitDB) (const szF)+ . over (eDB.sizeDLDB) (const dlF)++canonizeRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)+canonizeRange = traverse (\(x, eid) -> (x,) <$> canonical eid)++rebuildRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)+rebuildRange rt = go Set.empty Empty <$> canonizeRange rt+ where+ go :: Set.HashSet EClassId -> RangeTree Double -> RangeTree Double -> RangeTree Double+ go seen root Empty = root+ go seen root (xs :|> (x,eid)) = go (Set.insert eid seen)+ (if Set.member eid seen+ then root+ else (x, eid) :<| root)+ xs -- (Prelude.filter ((/= eid) . snd) xs)
src/Algorithm/EqSat/Simplify.hs view
@@ -12,7 +12,7 @@ -- Module containing the algebraic rules and simplification function. -- ------------------------------------------------------------------------------module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewriteBasic, rewritesFun ) where+module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple ) where import Algorithm.EqSat (eqSat, applySingleMergeOnlyEqSat) import Algorithm.EqSat.Egraph@@ -63,7 +63,7 @@ isNotZero :: ConstrFun isNotZero = constrainOnVal $ \case- ConstVal x -> abs x < 1e-9+ ConstVal x -> abs x > 1e-9 _ -> True -- check if the matched pattern is even @@ -98,11 +98,10 @@ rewriteBasic :: [Rule] rewriteBasic = [- "x" * "x" :=> "x" ** 2- , "x" * "y" :=> "y" * "x"+ "x" * "y" :=> "y" * "x" , "x" + "y" :=> "y" + "x"- , ("x" ** "y") * "x" :=> "x" ** ("y" + 1) :| isConstPt "y"- , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"+ --, ("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") -- TODO: check that I don't need that , ("x" * "y") * "z" :=> "x" * ("y" * "z")@@ -117,6 +116,7 @@ , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y") -- TODO handle with power :| isConstPt "w" :| isConstPt "z" :| isNotZero "z" -- TODO: a + b*y :=> b * (a/b + y) :| isNotZero b , (("x" * "y") + ("z" * "w")) :=> "x" * ("y" + ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"+ -- , "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" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL @@ -127,15 +127,13 @@ rewritesFun :: [Rule] rewritesFun = [- log (sqrt "x") :=> 0.5 * log "x" :| isNotParam "x"- , log (exp "x") :==: exp (log "x")+ log (exp "x") :==: exp (log "x") , log (exp "x") :=> "x" -- , exp (log "x") :=> "x" -- :| isPositive "x" ??? exp(log(x)), x, log(exp(0))- , "x" ** (1/2) :==: sqrt "x" -- <==>- , "x" ** (1/3) :==: Fixed (Uni Cbrt "x") , log ("x" * "y") :=> log "x" + log "y" :| isConstPos "x" :| isConstPos "y" -- , log ("x" / "y") :=> log "x" - log "y" :| isConstPos "x" :| isConstPos "y" , log ("x" ** "y") :=> "y" * log "x"+ , log (powabs "x" "y") :=> "y" * log (abs "x") --, sqrt ("x" ** "y") :=> "x" ** ("y" / 2) :| isEven "y" -- , sqrt ("y" * "x") :=> sqrt "y" * sqrt "x" -- --, sqrt ("y" / "x") :=> sqrt "y" / sqrt "x"@@ -158,30 +156,91 @@ [ 0 + "x" :=> "x" -- , "x" - 0 :=> "x"- , 1 * "x" :=> "x"- , 0 * "x" :=> 0 :| isValid "x" -- :| isNotParam "x"+ --, 1 * "x" :=> "x" -- , 0 / "x" :=> 0 :| isNotZero "x" --, "x" - "x" :=> 0 :| isNotParam "x" --, "x" / "x" :=> 1 :| isNotZero "x" :| isNotParam "x" , "x" ** 1 :=> "x"- , 0 ** "x" :=> 0 :| isPositive "x"- , 1 ** "x" :=> 1+ , powabs "x" 1 :=> abs "x"+ -- , "x" * (1 / "x") :=> 1 :| isNotParam "x" :| isNotZero "x"- , 0 - "x" :=> negate "x"- , "x" + negate "y" :==: "x" - "y" -- , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x"- , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isPositive "x"+ , "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") :| isPositive "x"+ , powabs (powabs "x" "y") "z" :=> powabs "x" ("y" * "z") , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isPositive "x" :| isPositive "y" - , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isInteger "y" :| isNotZero "x"- , "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"- , ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"- , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isInteger "z" :| isNotZero "x" :| isNotZero "y"+ --, "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"+ --, ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"+ --, ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isInteger "z" :| isNotZero "x" :| isNotZero "y" ] +rewritesWithConstant :: [Rule]+rewritesWithConstant =+ [+ "x" * "x" :=> "x" ** 2+ , "x" - "x" :=> 0+ , "x" / "x" :=> 1 :| isNotZero "x"+ , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isPositive "x"+ , 1 ** "x" :=> 1+ , powabs 1 "x" :=> 1+ , log (sqrt "x") :=> 0.5 * log "x" :| isNotParam "x"+ , "x" ** (1/2) :==: sqrt "x" -- <==>+ , powabs "x" (1/2) :=> sqrt (abs "x")+ , "x" ** (1/3) :==: Fixed (Uni Cbrt "x")+ , 0 * "x" :=> 0 :| isValid "x" -- :| isNotParam "x"+ , 0 ** "x" :=> 0 :| isPositive "x"+ , powabs 0 "x" :=> 0+ , 0 - "x" :=> negate "x"+ , "x" + negate "y" :==: "x" - "y"+ ]+rewritesWithParam :: [Rule]+rewritesWithParam =+ [+ -- "x" * "x" :=> "x" ** Fixed (Param 0)+ "x" - "x" :=> Fixed (Param 0)+ , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"+ , 1 ** "x" :=> Fixed (Param 0)+ , powabs 1 "x" :=> Fixed (Param 0)+ -- , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"+ ]++rewritesSimple :: [Rule]+rewritesSimple =+ [+ "x" * "y" :=> "y" * "x"+ , "x" + "y" :=> "y" + "x"+ , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"+ , ("x" + "y") + "z" :=> "x" + ("y" + "z")+ , ("x" * "y") * "z" :=> "x" * ("y" * "z")+ , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")+ , "x" - ("y" + "z") :=> ("x" - "y") - "z" -- TODO: check that I don't this+ , "x" - ("y" - "z") :=> ("x" - "y") + "z" -- TODO+ , ("x" * "y") / "z" :=> ("x" / "z") * "y" :| isNotZero "z" -- TODO: inv(x) <=> x^-1 , x/y <=> x*y^-1+ , "x" * ("y" / "z") :=> ("x" / "z") * "y" :| isNotZero "z" -- ^+ , "x" / ("y" * "z") :=> ("x" / "z") / "y" :| isNotZero "z" -- ^ TODO: 0 ^-1 check+ , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") - ("z" * "x") :=> ("w" - "z") * "x" -- TODO: handle sub :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y")+ , log (exp "x") :=> "x"+ , exp (log "x") :=> "x"+ , log ("x" * "y") :=> log "x" + log "y"+ , log ("x" ** "y") :=> "y" * log "x"+ , abs ("x" * "y") :=> abs "x" * abs "y"+ , abs ("x" ** "y") :=> abs "x" ** "y"+ , abs ("x" - "y") :=> abs ("y" - "x")+ , recip (recip "x") :=> "x" :| isNotZero "x"+ , "x" * "x" :=> "x" ** Fixed (Param 0)+ , "x" - "x" :=> Fixed (Param 0)+ , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"+ , 1 ** "x" :=> Fixed (Param 0)+ , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"+ ]+powabs l r = Fixed (Bin PowerAbs l r)+ -- | default cost function for simplification -- TODO: -- num_params:@@ -192,14 +251,16 @@ -- univariates myCost :: SRTree Int -> Int myCost (Var _) = 1-myCost (Const _) = 1-myCost (Param _) = 1+myCost (Const _) = 3+myCost (Param _) = 3 myCost (Bin op l r) = 2 + l + r myCost (Uni _ t) = 3 + t -- all rewrite rules rewrites :: [Rule]-rewrites = rewriteBasic <> constReduction <> rewritesFun+rewrites = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithConstant+rewritesParams :: [Rule]+rewritesParams = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithParam -- | simplify using the default parameters simplifyEqSatDefault :: Fix SRTree -> Fix SRTree
src/Algorithm/SRTree/AD.hs view
@@ -4,6 +4,9 @@ {-# language ViewPatterns #-} {-# language FlexibleContexts #-} {-# language BangPatterns #-}+{-# language TypeApplications #-}+{-# language MultiWayIf #-}+ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.AD @@ -18,14 +21,12 @@ ----------------------------------------------------------------------------- module Algorithm.SRTree.AD- ( forwardMode- , forwardModeUnique- , reverseModeUnique- , reverseModeUniqueArr+ ( reverseModeArr+ , reverseModeGraph , forwardModeUniqueJac ) where -import Control.Monad (forM_, foldM)+import Control.Monad (forM_, foldM, when) import Control.Monad.ST ( runST ) import Data.Bifunctor (bimap, first, second) import qualified Data.DList as DL@@ -44,430 +45,272 @@ import GHC.IO (unsafePerformIO) import qualified Data.IntMap.Strict as IntMap import Data.List ( foldl' )--applyUni :: (Index ix, Source r e, Floating e, Floating b) => Function -> Either (Array r ix e) b -> Either (Array D ix e) b-applyUni f (Left t) =- Left $ M.map (evalFun f) t-applyUni f (Right t) =- Right $ evalFun f t-{-# INLINE applyUni #-}--applyDer :: (Index ix, Source r e, Floating e, Floating b) => Function -> Either (Array r ix e) b -> Either (Array D ix e) b-applyDer f (Left t) =- Left $ M.map (derivative f) t-applyDer f (Right t) =- Right $ derivative f t-{-# INLINE applyDer #-}+import qualified Data.Vector.Storable as VS+import Control.Scheduler +import Data.Maybe ( fromJust ) -negate' :: (Index ix, Source r e, Num e, Num b) => Either (Array r ix e) b -> Either (Array D ix e) b-negate' (Left t) = Left $ M.map negate t-negate' (Right t) = Right $ negate t-{-# INLINE negate' #-}+import Control.Monad.State.Strict -applyBin :: (Index ix, Floating b) => Op -> Either (Array D ix b) b -> Either (Array D ix b) b -> Either (Array D ix b) b-applyBin op (Left ly) (Left ry) =- Left $ case op of- Add -> ly !+! ry- Sub -> ly !-! ry- Mul -> ly !*! ry- Div -> ly !/! ry- Power -> ly .** ry- PowerAbs -> M.map abs (ly .** ry)- AQ -> ly !/! (M.map sqrt (M.map (+1) (ry !*! ry)))+--import UnliftIO.Async -applyBin op (Left ly) (Right ry) =- Left $ unsafeLiftArray (\ x -> evalOp op x ry) ly-applyBin op (Right ly) (Left ry) =- Left $ unsafeLiftArray (\ x -> evalOp op ly x) ry-applyBin op (Right ly) (Right ry) =- Right $ evalOp op ly ry-{-# INLINE applyBin #-}+import qualified Data.Map.Strict as Map --- | get the value of a certain index if it is an array (Left) --- or returns the value itself if it is a scalar.-(!??) :: (Manifest r e, Index ix) => Either (Array r ix e) e -> ix -> e-(Left y) !?? ix = y ! ix-(Right y) !?? ix = y-{-# INLINE (!??) #-}+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+ , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]])+ where+ yErr = fromJust mYErr+ --ys = delay ys'+ m = M.size ys+ p = VS.length theta+ comp = M.getComp xss+ one :: Array S Ix1 Double+ one = M.replicate comp m 1+ (key2int, int2key, cachedVal, (subtract 1) -> root) = cataM leftToRight alg tree `execState` (Map.empty, IntMap.empty, IntMap.empty, 0)+ (key2int', int2key', cachedVal', cachedGrad) = calcGrad root one `execState` (key2int, int2key, cachedVal, Map.empty) --- | Calculates the results of the error vector multiplied by the Jacobian of an expression using forward mode--- provided a vector of variable values `xss`, a vector of parameter values `theta` and--- a function that changes a Double value to the type of the variable values.--- uses unsafe operations to use mutable array instead of a tape-forwardMode :: Array S Ix2 Double -> Array S Ix1 Double -> SRVector -> Fix SRTree -> (Array D Ix1 Double, Array S Ix1 Double)-forwardMode xss theta err tree = let (yhat, jacob) = runST $ cataM lToR alg tree- in (fromEither yhat, computeAs S err ><! jacob)- where - (Sz p) = M.size theta- (Sz (m :. n)) = M.size xss- cmp = getComp xss- -- | if the tree does not use a variable - -- it will return a single scalar, fromEither fixes this- fromEither (Left y) = y- fromEither (Right y) = M.replicate cmp (Sz m) y+ calcGrad :: Int -> Array S Ix1 Double -> State (Map.Map (SRTree Int) Int, IntMap.IntMap (SRTree Int), IntMap.IntMap (Array S Ix1 Double), Map.Map (SRTree Int) (Array S Ix1 Double)) ()+ calcGrad key v = do node <- gets ((IntMap.! key) . _int2key)+ case node of+ Bin op l r -> do xl <- gets (getVal l)+ xr <- gets (getVal r)+ (dl, dr) <- diff op v xl xr l r+ calcGrad l dl+ calcGrad r dr+ Uni f t -> do x <- gets (getVal t)+ calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x))+ Param ix -> modify' (insertGrad v (Param ix))+ _ -> pure ()+ where+ _int2key (_, b, _, _) = b+ insertGrad v k (a, b, c, g) = (a, b, c, Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g) - -- if it is a variable, returns the value of that variable and an array of zeros (Jacobian)- alg (Var ix) = do tape <- M.newMArray (Sz2 m p) 0 - >>= UMA.unsafeFreeze cmp- pure (Left (xss <! ix), tape)+ 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)+ getVal key (a, b, c, d) = c IntMap.! key+ getKey key (a, b, c, d) = a Map.! key - -- if it is a constant, returns the value of the constant and array of zeros - alg (Const c) = do tape <- M.newMArray (Sz2 m p) 0- >>= UMA.unsafeFreeze cmp- pure (Right c, tape)+ leftToRight (Uni f mt) = Uni f <$> mt;+ leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+ leftToRight (Var ix) = pure (Var ix)+ leftToRight (Param ix) = pure (Param ix)+ leftToRight (Const c) = pure (Const c) - -- if it is a parameter, returns the value of the parameter and the jacobian with a one in the corresponding column- alg (Param ix) = do tape <- M.makeMArrayS (Sz2 m p) (\(i :. j) -> pure $ if j==ix then 1 else 0)- >>= UMA.unsafeFreeze cmp- pure (Right (theta ! ix), tape)+ evalKey (Var ix) = pure $ if ix == -1+ then ys+ else if ix == -2+ then yErr+ else M.computeAs S $ xss <! ix+ evalKey (Const v) = pure $ M.replicate comp m v+ evalKey (Param ix) = pure $ M.replicate comp m (theta VS.! ix)+ evalKey (Uni f t) = M.computeAs S . M.map (evalFun f) <$> gets (getVal t)+ evalKey (Bin op l r) = M.computeAs S <$> (M.zipWith (evalOp op) <$> gets (getVal l) <*> gets (getVal r)) - -- 1. applies the derivative of f in the evaluated child - -- 2. replaces the value of the Jacobian at (i, j) with yi * J[i, j]- alg (Uni f (t, tape')) = do let y = computeAs S . fromEither $ applyDer f t- tape <- UMA.unsafeThaw tape'- forM_ [0 .. m-1] $ \i -> do- let yi = y ! i- forM_ [0 .. p-1] $ \j -> do- v <- UMA.unsafeRead tape (i :. j)- UMA.unsafeWrite tape (i :. j) (yi * v)- tapeF <- UMA.unsafeFreeze cmp tape- pure (applyUni f t, tapeF)- -- li, ri are the corresponding values of the evaluated left and right children - -- vl, vr are the corresponding value of the Jacobian at (i, j) - -- applies the corresponding derivative of each binary operator - alg (Bin op (l, tl') (r, tr')) = do- tl <- UMA.unsafeThaw tl'- tr <- UMA.unsafeThaw tr'- let l' = case l of- Left y -> Left $ computeAs S y- Right v -> Right v- r' = case r of- Left y -> Left $ computeAs S y- Right v -> Right v- forM_ [0 .. m-1] $ \i -> do - let li = l' !?? i- ri = r' !?? i- forM_ [0 .. p-1] $ \j -> do - vl <- UMA.unsafeRead tl (i :. j)- vr <- UMA.unsafeRead tr (i :. j)- UMA.unsafeWrite tl (i :. j) $ case op of- Add -> (vl+vr)- Sub -> (vl-vr)- Mul -> (vl * ri + vr * li)- Div -> ((vl * ri - vr * li) / ri^2)- Power -> (li ** (ri - 1) * (ri * vl + li * log li * vr))- PowerAbs -> (abs li ** ri) * (vr * log (abs li) + ri * vl / li)- AQ -> ((1 + ri*ri) * vl - li * ri * vr) / (1 + ri*ri) ** 1.5- tlF <- UMA.unsafeFreeze cmp tl- pure (applyBin op l r, tlF)+ alg (Var ix) = insertKey (Var ix)+ alg (Param ix) = insertKey (Param ix)+ alg (Const v) = insertKey (Const v)+ alg (Uni f t) = insertKey (Uni f t)+ alg (Bin op l r) = insertKey (Bin op l r) + --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+ k <- gets (getKey (Bin Div l r))+ v <- gets (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+ k <- gets (getKey (Bin Power l r))+ v <- gets (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) - lToR (Var ix) = pure (Var ix)- lToR (Param ix) = pure (Param ix)- lToR (Const c) = pure (Const c)- lToR (Uni f mt) = Uni f <$> mt- lToR (Bin op ml mr) = Bin op <$> ml <*> mr+ diff PowerAbs dx fx gy l r = do+ k <- gets (getKey (Bin PowerAbs l r))+ v <- gets (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) --- | The function `forwardModeUnique` calculates the numerical gradient of the tree and evaluates the tree at the same time. It assumes that each parameter has a unique occurrence in the expression. This should be significantly faster than `forwardMode`.-forwardModeUnique :: SRMatrix -> PVector -> SRVector -> Fix SRTree -> (SRVector, Array S Ix1 Double)-forwardModeUnique xss theta err = second (toGrad . DL.toList) . cata alg- where- (Sz n) = M.size theta- one = replicateAs xss 1- toGrad grad = M.fromList (getComp xss) [g !.! err | g <- grad]+ 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) - alg (Var ix) = (xss <! ix, DL.empty)- alg (Param ix) = (replicateAs xss $ theta ! ix, DL.singleton one)- alg (Const c) = (replicateAs xss c, DL.empty)- alg (Uni f (v, gs)) = let v' = evalFun f v- dv = derivative f v- in (v', DL.map (*dv) gs)- alg (Bin Add (v1, l) (v2, r)) = (v1+v2, DL.append l r)- alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, DL.append l (DL.map negate r))- alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, DL.append (DL.map (*v2) l) (DL.map (*v1) r))- alg (Bin Div (v1, l) (v2, r)) = let dv = ((-v1)/(v2*v2)) - in (v1/v2, DL.append (DL.map (/v2) l) (DL.map (*dv) r))- alg (Bin Power (v1, l) (v2, r)) = let dv1 = v1 ** (v2 - one)- dv2 = v1 * log v1- in (v1 ** v2, DL.map (*dv1) (DL.append (DL.map (*v2) l) (DL.map (*dv2) r)))- alg (Bin PowerAbs (v1, l) (v2, r)) = let dv1 = abs v1 ** v2- dv2 = DL.map (* (log (abs v1))) r- dv3 = DL.map (*(v2 / v1)) l- in (abs v1 ** v2, DL.map (*dv1) (DL.append dv2 dv3))- alg (Bin AQ (v1, l) (v2, r)) = let dv1 = DL.map (*(1 + v2*v2)) l- dv2 = DL.map (*(-v1*v2)) r- in (v1/sqrt(1 + v2*v2), DL.map (/(1 + v2*v2)**1.5) $ DL.append dv1 dv2)+ fixNaN x = if isNaN x then 0 else x -data TupleF a b = Single a | T a b | Branch a b b deriving Functor -- hi, I'm a tree-type Tuple a = Fix (TupleF a)+ insertKey key = do+ isCached <- gets ((key `Map.member`) . graph)+ when (not isCached) $ do+ ev <- evalKey key+ modify' (insKey key ev)+ gets (getKey key) --- | Same as above, but using reverse mode, that is even faster.-reverseModeUnique :: SRMatrix+-- | Same as above, but using reverse mode with the tree encoded as an array, that is even faster.+reverseModeArr :: SRMatrix -> PVector- -> SRVector- -> (SRVector -> SRVector)- -> Fix SRTree+ -> Maybe PVector+ -> VS.Vector Double -- PVector+ -> [(Int, (Int, Int, Int, Double))] -- arity, opcode, ix, const val+ -> IntMap.IntMap Int -> (Array D Ix1 Double, Array S Ix1 Double)-reverseModeUnique xss theta ys f t = unsafePerformIO $- do jacob <- M.newMArray (Sz p) 0- let !_ = accu reverse (combine jacob) t ((Right 1), fwdMode)- j <- freezeS jacob- pure (v, j)- where- fwdMode = cata forward t- v = fromEither $ getTop fwdMode- err = f v - ys- (Sz2 m _) = M.size xss- p = countParams t- fromEither (Left x) = x- fromEither (Right x) = M.replicate (getComp xss) (Sz1 m) x-- oneTpl x = Fix $ Single x- tuple x y = Fix $ T x y- branch x y z = Fix $ Branch x y z-- getTop (Fix (Single x)) = x- getTop (Fix (T x y)) = x- getTop (Fix (Branch x y z)) = x-- unCons (Fix (T x y)) = y- getBranches (Fix (Branch x y z)) = (y,z)-- -- forward just creates a new tree with the partial- -- evaluation of the nodes- forward (Var ix) = oneTpl (Left $ xss <! ix)- forward (Param ix) = oneTpl (Right $ theta ! ix)- forward (Const c) = oneTpl (Right c)- forward (Uni g t) = let v = getTop t- in tuple (applyUni g v) t- forward (Bin op l r) = let vl = getTop l- vr = getTop r- in branch (applyBin op vl vr) l r---- -- reverse walks from the root to the leaf calculating the- -- partial derivative with respect to an arbitrary variable- -- up to that point- reverse (Var ix) (dx, _) = Var ix- reverse (Param ix) (dx, _) = Param ix- reverse (Const v) (dx, _) = Const v- reverse (Uni f t) (dx, unCons -> v) =- let g' = applyDer f (getTop v)- in Uni f (t, ( applyBin Mul dx g', v ))- reverse (Bin op l r) (dx, getBranches -> (vl, vr)) =- let (dxl, dxr) = diff op dx (getTop vl) (getTop vr)- in Bin op (l, (dxl, vl)) (r, (dxr, vr))-- -- dx is the current derivative so far- -- fx is the evaluation of the left branch- -- gx is the evaluation of the right branch- --- -- this should return a tuple, where the left element is- -- dx * d op(f(x), g(x)) / d f(x) and- -- the right branch dx * d op (f(x), g(x)) / d g(x)- diff Add dx fx gy = (dx, dx)- diff Sub dx fx gy = (dx, negate' dx)- diff Mul dx fx gy = (applyBin Mul dx gy, applyBin Mul dx fx)- diff Div dx fx gy = (applyBin Div dx gy, applyBin Mul dx (applyBin Div (negate' fx) (applyBin Mul gy gy)))- diff Power dx fx gy = let dxl = applyBin Mul dx (applyBin Power fx (applyBin Sub gy (Right 1)))- dv2 = applyBin Mul fx (applyUni Log fx)- in (applyBin Mul dxl gy, applyBin Mul dxl dv2)- diff PowerAbs dx fx gy = let dxl = applyBin Mul (applyBin Mul gy fx) (applyBin PowerAbs fx (applyBin Sub gy (Right 2)))- dxr = applyBin Mul (applyUni LogAbs fx) (applyBin PowerAbs fx gy)- in (applyBin Mul dxl dx, applyBin Mul dxr dx)- diff AQ dx fx gy = let dxl = applyUni Recip (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1)))- dxy = applyBin Div (applyBin Mul fx gy) (applyUni Cube (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1))))- in (applyBin Mul dxl dx, applyBin Mul dxy dx)--- -- once we reach a leaf with a parameter, we return a singleton- -- with that derivative upwards until the root- --combine :: (forall s . MArray (PrimState (ST s)) S Int Double) -> SRTree () -> (Either SRVector Double, a) -> ()- combine j (Var ix) s = 0- combine j (Const _) s = 0- combine j (Param ix) s = unsafePerformIO $ do- case fst s of- Left v -> do v' <- dotM v err- UMA.unsafeWrite j ix v'- Right v -> UMA.unsafeWrite j ix $ M.foldrS (\x acc -> x*v + acc) 0 err- UMA.unsafeRead j ix- combine j (Uni f gs) s = gs- combine j (Bin op l r) s = l+r---- | Same as above, but using reverse mode with the tree encoded as an array, that is even faster.---reverseModeUniqueArr :: SRMatrix--- -> PVector--- -> SRVector--- -> (SRVector -> SRVector)--- -> Array S Ix1 (Int, Int, Int, Double) -- arity, opcode, ix, const val--- -> (Array D Ix1 Double, Array S Ix1 Double)-reverseModeUniqueArr xss theta ys f t j2ix =- {-let fwd = forward- v = fwd IntMap.! 0- err = f v - delay ys- partial = reverseMode fwd- in -}+reverseModeArr xss ys mYErr theta t j2ix = unsafePerformIO $ do- fwd <- M.newMArray (Sz2 m n) 0- partial <- M.newMArray (Sz2 m n) 0+ fwd <- M.newMArray (Sz2 n m) 0+ partial <- M.newMArray (Sz2 n m) 0 jacob <- M.newMArray (Sz p) 0- fwd' <- UMA.unsafeFreeze (getComp xss) fwd- let v = fwd' M.<! 0- err = M.computeAs S $ f v - delay ys- forward fwd- combine partial jacob err+ val <- M.newMArray (Sz m) 0+ let+ stps = 2+ --delta = m `div` stps+ --rngs = [(i*delta, min m $ (i+1)*delta) | i <- [0..stps] ]+ (a, b) = (0, m)++ forward (a, b) fwd+ calculateYHat (a, b) fwd val+ reverseMode (a, b) fwd partial+ combine (a, b) partial jacob j <- UMA.unsafeFreeze (getComp xss) jacob- pure (v, j)+ v <- UMA.unsafeFreeze (getComp xss) val+ pure (delay v, j) where (Sz2 m _) = M.size xss- (Sz p) = M.size theta+ p = VS.length theta n = length t+ toLin i j = i*m + j+ yErr = fromJust mYErr+ eps = 1e-8 - forward :: MArray (PrimState IO) S Ix2 Double -> IO ()- forward fwd = forM_ (Prelude.reverse t) makeFwd+ myForM_ [] _ = pure ()+ myForM_ (!x:xs) f = do f x+ myForM_ xs f+ {-# INLINE myForM_ #-}++ calculateYHat :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO ()+ calculateYHat (a, b) fwd yhat = myForM_ [a..b-1] $ \i -> do+ vi <- UMA.unsafeRead fwd (0 :. i)+ UMA.unsafeWrite yhat i vi+ {-# INLINE calculateYHat #-}++ forward :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> IO ()+ forward (a, b) fwd = do+ let t' = Prelude.reverse t+ myForM_ t' makeFwd where- makeFwd (j, (0, 0, ix, _)) = do let j' = j2ix IntMap.! j- forM_ [0..m-1] $ \i -> do- let val = xss M.! (i :. ix)- UMA.unsafeWrite fwd (i :. j') val+ makeFwd (j, (0, 0, ix, _)) =+ do let j' = j2ix IntMap.! j+ myForM_ [a..b-1] $ \i -> do+ --let val = xss M.! (i :. ix)+ UMA.unsafeWrite fwd (j' :. i) $ case ix of+ (-1) -> ys M.! i+ (-2) -> yErr M.! i+ _ -> xss M.! (i :. ix) makeFwd (j, (0, 1, ix, _)) = do let j' = j2ix IntMap.! j- v = theta M.! ix- forM_ [0..m-1] $ \i -> do- UMA.unsafeWrite fwd (i :. j') v+ v = theta VS.! ix+ myForM_ [a..b-1] $ \i -> do+ UMA.unsafeWrite fwd (j' :. i) v makeFwd (j, (0, 2, _, x)) = do let j' = j2ix IntMap.! j- forM_ [0..m-1] $ \i -> do- UMA.unsafeWrite fwd (i :. j') x+ myForM_ [a..b-1] $ \i -> do+ UMA.unsafeWrite fwd (j' :. i) x makeFwd (j, (1, f, _, _)) = do let j' = j2ix IntMap.! j j2 = j2ix IntMap.! (2*j + 1)- forM_ [0..m-1] $ \i -> do- v <- UMA.unsafeRead fwd (i :. j2)- let val = evalFun (toEnum f) v- UMA.unsafeWrite fwd (i :. j') val+ myForM_ [a..b-1] $ \i -> do+ v <- UMA.unsafeRead fwd (j2 :. i)+ UMA.unsafeWrite fwd (j' :. i) (evalFun (toEnum f) v) makeFwd (j, (2, op, _, _)) = do let j' = j2ix IntMap.! j j2 = j2ix IntMap.! (2*j + 1) j3 = j2ix IntMap.! (2*j + 2)- forM_ [0..m-1] $ \i -> do- l <- UMA.unsafeRead fwd (i :. j2)- r <- UMA.unsafeRead fwd (i :. j3)- let val = evalOp (toEnum op) l r- UMA.unsafeWrite fwd (i :. j') val- {-- forward = foldr (makeFwd) IntMap.empty (IntMap.toAscList t)- where- makeFwd (j, (0, 0, ix, _)) fwd = IntMap.insert j (xss M.<! ix) fwd- makeFwd (j, (0, 1, ix, _)) fwd = IntMap.insert j (M.replicate (getComp xss) (M.Sz m) (theta M.! ix)) fwd- makeFwd (j, (0, 2, _, x)) fwd = IntMap.insert j (M.replicate (getComp xss) (M.Sz m) x) fwd- makeFwd (j, (1, f, _, _)) fwd = let v = fwd IntMap.! (2*j + 1)- val = M.map (evalFun (toEnum f)) v- in IntMap.insert j val fwd- makeFwd (j, (2, op, _, _)) fwd = let l = fwd IntMap.! (2*j + 1)- r = fwd IntMap.! (2*j + 2)- val = M.zipWith (evalOp (toEnum op)) l r- in IntMap.insert j val fwd- -}-+ myForM_ [a..b-1] $ \i -> do+ l <- UMA.unsafeRead fwd (j2 :. i)+ r <- UMA.unsafeRead fwd (j3 :. i)+ UMA.unsafeWrite fwd (j' :. i) (evalOp (toEnum op) l r)+ makeFwd _ = pure ()+ {-# INLINE makeFwd #-}+ {-# INLINE forward #-} - -- reverse walks from the root to the leaf calculating the- -- partial derivative with respect to an arbitrary variable- -- up to that point- reverseMode :: MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix2 Double -> IO ()- reverseMode fwd partial = do forM_ [0..m-1] $ \i -> UMA.unsafeWrite partial (i :. 0) 1- forM_ t makeRev+ reverseMode :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix2 Double -> IO ()+ reverseMode (a, b) fwd partial =+ do myForM_ [a..b-1] $ \i -> UMA.unsafeWrite partial (0 :. i) 1+ myForM_ t makeRev where- makeRev (j, (1, f, _, _)) = do forM_ [0..m-1] $ \i -> do- let dxj = j2ix IntMap.! j- vj = j2ix IntMap.! (2*j + 1)- v <- UMA.unsafeRead fwd (i :. vj)- dx <- UMA.unsafeRead partial (i :. dxj)- let val = dx * derivative (toEnum f) v- UMA.unsafeWrite partial (i :. vj) val- makeRev (j, (2, op, _, _)) = do forM_ [0..m-1] $ \i -> do- let dxj = j2ix IntMap.! j- lj = j2ix IntMap.! (2*j + 1)- rj = j2ix IntMap.! (2*j + 2)- l <- UMA.unsafeRead fwd (i :. lj)- r <- UMA.unsafeRead fwd (i :. rj)- dx <- UMA.unsafeRead partial (i :. dxj)+ makeRev (j, (1, f, _, _)) = do let dxj = j2ix IntMap.! j+ vj = j2ix IntMap.! (2*j + 1)+ myForM_ [a..b-1] $ \i -> do+ v <- UMA.unsafeRead fwd (vj :. i)+ dx <- UMA.unsafeRead partial (dxj :. i)+ --let val = dx * derivative (toEnum f) v+ UMA.unsafeWrite partial (vj :. i) (dx * derivative (toEnum f) v)+ makeRev (j, (2, op, _, _)) = do let dxj = j2ix IntMap.! j+ lj = j2ix IntMap.! (2*j + 1)+ rj = j2ix IntMap.! (2*j + 2)+ myForM_ [a..b-1] $ \i -> do+ l <- UMA.unsafeRead fwd (lj :. i)+ r <- UMA.unsafeRead fwd (rj :. i)+ dx <- UMA.unsafeRead partial (dxj :. i) let (dxl, dxr) = diff (toEnum op) dx l r- UMA.unsafeWrite partial (i :. lj) dxl- UMA.unsafeWrite partial (i :. rj) dxr+ UMA.unsafeWrite partial (lj :. i) dxl+ UMA.unsafeWrite partial (rj :. i) dxr makeRev _ = pure ()- {-- reverseMode fwd = foldr (makeRev) rev0 (IntMap.toDescList t)- where- rev0 = IntMap.insert 0 (M.replicate (getComp xss) (M.Sz m) 1) IntMap.empty+ {-# INLINE makeRev #-}+ {-# INLINE reverseMode #-} + --f(x)^g(x)+ --d f(x)^g(x) / d f(x) = f(x)^(g(x)-1)+ -- f(x) + g(x) = 1, 1+ -- f(x) - g(x) = 1, -1+ -- f(x) * g(x) = g(x), f(x)+ -- f(x) / g(x) = 1/g(x), -f(x)/g(x)^2+ -- f(x) ^ g(x) = g(x) * f(x) ^ (g(x) - 1), f(x) ^ g(x) * log f(x)+ -- |f(x)| ^ g(x) = g(x) * |f(x)| ^ (g(x) - 2) * f(x), |f(x)| ^ g(x) * log |f(x)| - makeRev (j, (1, f, _, _)) rev = let v = fwd IntMap.! (2*j + 1)- dx = rev IntMap.! j- val = dx !*! (M.map (derivative (toEnum f)) v)- in IntMap.insert (2*j + 1) val rev- makeRev (j, (2, op, _, _)) rev = let l = fwd IntMap.! (2*j + 1)- r = fwd IntMap.! (2*j + 2)- dx = rev IntMap.! j- (dxl, dxr) = diff (toEnum op) dx l r- in IntMap.insert (2*j + 2) dxr $ IntMap.insert (2*j + 1) dxl rev- makeRev (j, _) rev = rev- -}+ -- |f(x)| ^ g(x) = exp (log |f(x)| * g(x))+ -- => |f(x)| ^ (g(x) - 1) * g(x)+ -- => |f(x)| ^ g(x) * log |f(x)| * 1 - -- dx is the current derivative so far- -- fx is the evaluation of the left branch- -- gx is the evaluation of the right branch- --- -- this should return a tuple, where the left element is- -- dx * d op(f(x), g(x)) / d f(x) and- -- the right branch dx * d op (f(x), g(x)) / d g(x)- arr1 !**! arr2 = M.zipWith (**) arr1 arr2+ fixNaN x | isNaN x = 0+ | otherwise = x - diff Add dx fx gy = (dx, dx)- diff Sub dx fx gy = (dx, negate dx)- diff Mul dx fx gy = (dx * gy, dx * fx)- diff Div dx fx gy = (dx / gy, dx * (negate fx / (gy * gy)))- diff Power dx fx gy = let dxl = dx * (fx ** (gy-1))- dv2 = fx * log fx- in (dxl * gy, dxl * dv2)- diff PowerAbs dx fx gy = let dxl = (gy * fx) * (fx ** abs (gy - 2))- dxr = (log (abs fx)) * (fx ** abs gy)- in (dxl * dx, dxr * dx)+ diff :: Op -> Double -> Double -> Double -> (Double, Double)+ diff Add dx fx gy = (dx, dx)+ diff Sub dx fx gy = (dx, negate dx)+ diff Mul dx fx gy = (dx * gy, dx * fx)+ diff Div dx fx gy = (dx / gy, dx * (negate fx / (gy * gy)))+ --diff Power dx fx gy = (fixNaN $ dx * ((fx+eps)**gy - fx**gy)/eps, fixNaN $ dx * (fx**(gy+eps) - fx**gy)/eps)+ --diff PowerAbs dx fx gy = (fixNaN $ dx * (abs (fx+eps)**gy - abs fx**gy)/eps, fixNaN $ dx * (abs fx**(gy+eps) - abs fx**gy)/eps)+ {--}+ diff Power 0 _ _ = (0, 0)+ diff Power dx 0 0 = (0, 0)+ diff Power dx fx 0 = (0, fixNaN $ dx * log fx)+ diff Power dx 0 gy = (fixNaN $ dx * gy * if gy < 1 then eps ** (gy - 1) else 0+ , 0) --dx * fx ** gy * log fx)+ diff Power dx fx gy = (fixNaN $ dx * gy * fx ** (gy - 1), fixNaN $ dx * fx ** gy * log fx)++ diff PowerAbs 0 fx gy = (0, 0)+ diff PowerAbs 0 0 0 = (0, 0)+ diff PowerAbs dx fx 0 = (0, fixNaN $ dx * log (abs fx))+ diff PowerAbs dx 0 gy = (0, fixNaN $ dx * if gy < 0 then eps ** gy else 0)+ diff PowerAbs dx fx gy = (fixNaN $ dx * gy * fx * abs fx ** (gy - 2), fixNaN $ dx * abs fx ** gy * log (abs fx))+ {--} diff AQ dx fx gy = let dxl = recip ((sqrt . (+1)) (gy * gy)) dxy = fx * gy * (dxl^3) -- / (sqrt (gy*gy + 1)) in (dxl * dx, dxy * dx)- {-- diff Mul dx fx gy = (dx !*! gy, dx !*! fx)- diff Div dx fx gy = (dx !/! gy, dx !*! (M.map negate fx !/! (gy !*! gy)))- diff Power dx fx gy = let dxl = dx !*! (fx !**! (M.map (subtract 1) gy))- dv2 = fx !*! M.map log fx- in (dxl !*! gy, dxl !*! dv2)- diff PowerAbs dx fx gy = let dxl = (gy !*! fx) !*! (fx !**! M.map abs (M.map (subtract 2) gy))- dxr = (M.map log (M.map abs fx)) !*! (fx !**! M.map abs gy)- in (dxl !*! dx, dxr !*! dx)- diff AQ dx fx gy = let dxl = M.map recip (M.map (sqrt . (+1)) (gy !*! gy))- dxy = fx !*! gy !*! (M.map (^3) dxl) -- / (sqrt (gy*gy + 1))- in (dxl !*! dx, dxy !*! dx)- -} - -- once we reach a leaf with a parameter, we return a singleton- -- with that derivative upwards until the root- combine :: MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> Array S Ix1 Double -> IO ()- combine partial jacob err = forM_ t makeJacob+ {-# INLINE diff #-}++ combine :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO ()+ combine (lo, hi) partial jacob = myForM_ t makeJacob where- makeJacob (j, (0, 1, ix, _)) = do let j' = j2ix IntMap.! j- addI a b acc = do let v1 = err M.! a- v2 <- UMA.unsafeRead partial (a :. b)- pure (v1*v2 + acc)- acc <- foldM (\a i -> addI i j' a) 0 [0..m-1]+ makeJacob (j, (0, 1, ix, _)) = do val <- UMA.unsafeRead jacob ix+ let j' = j2ix IntMap.! j+ addI a b acc = do v2 <- UMA.unsafeRead partial (b :. a)+ pure (v2 + acc)+ acc <- foldM (\a i -> addI i j' a) val [lo..hi-1] UMA.unsafeWrite jacob ix acc makeJacob _ = pure ()- {-- combine :: IntMap.IntMap (Array D Ix1 Double) -> MArray (PrimState IO) S Ix1 Double -> Array D Ix1 Double -> IO ()- combine partial jacob err = forM_ (IntMap.toAscList t) makeJacob- where- makeJacob (j, (0, 1, ix, _)) = do v <- dotM (partial IntMap.! j) err- UMA.unsafeWrite jacob ix v- makeJacob _ = pure ()- -}+ {-# INLINE combine #-} -- | The function `forwardModeUnique` calculates the numerical gradient of the tree and evaluates the tree at the same time. It assumes that each parameter has a unique occurrence in the expression. This should be significantly faster than `forwardMode`. forwardModeUniqueJac :: SRMatrix -> PVector -> Fix SRTree -> [PVector]
src/Algorithm/SRTree/ConfidenceIntervals.hs view
@@ -25,7 +25,7 @@ import Data.SRTree.Recursion ( cata ) import Algorithm.SRTree.Likelihoods import Algorithm.SRTree.Opt- ( minimizeNLLNonUnique, minimizeNLLWithFixedParam )+ ( minimizeNLL, minimizeNLLWithFixedParam ) import Data.List ( sortOn, nubBy ) import Data.Maybe ( fromMaybe ) import Algorithm.SRTree.NonlinearOpt@@ -179,8 +179,8 @@ Right vr -> Right $ evalOp op vl vr -- calculate the profile likelihood of every parameter -getAllProfiles :: PType -> Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> PVector -> [CI] -> Double -> [ProfileT]-getAllProfiles ptype dist mSErr xss ys tree theta stdErr estCIs alpha = reverse (getAll 0 [])+getAllProfiles :: PType -> Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> PVector -> [CI] -> Double -> [ProfileT]+getAllProfiles ptype dist mYerr xss ys tree theta stdErr estCIs alpha = reverse (getAll 0 []) where (A.Sz k) = A.size theta (A.Sz n) = A.size ys@@ -188,18 +188,18 @@ tau_max' = sqrt $ quantile (fDistribution k (n - k)) (1 - alpha) profFun ix = case ptype of- Bates -> getProfile dist mSErr xss ys tree theta (stdErr A.! ix) tau_max ix- ODE -> getProfileODE dist mSErr xss ys tree theta (stdErr A.! ix) (estCIs !! ix) tau_max ix- Constrained -> getProfileCnstr dist mSErr xss ys tree theta (stdErr A.! ix) tau_max' ix+ Bates -> getProfile dist mYerr xss ys tree theta (stdErr A.! ix) tau_max ix+ ODE -> getProfileODE dist mYerr xss ys tree theta (stdErr A.! ix) (estCIs !! ix) tau_max ix+ Constrained -> getProfileCnstr dist mYerr xss ys tree theta (stdErr A.! ix) tau_max' ix getAll ix acc | ix == k = acc | otherwise = case profFun ix of- Left t -> getAllProfiles ptype dist mSErr xss ys tree t stdErr estCIs alpha+ Left t -> getAllProfiles ptype dist mYerr xss ys tree t stdErr estCIs alpha Right p -> getAll (ix + 1) (p : acc) -- calculates the profile likelihood of a single parameter getProfile :: Distribution- -> Maybe Double+ -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree@@ -208,7 +208,7 @@ -> Double -> Int -> Either PVector ProfileT-getProfile dist mSErr xss ys tree theta stdErr_i tau_max ix+getProfile dist mYerr xss ys tree theta stdErr_i tau_max ix | stdErr_i == 0.0 = pure $ ProfileT (A.fromList compMode [-tau_max, tau_max]) (A.fromLists' compMode [theta', theta']) (theta A.! ix) (const (theta A.! ix)) (const tau_max) | otherwise = do negDelta <- go kmax (-stdErr_i / 8) 0 1 mempty@@ -220,10 +220,10 @@ theta' = A.toList theta p0 = ([0], [theta_opt]) kmax = 300- nll_opt = nll dist mSErr xss ys tree theta_opt- theta_opt = fst $ minimizeNLLNonUnique dist mSErr 100 xss ys tree theta+ nll_opt = nll dist mYerr xss ys tree theta_opt+ (theta_opt, _, _) = minimizeNLL dist mYerr 100 xss ys tree theta optTh = theta_opt A.! ix- minimizer = minimizeNLLWithFixedParam dist mSErr 100 xss ys tree ix+ minimizer = minimizeNLLWithFixedParam dist mYerr 100 xss ys tree ix -- after k iterations, interpolates to the endpoint go 0 delta _ _ acc = Right acc@@ -236,10 +236,10 @@ t_delta = (theta_opt A.! ix) + delta * (t + inv_slope) theta_delta = updateS theta_opt [(ix, t_delta)] theta_t = minimizer theta_delta- zv = A.computeAs A.S (snd $ gradNLL dist mSErr xss ys tree theta_t) A.! ix- zvs = snd $ gradNLL dist mSErr xss ys tree theta_t+ zv = A.computeAs A.S (snd $ gradNLL dist mYerr xss ys tree theta_t) A.! ix+ zvs = snd $ gradNLL dist mYerr xss ys tree theta_t inv_slope' = min 4.0 . max 0.0625 . abs $ (tau / (stdErr_i * zv))- nll_cond = nll dist mSErr xss ys tree theta_t+ nll_cond = nll dist mYerr xss ys tree theta_t acc' = if nll_cond == nll_opt || ( (not.null) taus && tau == head taus ) || isNaN tau then acc else (tau:taus, theta_t:thetas)@@ -248,7 +248,7 @@ -- Based on https://insysbio.github.io/LikelihoodProfiler.jl/latest/ -- Borisov, Ivan, and Evgeny Metelkin. "Confidence intervals by constrained optimization—An algorithm and software package for practical identifiability analysis in systems biology." PLOS Computational Biology 16.12 (2020): e1008495. getProfileCnstr :: Distribution- -> Maybe Double+ -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree@@ -256,7 +256,7 @@ -> Double -> Double -> Int -> Either PVector ProfileT-getProfileCnstr dist mSErr xss ys tree theta stdErr_i tau_max ix+getProfileCnstr dist mYerr xss ys tree theta stdErr_i tau_max ix | stdErr_i == 0.0 = pure $ ProfileT taus thetas theta_i (const theta_i) (const tau_max) | otherwise = pure $ ProfileT taus thetas theta_i tau2theta (const tau_max) where@@ -264,24 +264,24 @@ theta' = A.toList theta thetas = A.fromLists' compMode [theta', theta'] theta_i = theta A.! ix- getPoint = getEndPoint dist mSErr xss ys tree theta tau_max ix+ getPoint = getEndPoint dist mYerr xss ys tree theta tau_max ix leftPt = getPoint True rightPt = getPoint False tau2theta tau = if tau < 0 then leftPt else rightPt -getEndPoint :: Distribution -> Maybe Double -> A.Array A.S Ix2 Double -> A.Array A.S A.Ix1 Double -> Fix SRTree -> A.Array A.S A.Ix1 Double -> Double -> Int -> Bool -> Double-getEndPoint dist mSErr xss ys tree theta tau_max ix isLeft =+getEndPoint :: Distribution -> Maybe PVector -> A.Array A.S Ix2 Double -> A.Array A.S A.Ix1 Double -> Fix SRTree -> A.Array A.S A.Ix1 Double -> Double -> Int -> Bool -> Double+getEndPoint dist mYerr xss ys tree theta tau_max ix isLeft = case minimizeAugLag problem (A.toStorableVector theta_opt) of Right sol -> solutionParams sol VS.! ix Left e -> traceShow e $ theta_opt A.! ix where (A.Sz1 n) = A.size theta - theta_opt = fst $ minimizeNLLNonUnique dist mSErr 100 xss ys tree theta- nll_opt = nll dist mSErr xss ys tree theta_opt+ (theta_opt, _, _) = minimizeNLL dist mYerr 100 xss ys tree theta+ nll_opt = nll dist mYerr xss ys tree theta_opt loss_crit = nll_opt + tau_max - loss = subtract loss_crit . nll dist mSErr xss ys tree . A.fromStorableVector compMode+ loss = subtract loss_crit . nll dist mYerr xss ys tree . A.fromStorableVector compMode obj = (if isLeft then id else negate) . (VS.! ix) stop = ObjectiveRelativeTolerance 1e-4 :| []@@ -296,7 +296,7 @@ -- Jian-Shen Chen & Robert I Jennrich (2002) Simple Accurate Approximation of Likelihood Profiles, -- Journal of Computational and Graphical Statistics, 11:3, 714-732, DOI: 10.1198/106186002493 getProfileODE :: Distribution- -> Maybe Double+ -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree@@ -306,27 +306,27 @@ -> Double -> Int -> Either PVector ProfileT-getProfileODE dist mSErr xss ys tree theta stdErr_i estCI tau_max ix+getProfileODE dist mYerr xss ys tree theta stdErr_i estCI tau_max ix | stdErr_i == 0.0 = pure dflt | otherwise = let (A.fromList compMode -> taus, A.fromLists' compMode . map A.toList -> thetas) = solLeft <> ([0], [theta_opt]) <> solRight (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix in pure $ ProfileT taus thetas optTh tau2theta theta2tau where dflt = ProfileT (A.fromList compMode [-tau_max, tau_max]) (A.fromLists' compMode [theta', theta']) (theta A.! ix) (const (theta A.! ix)) (const tau_max)- minimizer = fst . minimizeNLLNonUnique dist mSErr 100 xss ys tree- grader = snd . gradNLLNonUnique dist mSErr xss ys tree+ minimizer = (\(x, _, _) -> x) . minimizeNLL dist mYerr 100 xss ys tree+ grader = snd . gradNLL dist mYerr xss ys tree theta_opt = minimizer theta theta' = A.toList theta- nll_opt = nll dist mSErr xss ys tree theta_opt+ nll_opt = nll dist mYerr xss ys tree theta_opt optTh = theta_opt A.! ix p' = p+1 (A.Sz1 p) = A.size theta- sErr = fromMaybe 1 mSErr- getHess = hessianNLL dist mSErr xss ys tree+ --sErr = fromMaybe 1 mSErr+ getHess = hessianNLL dist mYerr xss ys tree odeFun gamma _ u = let grad = grader u- w = hessianNLL dist mSErr xss ys tree u+ w = hessianNLL dist mYerr xss ys tree u m = A.makeArray compMode (A.Sz (p' :. p')) (\ (i :. j) -> if | i<p && j<p -> w A.! (i :. j) | i==ix -> 1@@ -343,7 +343,7 @@ where f = if sig==1 then id else reverse solRight = scanOn 1 tsHi solLeft = scanOn (-1) tsLo- calcTau s t = let nll_i = nll dist mSErr xss ys tree $ snd t+ calcTau s t = let nll_i = nll dist mYerr xss ys tree $ snd t z = signum ((snd t A.! ix) - optTh) * sqrt (2 * nll_i - 2 * nll_opt) in if z == 0 || isNaN z then ([], []) else ([z], [snd t]) @@ -361,8 +361,8 @@ {-# INLINE rk #-} -- tau0, tau1 theta0, thetaX = tau1 theta0 / tau0-getStatsFromModel :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> BasicStats-getStatsFromModel dist mSErr xss ys tree theta = MkStats cov corr stdErr+getStatsFromModel :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> BasicStats+getStatsFromModel dist mYerr xss ys tree theta = MkStats cov corr stdErr where (A.Sz1 k) = A.size theta (A.Sz1 n) = A.size ys@@ -373,7 +373,7 @@ -- only for gaussian sErr = sqrt $ ssr / fromIntegral (n - k) - hess = hessianNLL dist mSErr xss ys tree theta+ hess = hessianNLL dist mYerr xss ys tree theta -- cov = catch (unsafePerformIO (invChol hess)) (\e -> trace "cov NegDef" $ pure ident) fexcept :: (A.PrimMonad m, A.MonadThrow m, A.MonadIO m) => A.SomeException -> m SRMatrix fexcept e = trace "cov NegDef" $ pure ident
src/Algorithm/SRTree/Likelihoods.hs view
@@ -1,4 +1,6 @@ {-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE TypeApplications #-}+ ----------------------------------------------------------------------------- -- | -- Module : Algorithm.SRTree.Likelihoods @@ -21,27 +23,39 @@ , r2 , nll , predict+ , buildNLL , gradNLL , gradNLLArr- , gradNLLNonUnique+ , gradNLLGraph , fisherNLL , getSErr , hessianNLL+ , tree2arr ) where -import Algorithm.SRTree.AD ( forwardMode, reverseModeUnique, reverseModeUniqueArr ) -- ( reverseModeUnique )+import Algorithm.SRTree.AD ( reverseModeArr, reverseModeGraph ) 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 import Data.Maybe (fromMaybe)-import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams)-import Data.SRTree.Derivative (deriveByParam)-import Data.SRTree.Eval (PVector, SRMatrix, SRVector, compMode, evalTree)+import Data.SRTree+import Data.SRTree.Recursion ( cata, accu )+import Data.SRTree.Derivative (deriveByParam, deriveByVar, derivative)+import Data.SRTree.Eval import qualified Data.IntMap.Strict as IntMap+import qualified Data.Vector.Storable as VS+import GHC.IO (unsafePerformIO)+import Data.Maybe +import Debug.Trace+import Data.SRTree.Print+ -- | Supported distributions for negative log-likelihood-data Distribution = Gaussian | Bernoulli | Poisson- deriving (Show, Read, Enum, Bounded)+-- 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+ deriving (Show, Read, Enum, Bounded, Eq) -- | Sum-of-square errors or Sum-of-square residues sse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double@@ -52,6 +66,14 @@ yhat = evalTree xss theta tree err = M.sum $ (delay ys - yhat) ^ (2 :: Int) +sseError :: SRMatrix -> PVector -> PVector -> Fix SRTree -> PVector -> Double+sseError xss ys yErr tree theta = err+ where+ (Sz m) = M.size ys+ cmp = getComp xss+ yhat = evalTree xss theta tree+ err = M.sum $ ((delay ys - yhat) ^ (2 :: Int) / (delay yErr))+ -- | Total Sum-of-squares sseTot :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double sseTot xss ys tree theta = err@@ -92,32 +114,46 @@ {-# inline negSum #-} -- | Negative log-likelihood-nll :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> Double+nll :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> Double --- | Gaussian distribution-nll Gaussian msErr xss ys t theta = 0.5*(ssr/s2 + m*log (2*pi*s2))+-- | Mean Squared error (not a distribution)+nll MSE _ xss ys t theta = mse xss ys t theta++-- | 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."+ | otherwise = 0.5*(sse xss ys t theta / s + m*log (2*pi*s)) where+ s = theta M.! (p' - 1) (Sz m') = M.size ys (Sz p') = M.size theta- m = fromIntegral m' - p = fromIntegral p'- ssr = sse xss ys t theta- mse' = mse xss ys t theta- est = sqrt (m - p) -- $ ssr / (m - p)- sErr = getSErr Gaussian est msErr- s2 = sErr ^ 2+ nParams = countParams t+ m = fromIntegral m'+ p = fromIntegral p' +-- | Gaussian with heteroscedasticity, it needs a valid mYerr+nll HGaussian mYerr xss ys t theta =+ case mYerr of+ Nothing -> error "For HGaussian, you must provide the measured error for the target variable."+ Just yErr -> 0.5*(sseError xss ys yErr t theta + M.sum (M.map (log . (2*) . (pi*)) yErr))+ where+ (Sz m') = M.size ys+ (Sz p') = M.size theta+ m = fromIntegral m'+ p = fromIntegral p'+ -- | Bernoulli distribution of f(x; theta) is, given phi = 1 / (1 + exp (-f(x; theta))), -- y log phi + (1-y) log (1 - phi), assuming y \in {0,1} nll Bernoulli _ xss ys tree theta | notValid ys = error "For Bernoulli distribution the output must be either 0 or 1."- | otherwise = negate . M.sum $ delay ys * yhat - log (1 + exp yhat)+ | otherwise = negate . M.sum $ delay ys * yhat - log (M.map (1+) $ exp yhat) where (Sz m) = M.size ys yhat = evalTree xss theta tree notValid = M.any (\x -> x /= 0 && x /= 1) -nll Poisson _ xss ys tree theta +nll Poisson _ xss ys tree theta | notValid ys = error "For Poisson distribution the output must be non-negative." -- | M.any isNaN yhat = error $ "NaN predictions " <> show theta | otherwise = negate . M.sum $ ys' * yhat - ys' * log ys' - exp yhat@@ -126,138 +162,267 @@ yhat = evalTree xss theta tree notValid = M.any (<0) -nll' :: Distribution -> Double -> SRVector -> SRVector -> Double-nll' Gaussian sErr yhat ys = 0.5*(ssr/s2 + m*log (2*pi*s2))- where - (Sz m') = M.size ys - m = fromIntegral m' - ssr = M.sum $ (ys - yhat)^2- s2 = sErr ^ 2-nll' Bernoulli _ yhat ys = negate . M.sum $ ys * yhat - log (1 + exp yhat)-nll' Poisson _ yhat ys = negate . M.sum $ ys * yhat - ys * log ys - exp yhat-{-# INLINE nll' #-}+nll ROXY mYerr xss ys tree theta+ | isNothing mYerr = error "Can't calculate ROXY nll without x,y-errors."+ | p < num_params + 3 = error "We need 3 additional parameters for ROXY."+ | n /= 1 && n/=5 = error "For ROXY dataset must contain a single variable, or 1 variable + 4 cached data."+ | otherwise = if isNaN negLL then (1.0/0.0) else negLL+ where+ (Sz p') = M.size theta+ (Sz2 m n) = M.size xss+ p = fromIntegral p'+ num_params = countParams tree + x0 = xss <! 0+ logX = xss <! 1+ logY = xss <! 2+ logXErr = xss <! 3+ logYErr = xss <! 4+++ yErr = fromJust mYerr+ one = M.replicate compMode (Sz m) 1+ zero = M.replicate compMode (Sz m) 0++ (sig, mu_gauss, w_gauss) = (theta ! num_params, theta ! (num_params + 1), theta ! (num_params + 2))++ applyDer :: Op -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double+ applyDer Add l dl r dr = dl+dr+ applyDer Sub l dl r dr = dl-dr+ applyDer Mul l dl r dr = l*dr + r*dl+ applyDer Div l dl r dr = (dl*r - dr*l) / (r^2)+ applyDer Power l dl r dr = l ** (r.-1) * (r*dl + l * log l * dr)+ applyDer PowerAbs l dl r dr = (abs l ** r) * (dr * log (abs l) + r * dl / l)+ applyDer AQ l dl r dr = ((1 +. r*r) * dl - l * r * dr) / M.map (**1.5) (1 +. r*r)++ (yhat, grad) = cata alg tree+ where+ alg (Var ix) = (x0, one)+ alg (Param ix) = (M.replicate compMode (Sz m) (theta M.! ix), zero)+ alg (Const x) = (M.replicate compMode (Sz m) x, zero)+ alg (Uni f (val, der)) = (M.map (evalFun f) val, M.map (derivative f) val * der)+ alg (Bin op (valL, derL) (valR, derR)) = (M.zipWith (evalOp op) valL valR, applyDer op valL derL valR derR)++ f = M.map (logBase 10) (abs yhat)+ fprime = grad / (log 10 *. yhat) * x0 .* log 10++ -- nll+ w_gauss2 = w_gauss ^ 2+ s2 = delay $ logYErr .+ sig^2+ den = fprime ^ 2 .* w_gauss2 * logXErr + s2 * (w_gauss2 +. logXErr)++ neglogP = log (2 * pi)+ +. log den+ + (w_gauss2 *. (f - logY) * (f - logY)+ + logXErr * (fprime * (mu_gauss -. logX) + f - logY)^2+ + s2 * (logX .- mu_gauss)^2) / den+ negLL = 0.5 * M.sum neglogP++-- WARNING: pass tree with parameters+-- TODO: handle error similar to ROXY+buildNLL MSE m tree = ((tree - var (-1)) ** 2) / constv m+buildNLL Gaussian m tree = (square(tree - var (-1)) / square (param p)) + log ((square (param p)))+ where+ square = Fix . Uni Square+ p = countParams tree+buildNLL HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2))+buildNLL Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree+buildNLL Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree+buildNLL ROXY m tree = neglogP+ where+ p = countParams tree+ f = log (abs tree) / log 10+ fprime = deriveByVar 0 tree / (log 10 * tree) * var 0 * log 10+ logX = var 1+ logY = var 2+ logXErr = var 3+ logYErr = var 4+ sig = param p+ mu_gauss = param (p+1)+ w_gauss = param (p+2)+ w_gauss2 = w_gauss ** 2+ s2 = logYErr + sig ** 2+ den = fprime ** 2 * w_gauss2 * logXErr + s2 * (w_gauss2 + logXErr)+ neglogP = log (2*pi)+ + log den+ + ( w_gauss2 * (f - logY) * (f - logY)+ + logXErr * (fprime *(mu_gauss - logX) + f - logY)**2+ + s2 * (logX - mu_gauss) ** 2+ ) / den+ -- | Prediction for different distributions predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector+predict MSE 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+predict ROXY tree theta xss = evalTree xss theta tree -- | Gradient of the negative log-likelihood-gradNLL :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector)-gradNLL Gaussian msErr xss ys tree theta =- (nll' Gaussian sErr yhat ys', delay grad ./ (sErr * sErr))+gradNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector)+gradNLL dist mYerr xss ys tree theta = (f, delay grad) -- gradNLLArr dist xss ys mYerr treeArr j2ix (toStorableVector theta) where- (Sz m) = M.size ys- (Sz p) = M.size theta- ys' = delay ys- (yhat, grad) = reverseModeUnique xss theta ys' id tree- -- err = yhat - delay ys- ssr = sse xss ys tree theta- est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p)- sErr = getSErr Gaussian est msErr+ grad :: PVector+ grad = M.fromList M.Seq [finitediff ix | ix <- [0..p-1]]+ (Sz p) = M.size theta -gradNLL Bernoulli _ xss (delay -> ys) tree theta- | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."- | otherwise = (nll' Bernoulli 1.0 yhat ys, delay grad)- where- (yhat, grad) = reverseModeUnique xss theta ys logistic tree- grad' = M.map nanTo0 grad- --err = logistic yhat - ys- nanTo0 x = if isNaN x then 0 else x+ disturb :: Int -> PVector+ disturb ix = M.fromList M.Seq $ Prelude.zipWith (\iy v -> if iy==ix then (v+eps) else v) [0..] (M.toList theta)+ eps :: Double+ 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+ 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..] -gradNLL Poisson _ xss (delay -> ys) tree theta- | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."- -- | M.any isNaN grad = error $ "NaN gradient " <> show grad- | otherwise = (nll' Poisson 1.0 yhat ys, delay grad)+ {-+ -- EXAMPLE OF FINITE DIFFERENCE+ -- Implement for debugging+gradNLL ROXY mXerr mYerr xss ys tree theta =+ (f, delay grad) where- (yhat, grad) = reverseModeUnique xss theta ys exp tree- --err = exp yhat - ys+ (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 #-}+ -- | Gradient of the negative log-likelihood---Array B Ix1 (Int, Int, Int, Double)-gradNLLArr :: Distribution -> Maybe Double -> SRMatrix -> PVector -> [(Int,(Int, Int, Int, Double))] -> IntMap.IntMap Int -> PVector -> (Double, SRVector)-gradNLLArr Gaussian msErr xss ys tree j2ix theta =- (nll' Gaussian sErr yhat ys', delay grad ./ (sErr * sErr))+gradNLLArr MSE xss ys mYerr tree j2ix theta =+ (M.sum yhat, delay grad') where- (Sz m) = M.size ys- (Sz p) = M.size theta- ys' = delay ys- (yhat, grad) = reverseModeUniqueArr xss theta ys' id tree j2ix- -- err = yhat - delay ys- --ssr = sse xss ys tree theta- est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p)- sErr = getSErr Gaussian est msErr--gradNLLArr Bernoulli _ xss (delay -> ys) tree j2ix theta+ (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+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr Bernoulli xss ys mYerr tree j2ix theta | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."- | otherwise = (nll' Bernoulli 1.0 yhat ys, delay grad)+ | otherwise = (M.sum yhat, delay grad') where- (yhat, grad) = reverseModeUniqueArr xss theta ys logistic tree j2ix+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix grad' = M.map nanTo0 grad- --err = logistic yhat - ys- nanTo0 x = if isNaN x then 0 else x--gradNLLArr Poisson _ xss (delay -> ys) tree j2ix theta+gradNLLArr Poisson xss ys mYerr tree j2ix theta | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."- -- | M.any isNaN grad = error $ "NaN gradient " <> show grad- | otherwise = (nll' Poisson 1.0 yhat ys, delay grad)+ | otherwise = (M.sum yhat, delay grad') where- (yhat, grad) = reverseModeUniqueArr xss theta ys exp tree j2ix- --err = exp yhat - ys+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr ROXY xss ys mYerr tree j2ix theta =+ ((*0.5) $ M.sum yhat, M.map (*(0.5)) $ delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad -- | Gradient of the negative log-likelihood-gradNLLNonUnique :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector)-gradNLLNonUnique Gaussian msErr xss ys tree theta =- (nll' Gaussian sErr yhat ys', delay grad ./ (sErr * sErr))+gradNLLGraph MSE xss ys mYerr tree theta =+ (M.sum yhat, grad') where- (Sz m) = M.size ys- (Sz p) = M.size theta- ys' = delay ys- (yhat, grad) = forwardMode xss theta err tree- err = predict Gaussian tree theta xss - ys'- ssr = sse xss ys tree theta- est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p)- sErr = getSErr Gaussian est msErr--gradNLLNonUnique Bernoulli _ xss (delay -> ys) tree theta+ (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+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph Bernoulli xss ys mYerr tree theta | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."- | otherwise = (nll' Bernoulli 1.0 yhat ys, delay grad)+ | otherwise = (M.sum yhat, grad') where- (yhat, grad) = forwardMode xss theta err tree- grad' = M.map nanTo0 grad- err = predict Bernoulli tree theta xss - delay ys- nanTo0 x = if isNaN x then 0 else x--gradNLLNonUnique Poisson _ xss (delay -> ys) tree theta+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph Poisson xss ys mYerr tree theta | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."- -- | M.any isNaN grad = error $ "NaN gradient " <> show grad- | otherwise = (nll' Poisson 1.0 yhat ys, delay grad)+ | otherwise = (M.sum yhat, grad') where- (yhat, grad) = forwardMode xss theta err tree- err = predict Poisson tree theta xss - delay ys+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph ROXY xss ys mYerr tree theta =+ ((*0.5) $ M.sum yhat, VS.map (*(0.5)) $ grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad -- | Fisher information of negative log-likelihood-fisherNLL :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector-fisherNLL dist msErr xss ys tree theta = makeArray cmp (Sz p) build+fisherNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector+fisherNLL ROXY mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff where+ cmp = getComp xss+ (Sz m) = M.size ys+ (Sz p) = M.size theta+ f = nll ROXY mYerr xss ys tree theta+ eps = 1e-6+ finiteDiff ix = unsafePerformIO $ do+ theta' <- Mut.thaw theta+ v <- Mut.readM theta' ix+ Mut.writeM theta' ix (v + eps)+ thetaPlus <- Mut.freezeS theta'+ Mut.writeM theta' ix (v - eps)+ thetaMinus <- Mut.freezeS theta'+ let fPlus = nll ROXY mYerr xss ys tree thetaPlus+ fMinus = nll ROXY mYerr xss ys tree thetaMinus+ pure $ (fPlus + fMinus - 2*f)/(eps*eps)+fisherNLL Gaussian mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff+ where+ cmp = getComp xss+ (Sz m) = M.size ys+ (Sz p) = M.size theta+ f = nll Gaussian mYerr xss ys tree theta+ eps = 1e-6+ finiteDiff ix = unsafePerformIO $ do+ theta' <- Mut.thaw theta+ v <- Mut.readM theta' ix+ Mut.writeM theta' ix (v + eps)+ thetaPlus <- Mut.freezeS theta'+ Mut.writeM theta' ix (v - eps)+ thetaMinus <- Mut.freezeS theta'+ let fPlus = nll Gaussian mYerr xss ys tree thetaPlus+ fMinus = nll Gaussian mYerr xss ys tree thetaMinus+ pure $ (fPlus + fMinus - 2*f)/(eps*eps)+fisherNLL dist mYerr xss ys tree theta = makeArray cmp (Sz p) build+ where build ix = let dtdix = deriveByParam ix t' d2tdix2 = deriveByParam ix dtdix f' = eval dtdix f'' = eval d2tdix2 - in (/sErr^2) . M.sum $ phi' * f'^2 - res * f''+ in M.sum $ phi' * f'^2 - res * f''+ --case dist of+ -- Gaussian -> M.sum . (/delay (theta M.! (p-1))) $ phi' * f'^2 - res * f''+ -- _ -> M.sum $ phi' * f'^2 - res * f'' cmp = getComp xss (Sz m) = M.size ys (Sz p) = M.size theta t' = fst $ floatConstsToParam tree eval = evalTree xss theta- ssr = sse xss ys tree theta- sErr = getSErr dist est msErr- est = sqrt $ fromIntegral (m-p) -- $ ssr / fromIntegral (m - p) yhat = eval t' res = delay ys - phi+ yErr = case mYerr of+ Nothing -> M.replicate (getComp xss) (Sz m) est+ Just e -> e+ est = fromIntegral (m - p) (phi, phi') = case dist of+ MSE -> (yhat, M.replicate compMode (Sz m) 1) Gaussian -> (yhat, M.replicate compMode (Sz m) 1) Bernoulli -> (logistic yhat, phi*(M.replicate compMode (Sz m) 1 - phi)) Poisson -> (exp yhat, phi)@@ -266,8 +431,9 @@ -- -- Note, though the Fisher is just the diagonal of the return of this function -- it is better to keep them as different functions for efficiency-hessianNLL :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRMatrix-hessianNLL dist msErr xss ys tree theta = makeArray cmp (Sz (p :. p)) build +hessianNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRMatrix+hessianNLL ROXY mYerr xss ys tree theta = undefined+hessianNLL dist mYerr xss ys tree theta = makeArray cmp (Sz (p :. p)) build where build (ix :. iy) = let dtdix = deriveByParam ix t' dtdiy = deriveByParam iy t' @@ -275,15 +441,19 @@ fx = eval dtdix fy = eval dtdiy fxy = eval d2tdixy - in (/sErr^2) . M.sum $ phi' * fx * fy - res * fxy+ in case dist of+ Gaussian -> M.sum . (/delay yErr) $ phi' * fx * fy - res * fxy+ _ -> M.sum $ phi' * fx * fy - res * fxy+ cmp = getComp xss (Sz m) = M.size ys (Sz p) = M.size theta t' = tree -- relabelParams tree -- $ floatConstsToParam tree eval = evalTree xss theta- ssr = sse xss ys tree theta- sErr = getSErr dist est msErr- est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p)+ yErr = case mYerr of+ Nothing -> M.replicate compMode (Sz m) est+ Just e -> e+ est = fromIntegral (m - p) yhat = eval t' res = delay ys - phi @@ -292,3 +462,27 @@ Bernoulli -> (logistic yhat, phi*(M.replicate cmp (Sz m) 1 - phi)) Poisson -> (exp yhat, phi) +tree2arr :: Fix SRTree -> IntMap.IntMap (Int, Int, Int, Double)+tree2arr tree = IntMap.fromList listTree+ where+ height = cata alg+ where+ alg (Var ix) = 1+ alg (Const x) = 1+ alg (Param ix) = 1+ alg (Uni _ t) = 1 + t+ alg (Bin _ l r) = 1 + max l r+ listTree = accu indexer convert tree 0++ indexer (Var ix) iy = Var ix+ indexer (Const x) iy = Const x+ indexer (Param ix) iy = Param ix+ indexer (Bin op l r) iy = Bin op (l, 2*iy+1) (r, 2*iy+2)+ indexer (Uni f t) iy = Uni f (t, 2*iy+1)++ convert (Var ix) iy = [(iy, (0, 0, ix, -1))]+ convert (Const x) iy = [(iy, (0, 2, -1, x))]+ convert (Param ix) iy = [(iy, (0, 1, ix, -1))]+ convert (Uni f t) iy = (iy, (1, fromEnum f, -1, -1)) : t+ convert (Bin op l r) iy = (iy, (2, fromEnum op, -1, -1)) : (l <> r)+{-# INLINE tree2arr #-}
src/Algorithm/SRTree/ModelSelection.hs view
@@ -26,26 +26,27 @@ import Data.SRTree.Recursion (cata) import qualified Data.Vector.Storable as VS +import Debug.Trace -- | Bayesian information criterion-bic :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-bic dist mSErr xss ys theta tree = (p + 1) * log n + 2 * nll dist mSErr xss ys tree theta+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 where (A.Sz (fromIntegral -> p)) = A.size theta (A.Sz (fromIntegral -> n)) = A.size ys {-# INLINE bic #-} -- | Akaike information criterion-aic :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-aic dist mSErr xss ys theta tree = 2 * (p + 1) + 2 * nll dist mSErr xss ys tree theta+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 where (A.Sz (fromIntegral -> p)) = A.size theta (A.Sz (fromIntegral -> n)) = A.size ys {-# INLINE aic #-} -- | Evidence -evidence :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-evidence dist mSErr xss ys theta tree = (1 - b) * nll dist mSErr xss ys tree theta - p / 2 * log b+evidence :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+evidence dist mYerr xss ys theta tree = (1 - b) * nll dist mYerr xss ys tree theta - p / 2 * log b where (A.Sz (fromIntegral -> p)) = A.size theta (A.Sz (fromIntegral -> n)) = A.size ys@@ -54,34 +55,34 @@ -- | 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 Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-mdl dist mSErr xss ys theta tree = nll' dist mSErr xss ys theta' tree- + logFunctional tree- + logParameters dist mSErr xss ys theta tree+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 where- fisher = fisherNLL dist mSErr xss ys tree theta+ 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 isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 {-# INLINE mdl #-} -- | MDL Lattice as described in -- Bartlett, Deaglan, Harry Desmond, and Pedro Ferreira. "Priors for symbolic regression." Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 2023.-mdlLatt :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-mdlLatt dist mSErr xss ys theta tree = nll' dist mSErr xss ys theta' tree+mdlLatt :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+mdlLatt dist mYerr xss ys theta tree = nll' dist mYerr xss ys theta' tree + logFunctional tree- + logParametersLatt dist mSErr xss ys theta tree+ + logParametersLatt dist mYerr xss ys theta tree where- fisher = fisherNLL dist mSErr xss ys tree theta+ 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 isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 {-# INLINE mdlLatt #-} -- | same as `mdl` but weighting the functional structure by frequency calculated using a wiki information of -- physics and engineering functions-mdlFreq :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-mdlFreq dist mSErr xss ys theta tree = nll dist mSErr xss ys tree theta+mdlFreq :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+mdlFreq dist mYerr xss ys theta tree = nll dist mYerr xss ys tree theta + logFunctionalFreq tree- + logParameters dist mSErr xss ys theta tree+ + logParameters dist mYerr xss ys theta tree {-# INLINE mdlFreq #-} -- log of the functional complexity@@ -107,11 +108,11 @@ {-# INLINE logFunctionalFreq #-} -- log of the parameters complexity-logParameters :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-logParameters dist mSErr xss ys theta tree = -(p / 2) * log 3 + 0.5 * logFisher + logTheta+logParameters :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+logParameters dist mYerr xss ys theta tree = -(p / 2) * log 3 + 0.5 * logFisher + logTheta where -- p = fromIntegral $ VS.length theta- fisher = fisherNLL dist mSErr xss ys tree theta+ fisher = fisherNLL dist mYerr xss ys tree theta (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (A.toList theta) (A.toList fisher)@@ -123,11 +124,11 @@ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 -- same as above but for the Lattice -logParametersLatt :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-logParametersLatt dist mSErr xss ys theta tree = 0.5 * p * (1 - log 3) + 0.5 * log detFisher+logParametersLatt :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+logParametersLatt dist mYerr xss ys theta tree = 0.5 * p * (1 - log 3) + 0.5 * log detFisher where- fisher = fisherNLL dist mSErr xss ys tree theta- detFisher = det $ hessianNLL dist mSErr xss ys tree theta+ fisher = fisherNLL dist mYerr xss ys tree theta+ detFisher = det $ hessianNLL dist mYerr xss ys tree theta (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (A.toList theta) (A.toList fisher)@@ -139,8 +140,8 @@ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 -- flipped version of nll-nll' :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double-nll' dist mSErr xss ys theta tree = nll dist mSErr xss ys tree theta+nll' :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+nll' dist mYerr xss ys theta tree = nll dist mYerr xss ys tree theta {-# INLINE nll' #-} treeToNat :: Fix SRTree -> Double
src/Algorithm/SRTree/NonlinearOpt.hs view
@@ -644,9 +644,9 @@ solveProblem :: N.Opt -> Vector Double -> IO Solution solveProblem opt x0 = do- (N.Output outret outcost outx) <- N.optimize opt x0+ (N.Output outret outcost outx nevals) <- N.optimize opt x0 if (N.isSuccess outret)- then return $ Solution outcost outx outret+ then return $ Solution outcost outx outret nevals else Ex.throw $ NloptException outret minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution@@ -903,6 +903,8 @@ , solutionParams :: Vector Double -- ^ The parameter vector which -- minimizes the objective , solutionResult :: N.Result -- ^ Why the optimizer stopped++ , nEvals :: Int -- ^ Number of evaluations until stop } deriving (Eq, Show, Read) applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO ()
src/Algorithm/SRTree/Opt.hs view
@@ -18,7 +18,7 @@ import Algorithm.SRTree.NonlinearOpt import Data.Bifunctor (bimap, second) import Data.Massiv.Array-import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams, countNodes)+import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams, countNodes, convertProtectedOps) import Data.SRTree.Eval (evalTree, compMode) import qualified Data.Vector.Storable as VS import qualified Data.IntMap.Strict as IntMap@@ -26,115 +26,79 @@ import Debug.Trace -tree2arr :: Fix SRTree -> IntMap.IntMap (Int, Int, Int, Double)-tree2arr tree = IntMap.fromList listTree- where- height = cata alg- where- alg (Var ix) = 1- alg (Const x) = 1- alg (Param ix) = 1- alg (Uni _ t) = 1 + t- alg (Bin _ l r) = 1 + max l r- listTree = accu indexer convert tree 0 - indexer (Var ix) iy = Var ix- indexer (Const x) iy = Const x- indexer (Param ix) iy = Param ix- indexer (Bin op l r) iy = Bin op (l, 2*iy+1) (r, 2*iy+2)- indexer (Uni f t) iy = Uni f (t, 2*iy+1) - convert (Var ix) iy = [(iy, (0, 0, ix, -1))]- convert (Const x) iy = [(iy, (0, 2, -1, x))]- convert (Param ix) iy = [(iy, (0, 1, ix, -1))]- convert (Uni f t) iy = (iy, (1, fromEnum f, -1, -1)) : t- convert (Bin op l r) iy = (iy, (2, fromEnum op, -1, -1)) : (l <> r)- -- | minimizes the negative log-likelihood of the expression-minimizeNLL :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double)-minimizeNLL dist msErr niter xss ys tree t0- | niter == 0 = (t0, f)- | n == 0 = (t0, f)- | otherwise = (fromStorableVector compMode t_opt, f)+minimizeNLL' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeNLL' alg dist mYerr niter xss ys tree t0+ | niter == 0 = (t0, f, 0)+ | n == 0 = (t0, f, 0)+ | otherwise = (t_opt', nll dist mYerr xss ys tree t_opt', nEvs) where- tree' = relabelParams tree -- $ fst $ floatConstsToParam tree+ tree' = buildNLL dist (fromIntegral m) $ relabelParams $ tree -- convertProtectedOps t0' = toStorableVector t0 treeArr = IntMap.toAscList $ tree2arr tree' j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] (Sz n) = size t0 (Sz m) = size ys- funAndGrad = second (toStorableVector . computeAs S) . gradNLLArr dist msErr xss ys treeArr j2ix . fromStorableVector compMode- (f, _) = gradNLLArr dist msErr xss ys treeArr j2ix t0 -- if there's no parameter or no iterations-- algorithm = LBFGS funAndGrad Nothing- stop = ObjectiveRelativeTolerance 1e-6 :| [MaximumEvaluations (fromIntegral niter)]- problem = LocalProblem (fromIntegral n) stop algorithm- t_opt = case minimizeLocal problem t0' of- Right sol -> solutionParams sol- Left e -> t0'+ funAndGrad = gradNLLGraph dist xss ys mYerr tree' -- second (toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix --- | minimizes the likelihood assuming repeating parameters in the expression -minimizeNLLNonUnique :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double)-minimizeNLLNonUnique dist msErr niter xss ys tree t0- | niter == 0 = (t0, f)- | n == 0 = (t0, f)- | otherwise = (fromStorableVector compMode t_opt, f)- where- t0' = toStorableVector t0- (Sz n) = size t0- (Sz m) = size ys- funAndGrad = second (toStorableVector . computeAs S) . gradNLLNonUnique dist msErr xss ys tree . fromStorableVector compMode- (f, _) = gradNLLNonUnique dist msErr xss ys tree t0 -- if there's no parameter or no iterations+ (f, _) = gradNLLGraph dist xss ys mYerr tree' t0' -- if there's no parameter or no iterations+ -- gradNLL dist mYerr xss ys tree t0+ --debug1 = gradNLLArr dist msErr xss ys treeArr j2ix t0+ --debug2 = gradNLL dist msErr xss ys tree t0 - algorithm = LBFGS funAndGrad Nothing- stop = ObjectiveRelativeTolerance 1e-5 :| [MaximumEvaluations (fromIntegral niter)]+ algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) -- alg funAndGrad Nothing -- PRAXIS (fst . funAndGrad) [] Nothing -- TNEWTON funAndGrad Nothing+ stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)] problem = LocalProblem (fromIntegral n) stop algorithm- t_opt = case minimizeLocal problem t0' of- Right sol -> solutionParams sol- Left e -> t0'+ (t_opt, nEvs) = case minimizeLocal problem t0' of+ Right sol -> (solutionParams sol, nEvals sol) -- traceShow (">>>>>>>", nEvals sol) $+ Left e -> (t0', 0)+ t_opt' = fromStorableVector compMode t_opt+ debugGrad t = let g1 = gradNLL dist mYerr xss ys tree . fromStorableVector compMode $ t+ g2 = gradNLLArr dist xss ys mYerr treeArr j2ix t+ g3 = gradNLLGraph dist xss ys mYerr tree' t+ in traceShow (t, g1, g2, g3) $ g3 -- second (toStorableVector . computeAs S) g2 +minimizeNLL :: Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeNLL = minimizeNLL' TNEWTON+ -- | minimizes the function while keeping the parameter ix fixed (used to calculate the profile)-minimizeNLLWithFixedParam :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> Int -> PVector -> PVector-minimizeNLLWithFixedParam dist msErr niter xss ys tree ix t0+minimizeNLLWithFixedParam' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> Int -> PVector -> PVector+minimizeNLLWithFixedParam' alg dist mYerr niter xss ys tree ix t0 | niter == 0 = t0 | n == 0 = t0- | n > m = t0- | otherwise = fromStorableVector compMode t_opt+ | otherwise = t_opt' where+ tree' = buildNLL dist (fromIntegral m) $ relabelParams tree t0' = toStorableVector t0+ treeArr = IntMap.toAscList $ tree2arr tree'+ j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] (Sz n) = size t0 (Sz m) = size ys setTo0 = (VS.// [(ix, 0.0)])- funAndGrad = second (setTo0 . toStorableVector . computeAs S). gradNLLNonUnique dist msErr xss ys tree . fromStorableVector compMode- (f, _) = gradNLLNonUnique dist msErr xss ys tree t0 -- if there's no parameter or no iterations+ funAndGrad = second (setTo0 . toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix - algorithm = LBFGS funAndGrad Nothing- stop = ObjectiveRelativeTolerance 1e-5 :| [MaximumEvaluations (fromIntegral niter)]+ (f, _) = gradNLL dist mYerr xss ys tree t0 -- if there's no parameter or no iterations++ algorithm = alg funAndGrad Nothing -- PRAXIS (fst . funAndGrad) [] Nothing -- TNEWTON funAndGrad Nothing+ stop = ObjectiveRelativeTolerance 1e-8 :| [ObjectiveAbsoluteTolerance 1e-8, MaximumEvaluations (fromIntegral niter)] problem = LocalProblem (fromIntegral n) stop algorithm- t_opt = case minimizeLocal problem t0' of- Right sol -> solutionParams sol- Left e -> t0'+ (t_opt, nEvs) = case minimizeLocal problem t0' of+ Right sol -> (solutionParams sol, nEvals sol) -- traceShow (">>>>>>>", nEvals sol) $+ Left e -> (t0', 0)+ t_opt' = fromStorableVector compMode t_opt +minimizeNLLWithFixedParam = minimizeNLLWithFixedParam' TNEWTON+ -- | minimizes using Gaussian likelihood -minimizeGaussian :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double)+minimizeGaussian :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int) minimizeGaussian = minimizeNLL Gaussian Nothing -- | minimizes using Binomial likelihood -minimizeBinomial :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double)+minimizeBinomial :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int) minimizeBinomial = minimizeNLL Bernoulli Nothing -- | minimizes using Poisson likelihood -minimizePoisson :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double)+minimizePoisson :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int) minimizePoisson = minimizeNLL Poisson Nothing---- estimates the standard error if not provided -estimateSErr :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Int -> Maybe Double-estimateSErr Gaussian Nothing xss ys theta0 t nIter = Just err- where- theta = fst $ minimizeNLL Gaussian (Just 1) nIter xss ys t theta0- (Sz m) = size ys- (Sz p) = size theta- ssr = sse xss ys t theta- err = sqrt $ ssr / fromIntegral (m - p)-estimateSErr _ (Just s) _ _ _ _ _ = Just s-estimateSErr _ _ _ _ _ _ _ = Nothing
src/Data/SRTree.hs view
@@ -35,6 +35,8 @@ , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) ) where@@ -64,5 +66,7 @@ , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) )
src/Data/SRTree/Datasets.hs view
@@ -97,11 +97,11 @@ -- input variables. These will be renamed internally as x0, x1, ... in the order -- of this list. splitFileNameParams :: FilePath -> (FilePath, [B.ByteString])-splitFileNameParams (B.pack -> filename) = (B.unpack fname, take 4 params)+splitFileNameParams (B.pack -> filename) = (B.unpack fname, take 6 params) where (fname : params') = B.split ':' filename -- fill up the empty parameters with an empty string- params = params' <> replicate (4 - min 4 (length params')) B.empty+ params = params' <> replicate (6 - min 6 (length params')) B.empty {-# inline splitFileNameParams #-} -- | Tries to parse a string into an int@@ -114,8 +114,8 @@ -- | Given a map between PVector name and indeces, -- the target PVector and the variables SRMatrix, -- returns the indices of the variables SRMatrix and the target-getColumns :: [(B.ByteString, Int)] -> B.ByteString -> B.ByteString -> ([Int], Int)-getColumns headerMap target columns = (ixs, iy)+getColumns :: [(B.ByteString, Int)] -> B.ByteString -> B.ByteString -> B.ByteString -> ([Int], Int, Int)+getColumns headerMap target columns target_error = (ixs, iy, iy_error) where n_cols = length headerMap getIx c = case parseVal c of@@ -136,6 +136,10 @@ iy = if B.null target then n_cols - 1 else getIx $ B.unpack target+ -- if the target PVector is ommitted, use the last one+ iy_error = if B.null target_error+ then (-1)+ else getIx $ B.unpack target_error {-# inline getColumns #-} -- | Given the start and end rows, it returns the @@ -163,7 +167,7 @@ {-# inline getRows #-} -- | `loadDataset` loads a dataset with a filename in the format:--- filename.ext:start_row:end_row:target:features+-- filename.ext:start_row:end_row:target:features:y_err -- it returns the X_train, y_train, X_test, y_test, varnames, target name -- where varnames are a comma separated list of the name of the vars -- and target name is the name of the target@@ -176,7 +180,7 @@ -- of the target variable -- **features** is a comma separated list of SRMatrix names or indices to be used as -- input variables of the regression model.-loadDataset :: FilePath -> Bool -> IO ((SRMatrix, PVector, SRMatrix, PVector), String, String)+loadDataset :: FilePath -> Bool -> IO ((SRMatrix, PVector, SRMatrix, PVector), (Maybe PVector, Maybe PVector), String, String) loadDataset filename hasHeader = do csv <- readFileToLines fname pure $ processData csv params hasHeader@@ -184,8 +188,8 @@ (fname, params) = splitFileNameParams filename -- support function that does everything for loadDataset-processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> ((SRMatrix, PVector, SRMatrix, PVector), String, String)-processData csv params hasHeader = ((x_train, y_train, x_val, y_val) , varnames, targetname)+processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> ((SRMatrix, PVector, SRMatrix, PVector), (Maybe PVector, Maybe PVector), String, String)+processData csv params hasHeader = ((x_train, y_train, x_val, y_val) , (y_err_train, y_err_val), varnames, targetname) where ncols = length $ head csv nrows = length csv - fromEnum hasHeader@@ -197,8 +201,8 @@ ] targetname = if hasHeader then (B.unpack . fst . fromJust . find ((==iy).snd) $ header) else "y" -- get rows and SRMatrix indices- (st, end) = getRows (params !! 0) (params !! 1) nrows- (ixs, iy) = getColumns header (params !! 2) (params !! 3)+ (st, end) = getRows (params !! 0) (params !! 1) nrows+ (ixs, iy, iy_err) = getColumns header (params !! 2) (params !! 3) (params !! 4) -- load data and split sets datum = loadMtx content@@ -206,9 +210,14 @@ x = M.computeAs S $ M.throwEither $ M.stackInnerSlicesM $ map (datum <!) ixs y = datum <! iy+ y_err = datum <! iy_err+ x_train = M.computeAs S $ M.extractFromTo' (st :. 0) (end+1 :. p) x y_train = M.computeAs S $ M.extractFromTo' st (end+1) y x_val = M.computeAs S $ M.throwEither $ M.deleteRowsM st (Sz1 $ end - st + 1) x y_val = M.computeAs S $ M.throwEither $ M.deleteColumnsM st (Sz1 $ end - st + 1) y++ y_err_train = if iy_err == -1 then Nothing else Just $ M.computeAs S $ M.extractFromTo' st (end+1) y_err+ y_err_val = if iy_err == -1 then Nothing else Just $ M.computeAs S $ M.throwEither $ M.deleteColumnsM st (Sz1 $ end - st + 1) y_err {-# inline processData #-}
src/Data/SRTree/Derivative.hs view
@@ -44,7 +44,7 @@ alg1 (Bin Mul l r) = fst l * snd r + snd l * fst r alg1 (Bin Div l r) = (fst l * snd r - snd l * fst r) / snd r ** 2 alg1 (Bin Power l r) = snd l ** (snd r - 1) * (snd r * fst l + snd l * log (snd l) * fst r)- alg1 (Bin PowerAbs l r) = (abs (snd l) ** (snd r)) * (fst r * log (abs (snd l)) + snd r * fst l / snd l)+ alg1 (Bin PowerAbs l r) = (powabs (snd l) (snd r)) * (fst r * log (abs (snd l)) + snd r * fst l / snd l) alg1 (Bin AQ l r) = ((1 + snd r * snd r) * fst l - snd l * snd r * fst r) / (1 + snd r * snd r) ** 1.5 alg2 (Var ix) = var ix@@ -52,6 +52,8 @@ alg2 (Const c) = Fix (Const c) alg2 (Uni f t) = Fix (Uni f $ snd t) alg2 (Bin f l r) = Fix (Bin f (snd l) (snd r))+ --(abs (snd l) ** (snd r))+ powabs l r = Fix (Bin PowerAbs l r) -- | Derivative of each supported function -- For a function h(f) it returns the derivative dh/df
src/Data/SRTree/Eval.hs view
@@ -79,6 +79,7 @@ -- returns a vector with the same number of rows as xss and containing a single repeated value. replicateAs :: SRMatrix -> Double -> SRVector replicateAs xss c = let (Sz (m :. _)) = M.size xss in M.replicate (getComp xss) (Sz m) c+{-# INLINE replicateAs #-} -- | Evaluates the tree given a vector of variable values, a vector of parameter values and a function that takes a Double and change to whatever type the variables have. This is useful when working with datasets of many values per variables. evalTree :: SRMatrix -> PVector -> Fix SRTree -> SRVector@@ -184,6 +185,7 @@ evalInverse Abs = abs -- we assume abs(x) = sqrt(x^2) so y = sqrt(x^2) => x^2 = y^2 => x = sqrt(y^2) = x = abs(y) evalInverse Recip = recip evalInverse Cube = cbrt+{-# INLINE evalInverse #-} -- | evals the right inverse of an operator invright :: Floating a => Op -> a -> (a -> a)@@ -194,6 +196,7 @@ invright Power v = (**(1/v)) invright PowerAbs v = (**(1/v)) invright AQ v = (* sqrt (1 + v*v))+{-# INLINE invright #-} -- | evals the left inverse of an operator invleft :: Floating a => Op -> a -> (a -> a)@@ -204,7 +207,9 @@ invleft Power v = logBase v -- (/(log v)) . log -- y = v ^ r log y = r log v r = log y / log v invleft PowerAbs v = logBase v . abs invleft AQ v = (v/)+{-# INLINE invleft #-} -- | List of invertible functions invertibles :: [Function] invertibles = [Id, Sin, Cos, Tan, Tanh, ASin, ACos, ATan, ATanh, Sqrt, Square, Log, Exp, Recip]+{-# INLINE invertibles #-}
src/Data/SRTree/Internal.hs view
@@ -41,6 +41,8 @@ , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) ) where@@ -93,6 +95,32 @@ | Recip | Cube deriving (Show, Read, Eq, Ord, Enum)++removeProtectedOps :: Fix SRTree -> Fix SRTree +removeProtectedOps = cata alg + where + alg (Var ix) = var ix+ alg (Param ix) = param ix+ alg (Const x) = constv x+ alg (Bin PowerAbs l r) = l ** r+ alg (Bin op l r) = Fix $ Bin op l r+ alg (Uni SqrtAbs t) = Fix $ Uni Sqrt t+ alg (Uni LogAbs t) = Fix $ Uni Log t+ alg (Uni f t) = Fix $ Uni f t+{-# INLINE removeProtectedOps #-}++convertProtectedOps :: Fix SRTree -> Fix SRTree +convertProtectedOps = cata alg + where + alg (Var ix) = var ix+ alg (Param ix) = param ix+ alg (Const x) = constv x+ alg (Bin PowerAbs l r) = abs l ** r+ alg (Bin op l r) = Fix $ Bin op l r+ alg (Uni SqrtAbs t) = sqrt (abs t)+ alg (Uni LogAbs t) = log (abs t)+ alg (Uni f t) = Fix $ Uni f t+{-# INLINE convertProtectedOps #-} -- | create a tree with a single node representing a variable var :: Int -> Fix SRTree
src/Data/SRTree/Print.hs view
@@ -23,6 +23,7 @@ , printPython , showLatex , printLatex+ , showOp ) where @@ -89,8 +90,8 @@ showOp Mul = "*" showOp Div = "/" showOp Power = "^"-showOp AQ = "_aq_"-showOp PowerAbs = "||^"+showOp AQ = "aq"+showOp PowerAbs = "|^|" {-# INLINE showOp #-} -- | Displays a tree as a numpy compatible expression.
src/Numeric/Optimization/NLOPT/Bindings.hs view
@@ -613,6 +613,8 @@ , resultParameters :: V.Vector Double -- ^ Parameters corresponding -- to the minimum if -- optimization succeeded++ , nEvals :: Int -- ^ number of evaluations } foreign import ccall "nlopt.h nlopt_optimize"@@ -626,12 +628,14 @@ -> IO Output -- ^ Results of the optimization run optimize optimizer x0 = withOpt optimizer $ \opt -> do vmut <- V.thaw $ V.unsafeCast x0- alloca $ \costPtr -> do+ (result, outputCost, iceout) <- alloca $ \costPtr -> do result <- MV.unsafeWith vmut $ \xptr -> parseEnum <$> nlopt_optimize opt xptr costPtr outputCost <- peek . castPtr $ costPtr iceout <- V.unsafeFreeze (MV.unsafeCast vmut)- return $ Output result outputCost iceout+ return (result, outputCost, iceout)+ nEvals <- fromIntegral <$> get_numevals optimizer+ return $ Output result outputCost iceout nEvals {- Objective function setup -} @@ -962,6 +966,12 @@ get_maxeval :: Opt -> IO Word get_maxeval = getScalar nlopt_get_maxeval fromIntegral++foreign import ccall "nlopt.h nlopt_get_numevals"+ nlopt_get_numevals :: NloptOpt -> IO CInt++get_numevals :: Opt -> IO CInt+get_numevals = getScalar nlopt_get_numevals fromIntegral foreign import ccall "nlopt.h nlopt_set_maxtime" nlopt_set_maxtime :: NloptOpt -> CDouble -> IO CInt
src/Text/ParseSR.hs view
@@ -11,7 +11,7 @@ -- Functions to parse a string representing an expression -- ------------------------------------------------------------------------------module Text.ParseSR ( parseSR, showOutput, SRAlgs(..), Output(..) ) +module Text.ParseSR ( parseSR, parsePat, showOutput, SRAlgs(..), Output(..) ) where import Control.Applicative ((<|>))@@ -21,6 +21,7 @@ import Data.Char (toLower) import Data.List (sortOn) import Data.SRTree+import Algorithm.EqSat.DB import qualified Data.SRTree.Print as P import Debug.Trace (trace) @@ -30,6 +31,7 @@ -- numerical values represented as `Double`. The numerical values type -- can be changed with `fmap`. type ParseTree = Parser (Fix SRTree)+type ParsePat = Parser Pattern -- * Data types and caller functions @@ -51,15 +53,18 @@ -- >>> fmap (showOutput MATH) $ parseSR OPERON "lambda,theta" False "lambda ^ 2 - sin(theta*3*lambda)" -- Right "((x0 ^ 2.0) - Sin(((x1 * 3.0) * x0)))" parseSR :: SRAlgs -> B.ByteString -> Bool -> B.ByteString -> Either String (Fix SRTree)-parseSR HL header reparam = eitherResult . (`feed` "") . parse (parseHL reparam $ splitHeader header) . putEOL . B.strip-parseSR BINGO header reparam = eitherResult . (`feed` "") . parse (parseBingo reparam $ splitHeader header) . putEOL . B.strip-parseSR TIR header reparam = eitherResult . (`feed` "") . parse (parseTIR reparam $ splitHeader header) . putEOL . B.strip-parseSR OPERON header reparam = eitherResult . (`feed` "") . parse (parseOperon reparam $ splitHeader header) . putEOL . B.strip-parseSR GOMEA header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip-parseSR SBP header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip-parseSR EPLEX header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip-parseSR PYSR header reparam = eitherResult . (`feed` "") . parse (parsePySR reparam $ splitHeader header) . putEOL . B.strip+parseSR HL header reparam = eitherResult . (`feed` "") . parse (parseHL True reparam $ splitHeader header) . putEOL . B.strip+parseSR BINGO header reparam = eitherResult . (`feed` "") . parse (parseBingo True reparam $ splitHeader header) . putEOL . B.strip+parseSR TIR header reparam = eitherResult . (`feed` "") . parse (parseTIR True reparam $ splitHeader header) . putEOL . B.strip+parseSR OPERON header reparam = eitherResult . (`feed` "") . parse (parseOperon True reparam $ splitHeader header) . putEOL . B.strip+parseSR GOMEA header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR SBP header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR EPLEX header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR PYSR header reparam = eitherResult . (`feed` "") . parse (parsePySR True reparam $ splitHeader header) . putEOL . B.strip +parsePat :: B.ByteString -> Either String Pattern+parsePat = eitherResult . (`feed` "") . parse parsePatExpr . putEOL . B.strip+ eitherResult' :: Show r => Result r -> Either String r eitherResult' res = trace (show res) $ eitherResult res @@ -81,10 +86,11 @@ -- the name of the functions and operators of that SR algorithm, a list of parsers `binFuns` for binary functions -- a parser `var` for variables, a boolean indicating whether to change floating point values to free -- parameters variables, and a list of variable names with their corresponding indexes.-parseExpr :: [[Operator B.ByteString (Fix SRTree)]] -> [ParseTree -> ParseTree] -> ParseTree -> Bool -> [(B.ByteString, Int)] -> ParseTree-parseExpr table binFuns var reparam header = do e <- relabelParams <$> expr- many1' space- pure e+parseExpr :: Bool -> [[Operator B.ByteString (Fix SRTree)]] -> [ParseTree -> ParseTree] -> ParseTree -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseExpr relabel table binFuns var reparam header =+ do e <- if relabel then (relabelParams <$> expr) else expr+ many1' space+ pure e where term = parens expr <|> enclosedAbs expr <|> choice (map ($ expr) binFuns) <|> coef <|> varC <?> "term" expr = buildExpressionParser table term@@ -144,6 +150,15 @@ cbrt :: Fix SRTree -> Fix SRTree cbrt x = x ** (1/3) +cube :: Fix SRTree -> Fix SRTree+cube x = Fix $ Uni Cube x++sqrtabs :: Fix SRTree -> Fix SRTree+sqrtabs x = Fix $ Uni SqrtAbs x++logabs :: Fix SRTree -> Fix SRTree+logabs x = Fix $ Uni LogAbs x+ -- Parse `abs` functions as | x | enclosedAbs :: Num a => Parser a -> Parser a enclosedAbs expr = do char '|'@@ -164,8 +179,8 @@ -- * Custom parsers for SR algorithms -- | parser for Transformation-Interaction-Rational.-parseTIR :: Bool -> [(B.ByteString, Int)] -> ParseTree-parseTIR = parseExpr (prefixOps : binOps) binFuns var+parseTIR :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseTIR b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [ ] prefixOps = map (uncurry prefix)@@ -174,28 +189,35 @@ , ("sin", sin), ("cos", cos), ("tan", tan) , ("asinh", asinh), ("acosh", acosh), ("atanh", atanh) , ("asin", asin), ("acos", acos), ("atan", atan)- , ("sqrt", sqrt), ("cbrt", cbrt), ("square", (**2))- , ("log", log), ("exp", exp)+ , ("sqrtabs", sqrtabs), ("sqrt", sqrt), ("cbrt", cbrt), ("square", (**2))+ , ("logabs", logabs), ("log", log), ("exp", exp), ("cube", cube), ("recip", recip) , ("Id", id), ("Abs", abs) , ("Sinh", sinh), ("Cosh", cosh), ("Tanh", tanh) , ("Sin", sin), ("Cos", cos), ("Tan", tan) , ("ASinh", asinh), ("ACosh", acosh), ("ATanh", atanh) , ("ASin", asin), ("ACos", acos), ("ATan", atan)- , ("Sqrt", sqrt), ("Cbrt", cbrt), ("Square", (**2))- , ("Log", log), ("Exp", exp)+ , ("SqrtAbs", sqrtabs), ("Sqrt", sqrt), ("Cbrt", cbrt), ("Square", (**2))+ , ("LogAbs", logabs), ("Log", log), ("Exp", exp), ("Recip", recip), ("Cube", cube) ] binOps = [[binary "^" (**) AssocLeft], [binary "**" (**) AssocLeft] , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft] , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ , [binary "|**|" powabs AssocLeft], [binary "aq" aq AssocLeft] ]+ powabs l r = Fix $ Bin PowerAbs l r+ aq l r = Fix $ Bin AQ l r+ var = do char 'x' ix <- decimal pure $ Fix $ Var ix+ <|> do char 't'+ ix <- decimal+ pure $ Fix $ Param ix <?> "var" -- | parser for Operon.-parseOperon :: Bool -> [(B.ByteString, Int)] -> ParseTree-parseOperon = parseExpr (prefixOps : binOps) binFuns var+parseOperon :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseOperon b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [ binFun "pow" (**) ] prefixOps = map (uncurry prefix)@@ -216,8 +238,8 @@ <?> "var" -- | parser for HeuristicLab.-parseHL :: Bool -> [(B.ByteString, Int)] -> ParseTree-parseHL = parseExpr (prefixOps : binOps) binFuns var+parseHL :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseHL b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [ binFun "aq" aq ] prefixOps = map (uncurry prefix)@@ -237,8 +259,8 @@ <?> "var" -- | parser for Bingo-parseBingo :: Bool -> [(B.ByteString, Int)] -> ParseTree-parseBingo = parseExpr (prefixOps : binOps) binFuns var+parseBingo :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseBingo b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [] prefixOps = map (uncurry prefix)@@ -257,8 +279,8 @@ <?> "var" -- | parser for GOMEA-parseGOMEA :: Bool -> [(B.ByteString, Int)] -> ParseTree-parseGOMEA = parseExpr (prefixOps : binOps) binFuns var+parseGOMEA :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseGOMEA b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [] prefixOps = map (uncurry prefix)@@ -276,8 +298,8 @@ <?> "var" -- | parser for PySR-parsePySR :: Bool -> [(B.ByteString, Int)] -> ParseTree-parsePySR = parseExpr (prefixOps : binOps) binFuns var+parsePySR :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parsePySR b = parseExpr b (prefixOps : binOps) binFuns var where binFuns = [ binFun "pow" (**) ] prefixOps = map (uncurry prefix)@@ -300,3 +322,64 @@ ix <- decimal pure $ Fix $ Var ix <?> "var"++-- parse a pattern expression+parsePatExpr :: ParsePat+parsePatExpr = parsePattern (prefixOps : binOps) binFuns var+ where+ binFuns = [ ]+ prefixOps = map (uncurry prefix)+ [ ("id", id), ("abs", abs)+ , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh)+ , ("sin", sin), ("cos", cos), ("tan", tan)+ , ("asinh", asinh), ("acosh", acosh), ("atanh", atanh)+ , ("asin", asin), ("acos", acos), ("atan", atan)+ , ("sqrtabs", sqrtabs'), ("sqrt", sqrt), ("cbrt", cbrt'), ("square", (**2))+ , ("logabs", logabs'), ("log", log), ("exp", exp), ("cube", cube'), ("recip", recip')+ , ("Id", id), ("Abs", abs)+ , ("Sinh", sinh), ("Cosh", cosh), ("Tanh", tanh)+ , ("Sin", sin), ("Cos", cos), ("Tan", tan)+ , ("ASinh", asinh), ("ACosh", acosh), ("ATanh", atanh)+ , ("ASin", asin), ("ACos", acos), ("ATan", atan)+ , ("SqrtAbs", sqrtabs'), ("Sqrt", sqrt), ("Cbrt", cbrt'), ("Square", (**2))+ , ("LogAbs", logabs'), ("Log", log), ("Exp", exp), ("Recip", recip'), ("Cube", cube')+ , ("|log|", logabs'), ("|Log|", logabs'), ("|sqrt|", sqrtabs'), ("|Sqrt|", sqrtabs')+ , ("√", sqrt), ("|√|", sqrtabs')+ ]+ binOps = [[binary "^" (**) AssocLeft], [binary "**" (**) AssocLeft]+ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ , [binary "|**|" powabs AssocLeft], [binary "|^|" powabs AssocLeft]+ , [binary "aq" aq AssocLeft], [binary "|/|" aq AssocLeft]+ ]+ powabs l r = Fixed $ Bin PowerAbs l r+ aq l r = Fixed $ Bin AQ l r+ logabs' t = Fixed $ Uni LogAbs t+ sqrtabs' t = Fixed $ Uni SqrtAbs t+ cbrt' t = Fixed $ Uni Cbrt t+ cube' t = Fixed $ Uni Cube t+ recip' t = Fixed $ Uni Recip t++ var = do char 'x'+ ix <- decimal+ pure $ Fixed $ Var ix+ <|> do char 't'+ ix <- decimal+ pure $ Fixed $ Param ix+ <|> do char 'v'+ ix <- decimal+ pure $ VarPat (toEnum $ ix+65)+ <?> "var"++parsePattern :: [[Operator B.ByteString Pattern]] -> [ParsePat -> ParsePat] -> ParsePat -> ParsePat+parsePattern table binFuns var =+ do e <- expr+ many1' space+ pure e+ where+ term = parens expr <|> enclosedAbs expr <|> choice (map ($ expr) binFuns) <|> coef <|> var <?> "term"+ expr = buildExpressionParser table term+ coef = Fixed . Const <$> signed double <?> "const"++ getParserVar k v = (string k <|> enveloped k) >> pure (Fix $ Var v)+ enveloped s = (char ' ' <|> char '(') >> string s >> (char ' ' <|> char ')') >> pure ""
srtree.cabal view
@@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack name: srtree-version: 2.0.0.2+version: 2.0.0.3 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@@ -56,7 +56,7 @@ Paths_srtree hs-source-dirs: src- ghc-options: -fwarn-incomplete-patterns+ ghc-options: -fwarn-incomplete-patterns -threaded extra-lib-dirs: /usr/local/lib extra-libraries:@@ -64,7 +64,8 @@ build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -76,10 +77,13 @@ , list-shuffle >=1.0.0.1 && <1.1 , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8@@ -89,14 +93,16 @@ main-is: Main.hs other-modules: Random+ Util Paths_srtree hs-source-dirs: apps/egraphGP- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -109,27 +115,70 @@ , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 , optparse-applicative >=0.17 && <0.19- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8 default-language: Haskell2010 +executable egraphSearch+ main-is: Main.hs+ other-modules:+ Random+ Util+ Paths_srtree+ hs-source-dirs:+ apps/egraphSearch+ ghc-options: -threaded -rtsopts -with-rtsopts=-N -Wall+ build-depends:+ attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.8+ , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.0 && <1.1+ , mtl >=2.2 && <2.4+ , optparse-applicative >=0.17 && <0.19+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5+ , srtree+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.*+ , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8+ default-language: Haskell2010+ executable eqsatrepr main-is: Main.hs other-modules: Paths_srtree hs-source-dirs: apps/eqsatrepr- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -141,27 +190,35 @@ , list-shuffle >=1.0.0.1 && <1.1 , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8 default-language: Haskell2010 -executable ieeexplore+executable reggression main-is: Main.hs other-modules:+ Commands+ Random+ Util Paths_srtree hs-source-dirs:- apps/ieeexplore- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3+ apps/rEGGression+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends:- attoparsec >=0.14.4 && <0.15+ ansi-terminal >=1.1.2 && <1.2+ , attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -174,11 +231,16 @@ , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 , optparse-applicative >=0.17 && <0.19- , random ==1.2.*+ , random >=1.2 && <1.4+ , repline >=0.4.2.0 && <0.5+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17+ , table-layout >=1.0.0.1 && <1.1 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8@@ -190,11 +252,12 @@ Paths_srtree hs-source-dirs: apps/srsimplify- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -207,11 +270,14 @@ , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 , optparse-applicative >=0.17 && <0.19- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8@@ -226,11 +292,12 @@ Paths_srtree hs-source-dirs: apps/srtools- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -242,13 +309,15 @@ , list-shuffle >=1.0.0.1 && <1.1 , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4- , normaldistribution >=1.1.0.3 && <1.2 , optparse-applicative >=0.17 && <0.19- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8@@ -259,14 +328,16 @@ other-modules: GP Initialization+ Util Paths_srtree hs-source-dirs: apps/tinygp- ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3+ ghc-options: -threaded -rtsopts -with-rtsopts=-N build-depends: attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -279,11 +350,14 @@ , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 , optparse-applicative >=0.17 && <0.19- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8@@ -302,7 +376,8 @@ , ad , attoparsec >=0.14.4 && <0.15 , attoparsec-expr >=0.1.1.2 && <0.2- , base >=4.16 && <5+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9 , bytestring >=0.11 && <0.13 , containers >=0.6.7 && <0.8 , dlist ==1.0.*@@ -314,11 +389,14 @@ , list-shuffle >=1.0.0.1 && <1.1 , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3 , split >=0.2.5 && <0.3 , srtree , statistics >=0.16.2.1 && <0.17 , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1 , unordered-containers ==0.2.* , vector >=0.12 && <0.14 , zlib >=0.6.3 && <0.8