diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,20 @@ # Changelog for srtree +## 3.0.0.0 + +- **BREAKING**: Removed the Accelerate AD backend (`Algorithm.SRTree.AD.Accelerate`). + The `ADBackEnd` type now only has `SingleThread` and `MultiThread` constructors. + This removes the `accelerate` and `accelerate-llvm-native` dependencies. +- Out-of-core equality saturation with paged e-graph store (SQLite/PostgreSQL) +- Frontier re-saturation: mark changed classes and re-saturate only the frontier +- Streaming matcher for n-ary and cached genericJoin paths (O(1) memory on paged graphs) +- Cycle-safe and size-budgeted `getBestExpr` extraction +- Bounded cost/best fixpoints so recalc terminates on cyclic graphs +- Bounded node-to-class and canonical maps on paged graphs (LRU caches) +- Fast ByteString double parser for dataset loading +- Thread `Loss` (not `Distribution`) through fitness functions; add `readLoss` +- Multiset e-graph improvements + ## 2.0.1.7 - Added log10 MSE fitness function diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,5 @@ -Copyright Author name here (c) 2021 +Copyright (c) 2026, folivetti -All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: @@ -13,7 +12,7 @@ disclaimer in the documentation and/or other materials provided with the distribution. - * Neither the name of Author name here nor the names of other + * Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. @@ -21,7 +20,7 @@ "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY diff --git a/apps/Bench/Main.hs b/apps/Bench/Main.hs new file mode 100644 --- /dev/null +++ b/apps/Bench/Main.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE BangPatterns #-} + +import Criterion.Main +import Control.DeepSeq (force, NFData) +import Control.Exception (evaluate) +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as G +import qualified Data.Vector.Storable as VS + +import Data.SRTree +import Data.SRTree.Print +import Data.SRTree.Datasets +import Data.SRTree.Eval +import Data.SRTree.Random +import System.Random +import Control.Monad.State.Strict +import Algorithm.SRTree.NonlinearOpt +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.AD + +-- Assuming these are exported by your project modules: +-- import SRTree +-- import Compiler +-- import DatasetLoader + +-- Mock signatures based on your provided functions +-- randomTree :: Int -> Int -> Int -> IO Term -> IO NonTerm -> Bool -> IO Tree +-- loadDataset :: FilePath -> Bool -> IO [V.Vector Double] +-- evalTree :: Tree -> [V.Vector Double] -> V.Vector Double +-- compile :: [V.Vector Double] -> Tree -> (Theta -> V.Vector Double) + +genTerm = do coin <- tossBiased 0.4 + if coin then randomFrom [Fix $ Var ix | ix <- [0..8]] else randomFrom [Fix $ Param ix | ix <- [0..9]] +genNonTerm = randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Uni LogAbs (), Uni SqrtAbs ()] + +genMultipleTrees 0 = pure [] +genMultipleTrees n = do + t <- randomTree 5 10 150 genTerm genNonTerm False + ts <- genMultipleTrees (n-1) + pure (t:ts) + +getF (_, x, _) = x +{-# INLINE getF #-} +getT (t, _, _) = t +{-# INLINE getT #-} + +main :: IO () +main = do + -- 1. Initialization: Load the dataset + putStrLn "Loading dataset..." + ((dataset, y, _, _), _, _, _) <- loadDataset "data.tsv" True + + -- 2. Initialization: Generate the random expression tree + putStrLn "Generating random tree..." + -- Replace 'genTerm' and 'genNonTerm' with your actual generators + --g <- getStdGen + let g = mkStdGen 42 + -- tree <- evalStateT (randomTree 7 10 150 genTerm genNonTerm True) g + trees' <- evalStateT (genMultipleTrees 5) g + -- let trees' = [Fix (Uni LogAbs (Fix (Bin PowerAbs (param 0) (param 1 * var 0))))] :: [Fix SRTree] + + -- IMPORTANT: Force deep evaluation of the tree and dataset. + -- If we do not do this, GHC's lazy evaluation will cause the benchmark + -- to measure the time it takes to parse the CSV and build the tree in memory! + -- _ <- evaluate (force tree) + _ <- evaluate (force dataset) + + + -- 3. Initialization: Pre-compile the tree + -- We evaluate this strictly (!) so the one-time compilation cost + -- is not included in the runtime benchmark. + putStrLn "Compiling tree..." + let !compiledFn = [compile dataset tree | tree <- trees] + evalTree x th t = compile x t th + -- Mock theta (parameter vector) to pass into the closures + !theta = V.fromList [1.0, 0.5, 0.2, 0.3, 0.1, 0.5, 0.9, 0.3, 0.2, 0.4] + !theta1 = V.fromList [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + trees = map relabelParamsOrder $ filter (\t -> let v = V.sum (evalTree dataset theta t) in not (isInfinite v || isNaN v)) trees' + naiveEval = evalTree dataset theta + dataset' = map G.convert dataset + y' = G.convert y + theta1' = G.convert theta1 + + _ <- evaluate (force theta) + _ <- evaluate (force theta1) + print $ sum $ map (\t -> V.sum $ naiveEval t) trees + print $ sum $ map (\t -> V.sum $ t theta) compiledFn + print $ sum $ map (\t -> getF $ minimizeNLL MultiThread MSE Nothing 0 dataset y t theta1) trees + --print $ sum $ map (\t -> getF $ minimizeNLLCompiled MSE Nothing 0 dataset y t theta1) trees + + --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraph MSE dataset' y' Nothing t theta1') trees + --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraphO MSE dataset' y' Nothing t theta1') trees + --print $ sum $ map (\t -> VS.sum . snd $ compileGrad dataset' y' Nothing t 100 theta1') trees + --print $ sum $ map (\ct -> V.sum $ ct theta) compiledFn + --print $ sum $ map (\ct -> V.sum $ executeVM ct rowDataset theta) bytecodes + -- print $ V.sum $ evalTree dataset theta tree + -- print $ V.sum $ compiledFn theta + + putStrLn "Running benchmarks..." + + -- 4. The Benchmarks + defaultMain [ + bgroup "Tree Evaluation (Fixed Dataset)" [ + + -- The slow version: dynamically traversing the AST at runtime + bench "evalTree (Naive AST Traversal)" $ + nf (\ts -> sum [V.sum $ evalTree dataset theta1 t | t <- ts]) trees, + + + -- The fast version: executing the pre-compiled, stream-fused closure + bench "compile (Compiled Closure)" $ + nf (\t -> sum [V.sum (ct t) | ct <- compiledFn]) theta1, + + -- The fast version: executing the pre-compiled, stream-fused closure + bench "minimizeNLLCompiled (Compiled Closure)" $ + nf (\ts -> sum [V.sum . getT $ minimizeNLL MultiThread MSE Nothing 100 dataset' y' t theta1' | t <- ts]) trees + + --bench "minimizeNLLO (Naive optimized AST Traversal)" $ + -- nf (\ts -> sum [V.sum . getT $ minimizeNLLO MSE Nothing 100 dataset y t theta1 | t <- ts]) trees + + -- The slow version: dynamically traversing the AST at runtime + --bench "minimizeNLL (Naive AST Traversal)" $ + -- nf (\ts -> sum [V.sum . getT $ minimizeNLL (NLL MSE) Nothing 100 dataset y t theta1 | t <- ts]) trees + + ] + ] diff --git a/apps/BenchEqSat/Main.hs b/apps/BenchEqSat/Main.hs new file mode 100644 --- /dev/null +++ b/apps/BenchEqSat/Main.hs @@ -0,0 +1,259 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +import Criterion.Main +import qualified Data.Vector.Unboxed as VU +import qualified Data.IntMap as IntMap +import qualified Data.HashMap.Strict as HashMap +import qualified Data.HashSet as Set + +import Data.SRTree +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Control.Monad.State.Strict +import Control.Monad (replicateM, zipWithM_) +import Control.Monad.Identity + +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 + +evalEG :: EGraphST Identity a -> (a, EGraph) +evalEG m = runIdentity $ runStateT m emptyGraph + +runInEG :: EGraph -> EGraphST Identity a -> (a, EGraph) +runInEG eg m = runIdentity $ runStateT m eg + +-- Expression generators for benchmarking +chainAdd :: Int -> Fix SRTree +chainAdd 0 = var 0 +chainAdd n = chainAdd (n-1) + var n + +deepBinTree :: Int -> Fix SRTree +deepBinTree 0 = var 0 +deepBinTree n = deepBinTree (n-1) + constv (fromIntegral n) + +complexTree :: Int -> Fix SRTree +complexTree n = go n + where + go 0 = var 0 + go i = (var i + constv (fromIntegral i)) * (go (i-1) + constv (fromIntegral i)) + +simplifyRules :: [Rule] +simplifyRules = + [ "a" + 0 :=> "a" + , "a" * 1 :=> "a" + , "a" + "a" :=> 2 * "a" + , "a" * 0 :=> 0 + , 0 + "a" :=> "a" + , 1 * "a" :=> "a" + ] + +-- More rules including commutativity (triggers more merges) +moreRules :: [Rule] +moreRules = + [ "a" + 0 :=> "a" + , "a" * 1 :=> "a" + , "a" + "a" :=> 2 * "a" + , "a" * 0 :=> 0 + , 0 + "a" :=> "a" + , 1 * "a" :=> "a" + , "a" + "b" :=> "b" + "a" + , "a" * "b" :=> "b" * "a" + ] + +addZero :: Fix SRTree -> Fix SRTree -> Fix SRTree +addZero l r = Fix (Bin Add l r) + +main :: IO () +main = do + putStrLn "Generating benchmark expressions..." + let smallExpr = chainAdd 5 + mediumExpr = chainAdd 20 + largeExpr = chainAdd 100 + complex = complexTree 8 + + putStrLn "Running benchmarks..." + defaultMain [ + bgroup "E-graph Construction" [ + bench "fromTree (5-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) smallExpr, + bench "fromTree (20-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) mediumExpr, + bench "fromTree (100-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) largeExpr, + bench "fromTree (complex-ternary tree)" $ + whnf (\e -> evalEG $ fromTree myCost e) complex + ], + + bgroup "E-graph Add" [ + bench "add single e-node (Var)" $ + whnf (\eg -> runInEG eg $ add myCost (EVar 999)) (snd $ evalEG $ fromTree myCost smallExpr), + bench "add single e-node (Const)" $ + whnf (\eg -> runInEG eg $ add myCost (EConst 42.0)) (snd $ evalEG $ fromTree myCost smallExpr), + bench "add single e-node (Bin Add)" $ + whnf (\eg -> runInEG eg $ add myCost (ENAry EAdd (imFromList [0, 1]))) (snd $ evalEG $ fromTree myCost mediumExpr) + ], + + bgroup "Merge" [ + bench "merge two distinct eclasses (size 1)" $ + whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 1), + bench "merge two distinct eclasses (size 3)" $ + whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 3) + ], + + bgroup "Pattern Matching" [ + bench "match simple pattern (a+0)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMatchableEG), + bench "match commutative pattern (a+b)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMatchableEG), + bench "match triple pattern (a+b+c)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" + "c" :: Pattern)) (makeMatchableEG) + ], + + bgroup "Match After Merge" [ + bench "match (a+0) after merge (stale trie keys)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMergedEG), + bench "match (a+b) after merge (stale trie keys)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMergedEG) + ], + + bgroup "Rebuild" [ + bench "rebuild after 5 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 5), + bench "rebuild after 20 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 20), + bench "rebuild after 100 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 100) + ], + + bgroup "Cost Propagation" [ + bench "recalculateBest (10 eclasses)" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 10), + bench "recalculateBest (100 eclasses)" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 100) + ], + + bgroup "DB Operations" [ + bench "addToDB single enode" $ + whnf (\(en,eid,eg) -> runInEG eg $ addToDB en eid) (makeDBEntry), + bench "addToDB 10 enodes" $ + whnf (\(ens,eg) -> runInEG eg $ mapM_ (uncurry addToDB) ens) (makeDBEntries 10) + ], + + bgroup "Equality Saturation" [ + bench "eqSat small expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, simplifyRules), + bench "eqSat medium expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (mediumExpr, simplifyRules), + bench "eqSat small expr (8 rules, commutative)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, moreRules), + bench "eqSat large expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (largeExpr, simplifyRules) + ], + + bgroup "Extraction" [ + bench "getBestExpr (5-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 5), + bench "getBestExpr (20-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 20), + bench "getBestExpr (100-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 100) + ], + + bgroup "Fitness Operations" [ + bench "insertFitness single" $ + whnf (\(eid,eg) -> runInEG eg $ insertFitness eid 0.5 []) (makeExtractable 1), + bench "insertFitness 10 eclasses" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (\eid -> insertFitness eid 0.5 []) eids) (makeNEclasses 10), + bench "getTopFitEClassWithSize" $ + whnf (\(eids,eg) -> runInEG eg $ getTopFitEClassWithSize 1 3) (makeFitnessEG) + ] + ] + where + addZeroTree = addZero (var 0) (constv 0.0) + + makeMergePair :: Int -> (EClassId, EClassId, EGraph) + makeMergePair n = + let tree = deepBinTree n + (eid1, eg1) = evalEG $ fromTree myCost tree + (eid2, eg2) = runInEG eg1 $ fromTree myCost tree + in (eid1, eid2, eg2) + + makeMatchableEG :: (EGraph, EClassId) + makeMatchableEG = + let tree = complexTree 4 + (eid, eg) = evalEG $ do + eid' <- fromTree myCost tree + _ <- fromTree myCost (var 0 + constv 1.0) + _ <- fromTree myCost (var 1 * constv 2.0) + _ <- fromTree myCost (var 0 + constv 0.0) + _ <- fromTree myCost (var 1 * constv 1.0) + rebuild myCost + pure eid' + in (eg, eid) + + -- E-graph with merges applied, creating stale trie keys + makeMergedEG :: (EGraph, EClassId) + makeMergedEG = + let (_, eg) = evalEG $ do + eid1 <- fromTree myCost (var 0) + eid2 <- fromTree myCost (constv 0.0) + eid3 <- fromTree myCost (var 0 + constv 1.0) + _ <- fromTree myCost (var 1) + rebuild myCost + -- merge to create stale trie entries + merge myCost eid1 eid2 + merge myCost eid2 eid3 + rebuild myCost + pure eid1 + in (eg, 0) + + makeDirtyEG :: Int -> (EGraph, EClassId) + makeDirtyEG n = + let tree = deepBinTree n + (eid, eg) = evalEG $ do + eid' <- fromTree myCost tree + _ <- fromTree myCost (tree + var 999) + rebuild myCost + _ <- fromTree myCost (tree * var 998) + pure eid' + in (eg, eid) + + makeExtractable :: Int -> (EClassId, EGraph) + makeExtractable n = + let tree = deepBinTree n + in evalEG $ fromTree myCost tree + + makeNEclasses :: Int -> ([EClassId], EGraph) + makeNEclasses n = + evalEG $ replicateM n (fromTree myCost (constv (fromIntegral n))) + + makeFitnessEG :: ([EClassId], EGraph) + makeFitnessEG = evalEG $ do + eids <- mapM (fromTree myCost . constv . fromIntegral) [1..10] + zipWithM_ (\eid i -> insertFitness eid (fromIntegral i) []) eids [1..] + pure eids + + makeDBEntry :: (ENode, EClassId, EGraph) + makeDBEntry = + let (eid, eg) = evalEG $ do + eid <- fromTree myCost (var 999) + rebuild myCost + pure eid + in (EVar 777, eid, eg) + + makeDBEntries :: Int -> ([(ENode, EClassId)], EGraph) + makeDBEntries n = + let (eids, eg) = evalEG $ do + eids <- mapM (fromTree myCost . var) [999..(999 + n - 1)] + rebuild myCost + pure eids + in (zip (map EVar [1000..]) eids, eg) diff --git a/apps/Report/Main.hs b/apps/Report/Main.hs new file mode 100644 --- /dev/null +++ b/apps/Report/Main.hs @@ -0,0 +1,247 @@ +module Main (main) where + +import Options.Applicative +import qualified Data.ByteString.Char8 as B +import qualified Data.Vector.Unboxed as U +import Data.SRTree +import Data.SRTree.Eval (Target, Columns, compileLoss) +import Data.SRTree.Datasets (loadTrainingOnly) +import Data.SRTree.Print (showExpr) +import Text.ParseSR (parseSR, SRAlgs(..)) +import Algorithm.SRTree.Compile (compileTree, EvalTree(..), logParameters, logParametersLatt) +import Algorithm.SRTree.Likelihoods (Distribution(..), Loss(..), buildLoss, fisherNLL, hessianNLL) +import Algorithm.SRTree.ConfidenceIntervals + ( getStatsFromModel, paramCI, CIType(..), CI(..), BasicStats(..) + , ProfileT(..), PType(..), getAllProfiles, getCol + ) +import Algorithm.SRTree.ModelSelection (ModelEval(..), logFunctional, logFunctionalFreq) +import Statistics.Distribution (ContDistr(quantile)) +import Statistics.Distribution.FDistribution (fDistribution) +import Control.Exception (try, SomeException) +import Data.List.Split (splitOn) +import Text.Printf (printf) +import Control.Monad (forM_, when) + +---------------------------------------------------------------------- +-- CLI argument types +---------------------------------------------------------------------- +data CIMethod = LaplaceCI | ProfileCI deriving (Show) + +data ProfileTypeArg = BatesArg | ODEArg | ConstrainedArg deriving (Read) +instance Show ProfileTypeArg where + show BatesArg = "Bates" + show ODEArg = "ODE" + show ConstrainedArg = "Constrained" + +data ReportArgs = ReportArgs + { raExprs :: !FilePath + , raFormat :: !SRAlgs + , raData :: !FilePath + , raHeader :: !Bool + , raDist :: !Distribution + , raCriteria :: ![ModelEval] + , raCI :: !CIMethod + , raAlpha :: !Double + , raCIType :: !ProfileTypeArg + , raDbg :: !Bool + } + +---------------------------------------------------------------------- +-- Argument parser +---------------------------------------------------------------------- +argParser :: Parser ReportArgs +argParser = ReportArgs + <$> strOption ( long "exprs" <> short 'e' <> help "File with expressions, one per line" <> metavar "FILE" ) + <*> option auto ( long "format" <> short 'f' <> help "Expression format: TIR, HL, OPERON, BINGO, GOMEA, PYSR, SBP, EPLEX" <> metavar "FMT" ) + <*> strOption ( long "data" <> short 'd' <> help "Dataset file (optionally with :start:end:target:features:y_err)" <> metavar "FILE" ) + <*> switch ( long "header" <> help "Dataset has a header row" ) + <*> option auto ( long "dist" <> value Gaussian <> help "Distribution: Gaussian, Bernoulli, Poisson, LeastSquares" <> metavar "DIST" <> showDefault ) + <*> option parseCriteria ( long "criteria" <> short 'c' <> value [RMSE, R2, AIC, BIC] <> help "Comma-separated criteria" <> metavar "CRITERIA" <> showDefault ) + <*> option parseCI ( long "ci" <> value LaplaceCI <> help "CI method: Laplace, Profile" <> metavar "METHOD" <> showDefault ) + <*> option auto ( long "alpha" <> value 0.05 <> help "Significance level" <> metavar "ALPHA" <> showDefault ) + <*> option parseProfileType ( long "ci-type" <> value BatesArg <> help "Profile CI type: Bates, ODE, Constrained" <> metavar "TYPE" <> showDefault ) + <*> switch ( long "dbg" <> help "Debug: dump profile tau/theta spline points" ) + +parseCriteria :: ReadM [ModelEval] +parseCriteria = eitherReader $ \s -> + case traverse parseOne (splitOn "," s) of + Right es -> Right es + Left e -> Left e + where + parseOne "RMSE" = Right RMSE + parseOne "R2" = Right R2 + parseOne "AIC" = Right AIC + parseOne "BIC" = Right BIC + parseOne "Evidence" = Right Evidence + parseOne "FBF" = Right FBF + parseOne "MDL" = Right MDL + parseOne "MDLLatt" = Right MDLLatt + parseOne "MDLFreq" = Right MDLFreq + parseOne "NLL" = Right (EvalLoss (NLL Gaussian)) + parseOne s = Left ("unknown criterion: " ++ s) + +parseCI :: ReadM CIMethod +parseCI = eitherReader $ \s -> case s of + "Laplace" -> Right LaplaceCI + "Profile" -> Right ProfileCI + _ -> Left ("unknown CI method: " ++ s ++ " (use Laplace or Profile)") + +parseProfileType :: ReadM ProfileTypeArg +parseProfileType = eitherReader $ \s -> case s of + "Bates" -> Right BatesArg + "ODE" -> Right ODEArg + "Constrained" -> Right ConstrainedArg + _ -> Left ("unknown profile type: " ++ s ++ " (use Bates, ODE, or Constrained)") + +---------------------------------------------------------------------- +-- Report data +---------------------------------------------------------------------- +data ReportData = ReportData + { rdTree :: Fix SRTree + , rdTheta :: Target + , rdStdErr :: Target + , rdCriteria :: [(ModelEval, Double)] + , rdCIs :: [CI] + } + +---------------------------------------------------------------------- +-- Main +---------------------------------------------------------------------- +main :: IO () +main = do + args <- execParser (info (argParser <**> helper) fullDesc) + (xss, ys, mYerr) <- loadTrainingOnly (raData args) (raHeader args) + content <- B.readFile (raExprs args) + let exprs = filter (not . B.null) $ B.lines content + mapM_ (processOne args xss ys mYerr) (zip [(1 :: Int) ..] exprs) + +---------------------------------------------------------------------- +-- Process a single expression +---------------------------------------------------------------------- +processOne :: ReportArgs -> Columns -> Target -> Maybe Target -> (Int, B.ByteString) -> IO () +processOne args xss ys mYerr (idx, src) = do + result <- try $ do + tree <- case parseSR (raFormat args) B.empty True src of + Left e -> fail ("parse error: " ++ e) + Right t -> return $! relabelParams t + let dist = raDist args + nRows = U.length ys + nModelParams = countParamsUniq tree + nParams = nModelParams + + case dist of + Gaussian -> 1 + ROXY -> 3 + _ -> 0 + + let et = compileTree dist xss ys mYerr tree + theta0 = U.replicate nParams 1.0 + thetaOpt = ctOptimizer et theta0 + + when (any isNaN (U.toList thetaOpt)) $ + fail "optimisation returned NaN" + + let mseTree = buildLoss MSE (fromIntegral nRows) tree + mseLoss = compileLoss xss mseTree ys mYerr thetaOpt + nllLoss = ctNLL et thetaOpt + tss = ctVar et + + let fisherDiag = fisherNLL dist mYerr xss ys tree thetaOpt + hessCols = hessianNLL dist mYerr xss ys tree thetaOpt + hessLists = map U.toList hessCols + logP = logParameters fisherDiag thetaOpt + logPLatt = logParametersLatt hessLists fisherDiag thetaOpt + logF = logFunctional tree + logFFreq = logFunctionalFreq tree + nF = fromIntegral nRows + kF = fromIntegral nParams + crits = map (\c -> (c, evalOne c mseLoss nllLoss tss nF kF logP logPLatt logF logFFreq)) + (raCriteria args) + + let stats = getStatsFromModel dist mYerr xss ys tree thetaOpt + laplaceCIs = paramCI (Laplace stats) nRows thetaOpt (raAlpha args) + let ptype = case raCIType args of + BatesArg -> Bates + ODEArg -> ODE + ConstrainedArg -> Constrained + let kInt = U.length thetaOpt + nInt = U.length ys + profT = sqrt $ quantile (fDistribution (fromIntegral kInt) (fromIntegral $ nInt - kInt)) (1 - raAlpha args) + cis <- case raCI args of + LaplaceCI -> return laplaceCIs + ProfileCI -> do + let profiles = getAllProfiles ptype et thetaOpt (_stdErr stats) laplaceCIs (raAlpha args) + when (raDbg args) $ forM_ (zip [0..] profiles) $ \(i, ProfileT taus thetas _ tau2theta _) -> do + putStrLn $ "DEBUG Profile " ++ show i ++ " (opt=" ++ show (thetaOpt U.! i) ++ "):" + putStrLn $ " tau range: [" ++ show (if U.null taus then 0 else U.head taus) + ++ ", " ++ show (if U.null taus then 0 else U.last taus) ++ "]" + putStrLn $ " t=" ++ show profT + putStrLn $ " tau2theta(-t)=" ++ show (tau2theta (-profT)) + ++ " tau2theta(+t)=" ++ show (tau2theta profT) + putStrLn $ " profile points:" + let tausL = U.toList taus + thetasL = U.toList (getCol i thetas) + forM_ (zip tausL thetasL) $ \(tau, th) -> + putStrLn $ " tau=" ++ show tau ++ " theta=" ++ show th + return $ paramCI (Profile stats profiles) nRows thetaOpt (raAlpha args) + + return $! ReportData + { rdTree = tree + , rdTheta = thetaOpt + , rdStdErr = _stdErr stats + , rdCriteria = crits + , rdCIs = cis + } + + case result of + Right rd -> printReport idx src rd + Left e -> printFailure idx src (show (e :: SomeException)) + +---------------------------------------------------------------------- +-- Evaluate a single ModelEval from base quantities +---------------------------------------------------------------------- +evalOne :: ModelEval -> Double -> Double -> Double -> Double -> Double + -> Double -> Double -> Double -> Double -> Double +evalOne RMSE mse _ _ _ _ _ _ _ _ = sqrt mse +evalOne R2 mse _ tss n _ _ _ _ _ = 1 - n * mse / tss +evalOne AIC _ nll _ _ k _ _ _ _ = 2*k + 2*nll +evalOne BIC _ nll _ n k _ _ _ _ = k * log n + 2*nll +evalOne Evidence _ nll _ n k _ _ _ _ = (1 - b) * nll - k/2 * log b + where b = 1 / sqrt n +evalOne FBF _ nll _ n k _ _ _ _ = res + where b = 1 / sqrt n; nup = exp (1 - log 3) + res = (1 - b) * nll - k/2 * log b + k/2 * log (2*pi*nup) +evalOne MDL _ nll _ _ _ logP _ logF _ = nll + logF + logP +evalOne MDLLatt _ nll _ _ _ _ logPL logF _ = nll + logF + logPL +evalOne MDLFreq _ nll _ _ _ logP _ _ logFF = nll + logFF + logP +evalOne (EvalLoss (NLL Gaussian)) _ nll _ _ _ _ _ _ _ = nll +evalOne _ _ _ _ _ _ _ _ _ _ = 0 -- unreachable + +---------------------------------------------------------------------- +-- Output +---------------------------------------------------------------------- +printReport :: Int -> B.ByteString -> ReportData -> IO () +printReport idx src rd = do + putStrLn $ "=== Expression " ++ show idx ++ " ===" + putStrLn $ "Tree: " ++ showExpr (rdTree rd) + putStrLn "Parameters:" + let thetaList = U.toList (rdTheta rd) + ciList = rdCIs rd + forM_ (zip3 [0..] thetaList ciList) $ \(i, th, ci) -> + putStrLn $ " theta" ++ show i ++ ": " ++ fmt th + ++ " [" ++ fmt (lower_ ci) ++ ", " ++ fmt (upper_ ci) ++ "]" + putStrLn "Model Selection:" + forM_ (rdCriteria rd) $ \(c, v) -> + putStrLn $ " " ++ padRight 12 (show c) ++ ": " ++ fmt v + putStrLn "" + where + fmt x | abs x < 1e-10 = "0.0000" + | abs x >= 1e4 = printf "%.4e" x + | otherwise = printf "%.6f" x + padRight n s = s ++ replicate (max 0 (n - length s)) ' ' + +printFailure :: Int -> B.ByteString -> String -> IO () +printFailure idx src msg = do + putStrLn $ "=== Expression " ++ show idx ++ " ===" + putStrLn $ "Tree: " ++ B.unpack src + putStrLn $ "Error: " ++ msg + putStrLn "" diff --git a/apps/srsimplify/Main.hs b/apps/srsimplify/Main.hs deleted file mode 100644 --- a/apps/srsimplify/Main.hs +++ /dev/null @@ -1,103 +0,0 @@ -module Main (main) where - -import Options.Applicative -import Text.ParseSR.IO ( withInput, withOutput ) -import Text.ParseSR ( SRAlgs (..), Output (..) ) -import System.Random ( getStdGen, mkStdGen ) -import Text.Read ( readMaybe ) -import Data.Char ( toLower, toUpper ) -import Data.List ( intercalate ) - --- Data type to store command line arguments -data Args = Args - { from :: SRAlgs - , to :: Output - , infile :: String - , outfile :: String - , varnames :: String - } deriving Show - --- parser of command line arguments -opt :: Parser Args -opt = Args - <$> option sralgsReader - ( long "from" - <> short 'f' - <> metavar ("[" <> intercalate "|" sralgsHelp <> "]") - <> help "Input expression format" ) - <*> option srtoReader -- TODO - ( long "to" - <> short 't' - <> metavar ("[" <> intercalate "|" srtoHelp <> "]") - <> help "Output expression format" ) - <*> strOption - ( long "input" - <> short 'i' - <> metavar "INPUT-FILE" - <> showDefault - <> value "" - <> help "Input file containing expressions. \ - \ Empty string gets expression from stdin." ) - <*> strOption - ( long "output" - <> short 'o' - <> metavar "OUTPUT-FILE" - <> showDefault - <> value "" - <> help "Output file to store the stats in CSV format. \ - \ Empty string prints expressions to stdout." ) - <*> strOption - ( long "varnames" - <> short 'v' - <> metavar "VARNAMES" - <> showDefault - <> value "" - <> help "Comma separated string of variable names. \ - \ Empty string defaults to the algorithm default (x0, x1,..)." ) - --- helper functions to show the possible options -mkDescription :: Show a => [a] -> [String] -mkDescription = map (envelope '\'' . map toLower . show) - where - envelope :: a -> [a] -> [a] - envelope c xs = c : xs <> [c] -{-# INLINE mkDescription #-} - -sralgsHelp :: [String] -sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..] -{-# INLINE sralgsHelp #-} - -srtoHelp :: [String] -srtoHelp = mkDescription [toEnum 0 :: Output ..] -{-# INLINE srtoHelp #-} - --- helper functions to parse the options -mkReader :: Read a => String -> (a -> b) -> String -> ReadM b -mkReader err val sr = eitherReader - $ case readMaybe sr of - Nothing -> pure (Left err) - Just x -> pure (Right (val x)) - -sralgsReader :: ReadM SRAlgs -sralgsReader = - str >>= (mkReader errMsg id . map toUpper) - where - errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp - -srtoReader :: ReadM Output -srtoReader = - str >>= (mkReader errMsg id . map toUpper) - where - errMsg = "unknown algorithm. Available options are " <> intercalate "," srtoHelp - -main :: IO () -main = do - args <- execParser opts - withInput (infile args) (from args) (varnames args) False True - >>= withOutput (outfile args) (to args) - where - opts = info (opt <**> helper) - ( fullDesc <> progDesc "Simplify an expression\ - \ using equality saturation." - <> header "srsimplify - a CLI tool to simplify\ - \ symbolic regression expressions with equality saturation." ) diff --git a/apps/srtools/Args.hs b/apps/srtools/Args.hs deleted file mode 100644 --- a/apps/srtools/Args.hs +++ /dev/null @@ -1,184 +0,0 @@ -module Args where - -import Data.Char ( toLower, toUpper ) -import Data.List ( intercalate ) -import Algorithm.SRTree.Likelihoods ( Distribution (..) ) -import Algorithm.SRTree.ConfidenceIntervals ( PType (..) ) -import Options.Applicative -import Text.ParseSR ( SRAlgs (..) ) -import Text.Read ( readMaybe ) - --- Data type to store command line arguments -data Args = Args - { from :: SRAlgs - , infile :: String - , outfile :: String - , dataset :: String - , test :: String - , niter :: Int - , hasHeader :: Bool - , simpl :: Bool - , dist :: Distribution - , restart :: Bool - , rseed :: Int - , toScreen :: Bool - , useProfile :: Bool - , simple :: Bool - , sigma :: Double - , alpha :: Double - , ptype :: PType - } deriving Show - --- parser of command line arguments -opt :: Parser Args -opt = Args - <$> option sralgsReader - ( long "from" - <> short 'f' - <> metavar ("[" <> intercalate "|" sralgsHelp <> "]") - <> help "Input expression format" ) - <*> strOption - ( long "input" - <> short 'i' - <> metavar "INPUT-FILE" - <> showDefault - <> value "" - <> help "Input file containing expressions. \ - \ Empty string gets expression from stdin." ) - <*> strOption - ( long "output" - <> short 'o' - <> metavar "OUTPUT-FILE" - <> showDefault - <> value "" - <> help "Output file to store the stats in CSV format. \ - \ Empty string prints expressions to stdout." ) - <*> strOption - ( long "dataset" - <> short 'd' - <> metavar "DATASET-FILENAME" - <> help "Filename of the dataset used for optimizing the parameters. \ - \ Empty string omits stats that make use of the training data. \ - \ 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: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 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" - <> showDefault - <> value "" - <> help "Filename of the test dataset.\ - \ Empty string omits stats that make use of the training data.\ - \ It can have additional information as in the training set,\ - \ but the validation range will be discarded." ) - <*> option auto - ( long "niter" - <> metavar "NITER" - <> showDefault - <> value 10 - <> help "Number of iterations for the optimization algorithm.") - <*> switch - ( long "hasheader" - <> help "Uses the first row of the csv file as header.") - <*> switch - ( long "simplify" - <> help "Apply basic simplification." ) - <*> option distRead - ( long "distribution" - <> metavar ("[" <> intercalate "|" distHelp <> "]") - <> showDefault - <> value Gaussian - <> help "Minimize negative log-likelihood following one of\ - \ the avaliable distributions.\ - \ The default is Gaussian." - ) - <*> switch - ( long "restart" - <> help "If set, it samples the initial values of\ - \ the parameters using a Gaussian distribution N(0, 1),\ - \ otherwise it uses the original values of the expression." ) - <*> option auto - ( long "seed" - <> metavar "SEED" - <> showDefault - <> value (-1) - <> help "Random seed to initialize the parameters values.\ - \ Used only if restart is enabled.") - <*> switch - ( long "report" - <> help "If set, reports the analysis in a user-friendly\ - \ format instead of csv. It will also include\ - \ confidence interval for the parameters and predictions" ) - <*> 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 "sigma" - <> metavar "SIGMA" - <> showDefault - <> value 0.001 - <> help "Estimation of error for Guassian distribution.") - <*> option auto - ( long "alpha" - <> metavar "ALPHA" - <> showDefault - <> value 0.05 - <> help "Significance level for confidence intervals.") - <*> option auto - ( long "ptype" - <> metavar "[Bates | ODE | Constrained]" - <> showDefault - <> value Constrained - <> help "Profile Likelihood method. Default: Constrained. NOTE: Constrained method only calculates the endpoint." - ) - --- helper functions to show the possible options -mkDescription :: Show a => [a] -> [String] -mkDescription = map (envelope '\'' . map toLower . show) - where - envelope :: a -> [a] -> [a] - envelope c xs = c : xs <> [c] -{-# INLINE mkDescription #-} - -sralgsHelp :: [String] -sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..] -{-# INLINE sralgsHelp #-} - -distHelp :: [String] -distHelp = mkDescription [toEnum 0 :: Distribution ..] -{-# INLINE distHelp #-} - --- helper functions to parse the options -mkReader :: Read a => String -> (a -> b) -> String -> ReadM b -mkReader err val sr = eitherReader - $ case readMaybe sr of - Nothing -> pure (Left err) - Just x -> pure (Right (val x)) - -sralgsReader :: ReadM SRAlgs -sralgsReader = - str >>= (mkReader errMsg id . map toUpper) - where - errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp - ---s2Reader :: ReadM (Maybe Double) ---s2Reader = --- str >>= \s -> mkReader ("wrong format " <> s) Just s - -distRead :: ReadM Distribution -distRead = - str >>= \s -> mkReader ("unsupported distribution " <> s) id (capitalize s) - where - capitalize "" = "" - capitalize (c:cs) = toUpper c : if length cs == 2 then map toUpper cs else map toLower cs diff --git a/apps/srtools/IO.hs b/apps/srtools/IO.hs deleted file mode 100644 --- a/apps/srtools/IO.hs +++ /dev/null @@ -1,230 +0,0 @@ -{-# language BlockArguments #-} -{-# language LambdaCase #-} -module IO where - -import System.IO ( hClose, hPutStrLn, openFile, stderr, stdout, IOMode(WriteMode), Handle ) -import qualified Data.Massiv.Array as A -import Data.List ( intercalate ) -import Control.Monad ( unless, forM_ ) -import System.Random ( StdGen ) - -import Data.SRTree ( SRTree (..), Fix (..), var, floatConstsToParam, relabelVars ) -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,dist,niter,sigma) ) -import Report -import Data.SRTree.Recursion ( cata ) - -import Debug.Trace ( trace, traceShow ) - --- Header of CSV file -csvHeader :: String -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 -openWriteWithDefault _ fname = openFile fname WriteMode -{-# INLINE openWriteWithDefault #-} - --- procecss a single tree and return all the available stats -processTree :: 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, Info, (BasicStats, [CI], [CI], [CI], [CI])) -processTree args seed dset t ix = (basic, sseOrig, sseOpt, info, cis) - where - (tree, theta0') = floatConstsToParam t - theta0 = if dist args == Gaussian - then theta0' <> [sigma args] - else theta0' - - basic = getBasicStats args seed dset tree theta0 ix - treeVal = case (_xVal dset, _yVal dset) of - (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) - 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 - theta0 = if dist args == Gaussian - then theta0' <> [sigma args] - else theta0' - - basic = getBasicStats args seed dset tree theta0 ix - treeVal = case (_xVal dset, _yVal dset) of - (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 - hStat <- openWriteWithDefault stdout (outfile args) - hPutStrLn hStat csvHeader - forM_ (zip [0..] exprs) - \(ix, tree) -> - case tree of - Left err -> hPutStrLn stderr ("invalid expression: " <> err) - Right t -> let treeData = processTree args seed dset t ix - 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) - 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] - 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 -getTransformedFeatures :: Fix SRTree -> (Fix SRTree, [Fix SRTree]) -getTransformedFeatures = cata $ - \case - Var ix -> (Fix $ Var ix, []) - Param ix -> (Fix $ Param ix, []) - Const x -> (Fix $ Const x, []) - Uni f (t, vars) -> (Fix $ Uni f t, vars) - Bin op (l, vs1) (r, vs2) -> case (hasNoParam l, hasNoParam r) of - (False, True) -> let vs = vs1 <> vs2 - in (Fix $ Bin op l (var $ length vs), vs <> [r]) - (True, False) -> let vs = vs1 <> vs2 - in (Fix $ Bin op (var $ length vs) r, vs <> [l]) - ( _, _) -> (Fix $ Bin op l r, vs1 <> vs2) -- vs1 == vs2 == [] - - where - hasNoParam = cata $ - \case - Var ix -> True - Param ix -> False - Const x -> if floor x == ceiling x then True else False - Uni f t -> t - Bin op l r -> l && r - -allAreVars :: [Fix SRTree] -> Bool -allAreVars = all isOnlyVar - where - isOnlyVar (Fix (Var _)) = True - isOnlyVar _ = False - --- print the information on screen (including CIs) -printResultsScreen :: Args -> StdGen -> Datasets -> [String] -> String -> [Either String (Fix SRTree)] -> IO () -printResultsScreen args seed dset varnames targt exprs = do - forM_ (zip [0..] exprs) - \(ix, tree) -> - case tree of - Left err -> do putStrLn ("invalid expression: " <> err) - Right t -> let treeData = processTree args seed dset t ix - in printToScreen ix treeData - where - decim :: Int -> Double -> Double - decim n x = (fromIntegral . (round :: Double -> Integer)) (x * 10^n) / 10^n - sdecim n = show . decim n - nplaces = 4 - - - printToScreen ix (basic, _, sseOpt, info, (sts, cis, pis_tr, pis_val, pis_te)) = - do let (transformedT, newvars) = getTransformedFeatures (_expr basic) - varnames' = ['z': show ix | ix <- [0 .. length newvars - 1]] - putStrLn $ "=================== EXPR " <> show ix <> " ==================" - putStr $ targt <> " ~ f(" <> intercalate ", " varnames <> ") = " - putStrLn $ P.showExprWithVars varnames (_expr basic) - - unless (allAreVars newvars) do - putStrLn "\nExpression and transformed features: " - putStr $ targt <> " ~ f(" <> intercalate ", " varnames' <> ") = " - putStrLn $ P.showExprWithVars varnames' (relabelVars transformedT) - forM_ (zip varnames' newvars) \(vn, tv) -> do - putStrLn $ vn <> " = " <> P.showExprWithVars varnames tv - - putStrLn "\n---------General stats:---------\n" - putStrLn $ "Number of nodes: " <> show (_nNodes basic) - putStrLn $ "Number of params: " <> show (_nParams basic) - putStrLn $ "theta = " <> show (_params basic) - - putStrLn "\n----------Performance:--------\n" - putStrLn $ "SSE (train.): " <> sdecim nplaces (_sseTr sseOpt) - putStrLn $ "SSE (val.): " <> sdecim nplaces (_sseVal sseOpt) - putStrLn $ "SSE (test): " <> sdecim nplaces (_sseTe sseOpt) - putStrLn $ "NegLogLiklihood (train.): " <> sdecim nplaces (_nllTr info) - putStrLn $ "NegLogLiklihood (val.): " <> sdecim nplaces (_nllVal info) - putStrLn $ "NegLogLiklihood (test): " <> sdecim nplaces (_nllTe info) - - putStrLn "\n------Selection criteria:-----\n" - putStrLn $ "BIC: " <> sdecim nplaces (_bic info) - putStrLn $ "AIC: " <> sdecim nplaces (_aic info) - putStrLn $ "MDL: " <> sdecim nplaces (_mdl info) - putStrLn $ "MDL (freq.): " <> sdecim nplaces (_mdlFreq info) - putStrLn $ "Functional complexity: " <> sdecim nplaces (_cc info) - putStrLn $ "Parameter complexity: " <> sdecim nplaces (_cp info) - - putStrLn "\n---------Uncertainties:----------\n" - putStrLn "Correlation of parameters: " - putStrLn $ show $ A.map (decim 2) (_corr sts) - putStrLn $ "Std. Err.: " <> show (A.map (decim nplaces) (_stdErr sts)) - putStrLn "\nConfidence intervals:\n\nlower <= val <= upper" - mapM_ (printCI nplaces) cis - putStrLn "\nConfidence intervals (predictions training):\n\nlower <= val <= upper" - mapM_ (printCI nplaces) pis_tr - unless (null pis_val) do - putStrLn "\nConfidence intervals (predictions validation):\n\nlower <= val <= upper" - mapM_ (printCI nplaces) pis_val - unless (null pis_te) do - putStrLn "\nConfidence intervals (predictions test):\n\nlower <= val <= upper" - mapM_ (printCI nplaces) pis_te - putStrLn "=============================================================" diff --git a/apps/srtools/Main.hs b/apps/srtools/Main.hs deleted file mode 100644 --- a/apps/srtools/Main.hs +++ /dev/null @@ -1,34 +0,0 @@ -module Main (main) where - -import Data.ByteString.Char8 ( pack, unpack, split ) -import Options.Applicative -import System.Random ( getStdGen, mkStdGen ) -import Text.ParseSR.IO ( withInput ) - -import Args -import IO -import Report - -main :: IO () -main = do - args <- execParser opts - g <- getStdGen - (dset, varnames, tgname) <- getDataset args - - let seed = if rseed args < 0 - then g - else mkStdGen (rseed args) - varnames' = map unpack $ split ',' $ pack varnames - withInput (infile args) (from args) varnames False (simpl args) - >>= if toScreen args - then printResultsScreen args seed dset varnames' tgname -- full report on screen - 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\ - \ Symbolic Regression expressions." - <> header "srtools - a CLI tool to (re)optimize the numeric\ - \ parameters of symbolic regression expressions" - ) diff --git a/apps/srtools/Report.hs b/apps/srtools/Report.hs deleted file mode 100644 --- a/apps/srtools/Report.hs +++ /dev/null @@ -1,280 +0,0 @@ -module Report where - -import qualified Data.Vector.Storable as VS -import qualified Data.Massiv.Array as A -import Data.Massiv.Array ( Sz(..) ) -import Data.Maybe ( fromMaybe ) -import Statistics.Distribution.FDistribution ( fDistribution ) -import Statistics.Distribution.ChiSquared ( chiSquared ) -import Statistics.Distribution ( quantile ) -import System.Random ( StdGen, split, randomRs ) - -import Data.SRTree ( SRTree, Fix (..), floatConstsToParam, paramsToConst, countNodes ) -import Data.SRTree.Eval -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) -import Data.SRTree.Datasets ( loadDataset ) -import Data.SRTree.Print ( showExpr ) -import Debug.Trace ( trace, traceShow ) - -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 - , _yErrTr :: Maybe PVector - , _yErrVal :: Maybe PVector - , _yErrTe :: Maybe PVector - } - --- basic fields name -basicFields :: [String] -basicFields = [ "Index" - , "Filename" - , "Expression" - , "Number_of_nodes" - , "Number_of_parameters" - , "Parameters" - , "Number_of_evaluations" - ] - --- basic information about the tree -data BasicInfo = Basic { _index :: Int - , _fname :: String - , _expr :: Fix SRTree - , _nNodes :: Int - , _nParams :: Int - , _params :: [Double] - , _nEvals :: Int - } - --- optimization fields -optFields :: [String] -optFields = [ "SSE_train_orig" - , "SSE_val_orig" - , "SSE_test_orig" - , "SSE_train_opt" - , "SSE_val_opt" - , "SSE_test_opt" - ] - --- optimization information -data SSE = SSE { _sseTr :: Double - , _sseVal :: Double - , _sseTe :: Double - } - --- model selection fields -modelFields :: [String] -modelFields = [ "BIC" - , "BIC_val" - , "AIC" - , "AIC_val" - , "Evidence" - , "EvidenceVal" - , "MDL" - , "MDL_Freq" - , "MDL_Lattice" - , "MDL_val" - , "MDL_Freq_val" - , "MDL_Lattice_val" - , "NegLogLikelihood_train" - , "NegLogLikelihood_val" - , "NegLogLikelihood_test" - , "LogFunctional" - , "LogParameters" - , "Fisher" - ] - --- model selection information -data Info = Info { _bic :: Double - , _bicVal :: Double - , _aic :: Double - , _aicVal :: Double - , _evidence :: Double - , _evidenceVal :: Double - , _mdl :: Double - , _mdlFreq :: Double - , _mdlLatt :: Double - , _mdlVal :: Double - , _mdlFreqVal :: Double - , _mdlLattVal :: Double - , _nllTr :: Double - , _nllVal :: Double - , _nllTe :: Double - , _cc :: Double - , _cp :: Double - , _fisher :: [Double] - } - --- load the datasets -getDataset :: Args -> IO (Datasets, String, String) -getDataset args = do - ((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, 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 nEvs - where - -- (tree', theta0) = floatConstsToParam tree - thetas = if restart args - then A.fromList compMode $ take nParams (randomRs (-1.0, 1.0) seed) - else A.fromList compMode theta0 - (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 - params = A.toList t - anyNaN = A.any isNaN t - -getSSE :: Datasets -> Fix SRTree -> SSE -getSSE dset tree = SSE tr val te - where - (t, th) = floatConstsToParam tree - tr = sse (_xTr dset) (_yTr dset) t (A.fromList compMode th) - val = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - (Just xVal, Just yVal) -> sse xVal yVal t (A.fromList compMode th) - te = case (_xTe dset, _yTe dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - (Just xTe, Just yTe) -> sse xTe yTe t (A.fromList compMode th) - -getInfo :: Args -> Datasets -> Fix SRTree -> Fix SRTree -> Info -getInfo args dset tree treeVal = - Info { _bic = bic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt - , _bicVal = bicVal - , _aic = aic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt - , _aicVal = aicVal - , _evidence = evidence dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt - , _evidenceVal = evidenceVal - , _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 - , _nllTr = nllTr - , _nllVal = nllVal - , _nllTe = nllTe - , _cc = logFunctional tOpt - , _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) - (xVal, yVal) = case (_xVal dset, _yVal dset) of - (Nothing, _) -> (xTr, yTr) - (_, Nothing) -> (xTr, yTr) - (Just a, Just b) -> (a, b) - (tOpt, thetaOpt_nosig) = floatConstsToParam tree - thetaOpt = if dist args == Gaussian - then thetaOpt_nosig <> [sigma args] - else thetaOpt_nosig - thetaOpt' = A.fromList compMode thetaOpt - - (tOptVal, thetaOptVal_nosig) = floatConstsToParam treeVal - thetaOptVal = if dist args == Gaussian - then thetaOptVal_nosig <> [sigma args] - else thetaOptVal_nosig - thetaOptVal' = A.fromList compMode thetaOptVal - - dist' = dist args - - 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' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal - aicVal = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - _ -> aic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal - evidenceVal = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - _ -> evidence dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal - nllVal = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - _ -> 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' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal - mdlFreqVal = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - _ -> mdlFreq dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal - mdlLattVal = case (_xVal dset, _yVal dset) of - (Nothing, _) -> 0.0 - (_, Nothing) -> 0.0 - _ -> 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' (_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) - where - (Sz n) = A.size yTr - (tree, _) = floatConstsToParam (_expr basic) - theta = _params basic - tau_max = (quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - 0.01)) - tau_max' = sqrt $ quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - alpha') - (xTr, yTr) = (_xTr dset, _yTr dset) - dist' = dist args - 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, _, _) = 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' (_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) - jac xss = forwardModeUniqueJac xss (A.fromList compMode theta) tree -- FIX - - estCIs = paramCI (Laplace stats') n (A.fromList compMode theta) 0.001 - cis = paramCI method n (A.fromList compMode theta) alpha' - - estPIS_tr = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' [] - estPIS_val = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' [] - estPIS_te = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' [] - - pis_tr = predictionCI method dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' estPIS_tr - pis_val = case (_xVal dset, _yVal dset) of - (Nothing, _) -> [] - (Just xVal, _) -> predictionCI method dist' predFun jac prof xVal tree (A.fromList compMode theta) alpha' estPIS_val - pis_te = case (_xTe dset, _yTe dset) of - (Nothing, _) -> [] - (Just xTe, _) -> predictionCI method dist' predFun jac prof xTe tree (A.fromList compMode theta) alpha' estPIS_te diff --git a/apps/tinygp/GP.hs b/apps/tinygp/GP.hs deleted file mode 100644 --- a/apps/tinygp/GP.hs +++ /dev/null @@ -1,252 +0,0 @@ -{-# LANGUAGE ImportQualifiedPost #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE BangPatterns #-} -module GP where - -import Data.SRTree -import Algorithm.SRTree.Opt -import Algorithm.SRTree.Likelihoods -import Data.SRTree.Print -import Data.SRTree.Eval -import Data.SRTree.Recursion ( cata ) -import System.Random -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, maximumBy ) -import qualified Data.Vector.Mutable as MV - -data Method = Grow | Full | BTC -type Rng a = StateT StdGen IO a - -type GenUni = Fix SRTree -> Fix SRTree -type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree -type FitFun = Individual -> Rng Individual - -data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: [PVector] } - -instance Show Individual where - show (Individual t f p) = showExpr t <> "," <> show f <> "," <> show p - -toss :: Rng Bool -toss = state random -{-# INLINE toss #-} - -randomRange :: (Ord val, Random val) => (val, val) -> Rng val -randomRange rng = state (randomR rng) -{-# INLINE randomRange #-} - -randomFrom :: [a] -> Rng a -randomFrom funs = do n <- randomRange (0, length funs - 1) - pure $ funs !! n -{-# INLINE randomFrom #-} - -randomFromV :: V.Vector a -> Rng a -randomFromV funs = do n <- randomRange (0, length funs - 1) - pure $ funs V.! n -{-# INLINE randomFromV #-} - -countNodes' :: Fix SRTree -> Int -countNodes' = cata alg - where - alg (Var _) = 1 - alg (Param _) = 1 - alg (Const _) = 0 - alg (Bin _ l r) = 1 + l + r - alg (Uni Abs t) = t - alg (Uni _ t) = 1 + t -{-# INLINE countNodes' #-} - - -randomTree :: HyperParams -> Bool -> Rng (Fix SRTree) -randomTree hp grow - | depth <= 1 || size <= 2 = randomFrom term - | (min_depth >= 0 || (depth > 2 && not grow)) && size > 2 = genNonTerm - | otherwise = genTermOrNon - where - min_depth = _minDepth hp - depth = _maxDepth hp - size = _maxSize hp - term = _term hp - nonterm = _nonterm hp - - genNonTerm = - do et <- randomFrom nonterm - case et of - Left uniT -> uniT <$> randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow - Right binT -> do l <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow - r <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1 - countNodes' l} grow - pure (binT l r) - genTermOrNon = do r <- toss - if r - then randomFrom term - else genNonTerm - -data HyperParams = - HP { _minDepth :: Int - , _maxDepth :: Int - , _maxSize :: Int - , _popSize :: Int - , _tournSize :: Int - , _pc :: Double - , _pm :: Double - , _term :: [Fix SRTree] - , _nonterm :: [Either GenUni GenBin] - } - -tournament :: HyperParams -> V.Vector Individual -> Rng Individual -tournament hp pop = do - selection <- replicateM (_tournSize hp) (randomFromV $ V.filter (not.isNaN._fit) pop) - let maxFitness = maximum (fmap _fit selection) - champions = V.filter ((== maxFitness) . _fit) pop - if null selection - then randomFromV pop - else randomFromV champions - -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)) - fitFun $ Individual t 0.0 [] -- (M.fromList compMode theta' :: PVector) - --pure ind - --if isInfinite (_fit ind) - -- then randomIndividual hyperparams fitFun grow - -- else pure ind - -initialPop :: HyperParams -> FitFun -> Rng (V.Vector Individual) -initialPop hyperparams fitFun = do - let depths = [3 .. _maxDepth hyperparams] - pop <- forM depths $ \md -> - do let m = _popSize hyperparams `div` (_maxDepth hyperparams - 3 + 1) - g = V.fromList . take m $ cycle [True, False] - mapM (randomIndividual hyperparams{ _maxDepth = md} fitFun) g - pure (V.concat pop) - -fitnessMV :: Distribution -> [(SRMatrix, PVector, Maybe PVector)] -> Individual -> Rng Individual -fitnessMV dist datas ind = do - fs <- forM datas (fitness dist ind) - let fitOpt = minimum $ map fst fs - pure ind{_fit = fitOpt, _params = map snd fs} - -fitness :: Distribution -> Individual -> (SRMatrix, PVector, Maybe PVector) -> Rng (Double, PVector) -fitness dist ind (x, y, e) = do - let tree = relabelParams $ _tree ind - p = countParams tree - theta1' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1)) - theta2' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1)) - let (theta1, f1, _) = minimizeNLL dist e 50 x y tree theta1' - (theta2, f2, _) = minimizeNLL dist e 50 x y tree theta2' - fit1 = if isNaN f1 then (-1.0/0.0) else negate f1 - fit2 = if isNaN f1 then (-1.0/0.0) else negate f2 - thetaOpt = if fit1 > fit2 then theta1 else theta2 - fitOpt = max fit1 fit2 - pure (fitOpt, thetaOpt) - - -mutate :: HyperParams -> Individual -> Rng (Maybe Individual) -mutate hp ind = do - let sz = countNodes' (_tree 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 [] - else pure Nothing - where - 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 (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 - 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 dist ind dataTrains dataTests = do - let tree = relabelParams $ _tree ind - thetas = _params ind - mseTrain = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTrains - mseTest = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTests - r2Train = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTrains - r2Test = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTests - thetaStr = intercalate "_" $ map (intercalate ";" . map show . M.toList) thetas - putStrLn "id,Expression,theta,size,MSE_train,MSE_test,R2_train,R2_test" - putStr $ "0," <> showExpr tree <> "," <> thetaStr <> "," <> show (countNodes tree) <> "," <> show mseTrain <> "," <> show mseTest <> "," <> show r2Train <> "," <> show r2Test - -report :: Int -> V.Vector Individual -> IO () -report gen = mapM_ reportOne - where reportOne ind = do putStr (show gen) - putStr ": " - putStr (showExpr (_tree ind)) - putStr " - " - putStr (show (_fit ind)) - putStr " " - print (map M.toList $ _params ind) -{-# INLINE report #-} - -evolution :: Int -> HyperParams -> FitFun -> Rng (Individual) -evolution gen hp fitFun = do - pop <- initialPop hp fitFun - --liftIO $ report (-1) pop - go gen pop - where - go 0 !pop = pure $ pop V.! 0 - go n !pop = do - let best = V.maximumOn _fit $ V.filter (not.isNaN._fit) 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' diff --git a/apps/tinygp/Initialization.hs b/apps/tinygp/Initialization.hs deleted file mode 100644 --- a/apps/tinygp/Initialization.hs +++ /dev/null @@ -1,50 +0,0 @@ -module Initialization where - -data InitiMethod = GROW | FULL | BTC | HALFHALF -{- -btc = undefined - -def btc(pset_, depth_, length_, type_=None): - if type_ is None: - type_ = pset_.ret - - expr = [] - - arities = list(map(lambda x: x.arity, pset_.primitives[type_])) - minFunctionArity = min(arities) - maxFunctionArity = max(arities) - - # adapt length to restrictions of the primitive set - if length_ % 2 == 0 and minFunctionArity > 1: - length_ = length_ + 1 if np.random.random_sample(1) > 0.5 else length_ - 1 - - targetLength = length_ - 1 # don't count the root node - maxFunctionArity = min(maxFunctionArity, targetLength) - minFunctionArity = min(minFunctionArity, targetLength) - root = sampleChild(pset_, minFunctionArity, maxFunctionArity, type_) - - # inner lists of the form [node, depth, childIndex] - # childIndex is only used at the end to transform - # the representation from breadth to prefix - expr.append([root, 0, 1]) - - openSlots = root.arity - - for i in range(0, length_): - (node, nodeDepth, childIndex) = expr[i] - childDepth = nodeDepth + 1 - - for j in range(0, getArity(node)): - maxArity = 0 if childDepth == depth_ - 1 else min(maxFunctionArity, targetLength - openSlots) - minArity = min(minFunctionArity, maxArity) - child = sampleChild(pset_, minArity, maxArity, type_) - - if j == 0: - expr[i][2] = len(expr) - - expr.append([child, childDepth, 0]) - openSlots += getArity(child) - - nodes = breadthToPrefix(expr) - return nodes - -} diff --git a/apps/tinygp/Main.hs b/apps/tinygp/Main.hs deleted file mode 100644 --- a/apps/tinygp/Main.hs +++ /dev/null @@ -1,116 +0,0 @@ -module Main (main) where - -import GP ( HyperParams(HP), fitnessMV, evolution, printFinal ) -import Data.SRTree -import System.Random ( getStdGen ) -import Control.Monad.State.Strict ( evalStateT ) -import Data.SRTree.Datasets ( loadDataset ) -import Options.Applicative -import Data.Massiv.Array -import Util -import Algorithm.SRTree.Likelihoods -import Data.SRTree.Datasets - --- Data type to store command line arguments -data Args = Args - { dataset :: String, - _testData :: String, - popSize :: Int, - gens :: Int, - _maxSize :: Int, - pc :: Double, - pm :: Double, - _nonterminals :: String, - _nTournament :: Int, - _distribution :: Distribution - } - 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" - <> value "" - <> metavar "INPUT-FILE" - <> help "CSV dataset." ) - <*> option auto - ( long "population" - <> short 'p' - <> metavar "POP-SIZE" - <> showDefault - <> value 100 - <> help "Population size." ) - <*> option auto - ( long "generations" - <> short 'g' - <> metavar "GENS" - <> showDefault - <> 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 - <> value 0.9 - <> help "Crossover probability." ) - <*> option auto - ( long "probMut" - <> metavar "PM" - <> 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." - ) - <*> option auto - ( long "distribution" - <> value MSE - <> showDefault - <> help "distribution of the data.") - -nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> Fix $ Bin PowerAbs l r), Left recip, Left log, Left exp, Left (\t -> Fix $ Uni SqrtAbs t)] - -main :: IO () -main = do - args <- execParser opts - g <- getStdGen - --(x, y, _) <- loadTrainingOnly (dataset args) True - --(x_test, y_test, _) <- loadTrainingOnly (_testData args) True - - let datasets = words (dataset args) - dataTrains <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets - dataTests <- if null (_testData args) - then pure dataTrains - else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args) - - let hp = HP 3 10 (_maxSize args) (popSize args) (_nTournament args) (pc args) (pm args) terms (parseNonTerms $ _nonterminals args) - (Sz2 _ nFeats) = size . getX $ head dataTrains - terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]] - best <- evalStateT (evolution (gens args) hp (fitnessMV (_distribution args) dataTrains)) g - printFinal (_distribution args) best dataTrains dataTests - where - opts = info (opt <**> helper) - ( fullDesc <> progDesc "Very simple example of GP using SRTree." - <> header "tinyGP - a very simple example of GP using SRTRee." ) diff --git a/apps/tinygp/Util.hs b/apps/tinygp/Util.hs deleted file mode 100644 --- a/apps/tinygp/Util.hs +++ /dev/null @@ -1,63 +0,0 @@ -{-# LANGUAGE BlockArguments #-} -{-# LANGUAGE TupleSections #-} - -module Util where - -import qualified Data.Map.Strict as Map -import Data.Massiv.Array as MA hiding (forM_, forM) -import Data.SRTree -import Data.SRTree.Eval -import Algorithm.SRTree.Opt -import Algorithm.EqSat.Egraph -import Algorithm.EqSat.Build -import Algorithm.EqSat.Info - -import Algorithm.SRTree.NonlinearOpt -import System.Random -import 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 - diff --git a/src/Algorithm/EqSat.hs b/src/Algorithm/EqSat.hs --- a/src/Algorithm/EqSat.hs +++ b/src/Algorithm/EqSat.hs @@ -24,150 +24,344 @@ import Data.Function (on) import Data.IntMap (IntMap) import qualified Data.IntMap as IntMap -import Data.List (intercalate, minimumBy) +import qualified Data.IntSet as IntSet +import Data.List (intercalate) import Data.Map (Map) import qualified Data.Map as Map import Data.Maybe (mapMaybe) import Data.SRTree import Data.HashSet (HashSet) import qualified Data.HashSet as Set -import Control.Monad ( zipWithM ) - -import Debug.Trace +import Control.Monad ( zipWithM, forM_ ) -- | The `Scheduler` stores a map with the banned iterations of a certain rule . -- TODO: make it more customizable. type Scheduler a = State (IntMap Int) a -- to avoid importing -fromJust :: Maybe a -> a -fromJust (Just x) = x -fromJust _ = error "fromJust called with Nothing" -{-# INLINE fromJust #-} - -- | runs equality saturation from an expression tree, -- a given set of rules, and a cost function. -- Returns the tree with the smallest cost. -eqSat :: Monad m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) +eqSat :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) eqSat expr rules costFun maxIt = do root <- fromTree costFun expr - (end, it) <- runEqSat costFun rules maxIt - 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 modify' (const emptyGraph) >> eqSat best rules costFun it -- reapplies eqsat on the best so far - else pure best + _ <- runEqSat costFun rules maxIt + recalculateBest costFun root type CostMap = Map EClassId (Int, Fix SRTree) -- | recalculates the costs with a new cost function -recalculateBest :: Monad m => CostFun -> EClassId -> EGraphST m (Fix SRTree) +recalculateBest :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree) recalculateBest costFun eid = - do classes <- gets _eClass - let costs = fillUpCosts classes Map.empty + do ecls <- allClasses + let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls] + costs = fillUpCosts classes Map.empty eid' <- canonical eid - pure $ snd $ costs Map.! eid' + case Map.lookup eid' costs of + Just (_, t) -> pure t + Nothing -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid' + <> " nClasses=" <> show (IntMap.size classes) + <> " costSize=" <> show (Map.size costs) where - nodeCost :: CostMap -> ENode -> Maybe (Int, Fix SRTree) + nodeCost :: CostMap -> ENode -> (Int, Fix SRTree) nodeCost costMap enode = - do optChildren <- traverse (costMap Map.!?) (childrenOf enode) -- | gets the cost of the children, if one is missing, returns Nothing - let cc = map fst optChildren - nc = map snd optChildren - n = replaceChildren cc enode - c = costFun n - pure (c + sum cc, Fix $ replaceChildren nc enode) -- | otherwise, returns the cost of the node + children and the expression so far - - minimumBy' f [] = Nothing - minimumBy' f xs = Just $ minimumBy f xs + -- A child that has not been costed yet (a cycle, or a class whose + -- cost is computed later in this iteration) contributes a large + -- sentinel instead of 0: a 0 placeholder is cheaper than the real + -- cost, so the fixpoint below would keep the stale placeholder tree + -- (e.g. `x * 0.0` for `x * (y + z)`). Real costs always beat it. + let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (costMap Map.!? cid) | cid <- eChildren enode ] + c = case enode of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + _ -> costFun (replaceChildren cc (fromENode enode)) + in (c + sum cc, Fix $ case enode of + ENAry op _ -> unfix (naryTree op nc) + _ -> replaceChildren nc (fromENode enode)) -- | missing children (cyclic classes) get cost 0 so every class is costed + costSentinel :: Int + costSentinel = 1000000 fillUpCosts :: IntMap EClass -> CostMap -> CostMap - fillUpCosts classes m = - case IntMap.foldrWithKey costOfClass (False, m) classes of -- applies costOfClass to each class - (False, _) -> m - (True, m') -> fillUpCosts classes m' -- | if something changed, recurse + fillUpCosts classes = go (IntMap.size classes + 1) (IntMap.keysSet classes) + where + go 0 _ m = m + go n dirty m + | IntSet.null dirty = m + | otherwise = go (n - 1) dirty' m' + where + (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty + step (d, cm) eid = case IntMap.lookup eid classes of + Nothing -> (d, cm) + Just ecl -> + let currentCost = Map.lookup eid cm + minCost = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c') + ) Nothing (_eNodes ecl) + (changed, cm') = case (currentCost, minCost) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, Map.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, Map.insert eid new cm) + d' = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl) + else d + in d' `seq` cm' `seq` (d', cm') - costOfClass :: EClassId -> EClass -> (Bool, CostMap) -> (Bool, CostMap) - costOfClass eid ecl (b, m) = - let currentCost = m Map.!? eid - minCost = minimumBy' (compare `on` fst) -- get the minimum available cost of the nodes of this class - $ mapMaybe (nodeCost m) - $ map decodeEnode - $ Set.toList (_eNodes ecl) - in case (currentCost, minCost) of -- replace the costs accordingly - (_, Nothing) -> (b, m) - (Nothing, Just new) -> (True, Map.insert eid new m) - (Just old, Just new) -> if fst old <= fst new - then (b, m) - else (True, Map.insert eid new m) +-- | Recompute every e-class's cost-minimal @_best@/_cost@ bottom-up and write +-- it back into the graph. Needed after loading a graph whose best/cost were +-- not persisted (e.g. via srtree-db), where @_best@ may otherwise hold an +-- arbitrary (potentially large) e-node. +recalculateBestAll :: ClassStore m => CostFun -> EGraphST m () +recalculateBestAll costFun = do + ecls <- allClasses + let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls] + bests = fixpoint classes IntMap.empty + forM_ (IntMap.toList bests) $ \(eid, (c, en)) -> + case IntMap.lookup eid classes of + Nothing -> pure () + Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } } + where + nodeCost :: IntMap (Int, ENode) -> ENode -> (Int, ENode) + nodeCost cm en = + let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc + _ -> costFun (replaceChildren cc (fromENode en)) + sum cc + in (c, en) + costSentinel :: Int + costSentinel = 1000000 + fixpoint :: IntMap EClass -> IntMap (Int, ENode) -> IntMap (Int, ENode) + fixpoint classes0 = go (IntMap.size classes0 + 1) (IntMap.keysSet classes0) + where + go 0 _ m = m + go n dirty m + | IntSet.null dirty = m + | otherwise = go (n - 1) dirty' m' + where + (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty + step (d, cm) eid = case IntMap.lookup eid classes0 of + Nothing -> (d, cm) + Just ecl -> + let current = IntMap.lookup eid cm + minNode = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minNode) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, IntMap.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, IntMap.insert eid new cm) + d' = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl) + else d + in d' `seq` cm' `seq` (d', cm') + +-- | Like 'recalculateBestAll' but streamed: each e-class body is fetched on +-- demand through 'ClassStore' (so a paged graph never materializes every body +-- at once) and only the small @(cost, best e-node)@ map is kept resident. The +-- structural worklist fixpoint is identical. +recalculateBestAllStream :: ClassStore m => CostFun -> EGraphST m () +recalculateBestAllStream costFun = do + ids <- allKeys + let idSet = IntSet.fromList ids + costSentinel = 1000000 + nodeCost cm en = + let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc + _ -> costFun (replaceChildren cc (fromENode en)) + sum cc + in (c, en) + stepEid cm eid = do + mec <- readDirect eid + case mec of + Nothing -> pure (IntSet.empty, cm) + Just ecl -> do + let current = IntMap.lookup eid cm + minNode = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minNode) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, IntMap.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, IntMap.insert eid new cm) + dirty = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) IntSet.empty (_parents ecl) + else IntSet.empty + pure (dirty, cm') + fixpoint n dirty cm + | n <= 0 || IntSet.null dirty = pure cm + | otherwise = go (IntSet.toList dirty) IntSet.empty cm + where + go [] d acc = fixpoint (n - 1) d acc + go (e : es) d acc = do + (d', m') <- stepEid acc e + go es (IntSet.union d d') m' + cm <- fixpoint (IntSet.size idSet + 1) idSet IntMap.empty + forM_ (IntMap.toList cm) $ \(eid, (c, en)) -> do + mec <- readDirect eid + case mec of + Nothing -> pure () + Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } } + +-- | Streaming variant of 'recalculateBest': computes the cost-minimal tree for a +-- single root without materializing every e-class body at once. +recalculateBestStream :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree) +recalculateBestStream costFun eid = do + ids <- allKeys + let idSet = IntSet.fromList ids + costSentinel = 1000000 + nodeCost cm en = + let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (Map.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + _ -> costFun (replaceChildren cc (fromENode en)) + in (c + sum cc, Fix $ case en of + ENAry op _ -> unfix (naryTree op nc) + _ -> replaceChildren nc (fromENode en)) + stepEid cm eid' = do + mec <- lookupClass eid' + case mec of + Nothing -> pure (IntSet.empty, cm) + Just ecl -> do + let current = Map.lookup eid' cm + minCost = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minCost) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, Map.insert eid' new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, Map.insert eid' new cm) + dirty = if changed + then Set.foldl' (\acc (pid,_) -> IntSet.insert pid acc) IntSet.empty (_parents ecl) + else IntSet.empty + pure (dirty, cm') + fixpoint n dirty cm + | n <= 0 || IntSet.null dirty = pure cm + | otherwise = go (IntSet.toList dirty) IntSet.empty cm + where + go [] d acc = fixpoint (n - 1) d acc + go (e : es) d acc = do + (d', m') <- stepEid acc e + go es (IntSet.union d d') m' + cm <- fixpoint (IntSet.size idSet + 1) idSet Map.empty + eid' <- canonical eid + case Map.lookup eid' cm of + Just (_, t) -> pure t + Nothing -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid' + <> " costSize=" <> show (Map.size cm) + +-- | Run equality saturation and stream the final extraction (see +-- 'recalculateBestStream'), so a paged graph is never fully materialized. +eqSatStream :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) +eqSatStream expr rules costFun maxIt = do + root <- fromTree costFun expr + _ <- runEqSat costFun rules maxIt + recalculateBestAllStream costFun + recalculateBestStream costFun root + +-- | replaces the equality rules with two one-way rules +replaceEqRules :: Rule -> [Rule] +replaceEqRules (p1 :=> p2) = [p1 :=> p2] +replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1] +replaceEqRules (r :| cond) = map (:| cond) $ replaceEqRules r + +-- | Compile a rule source into a query, or `Nothing` for n-ary patterns that +-- use the direct multiset matcher instead. +compileSource :: Rule -> Maybe (Query, [ClassOrVar], ClassOrVar) +compileSource r = if hasNAry (source r) + then Nothing + else Just (compileToQuery (source r)) + +-- | Cap on the total number of rule matches applied in a single eqsat +-- iteration. Combined with the per-rule caps ('ruleBudget'/'ruleRootVisit' for +-- n-ary, 'ruleMatchBudget' for the cached path) and the persistent +-- mark-on-attempt seen-set (which makes each rule's budget advance to new +-- matches), this bounds a single iteration's apply/rebuild work regardless of +-- graph size. +iterMatchBudget :: Int +iterMatchBudget = 2000 + -- | run equality saturation for a number of iterations -runEqSat :: Monad m => CostFun -> [Rule] -> Int -> EGraphST m (Bool, Int) -runEqSat costFun rules maxIter = go maxIter IntMap.empty +runEqSat :: ClassStore m => CostFun -> [Rule] -> Int -> EGraphST m (Bool, Int) +runEqSat costFun rules maxIter = go maxIter IntMap.empty compiledRules where rules' = concatMap replaceEqRules rules - - -- replaces the equality rules with two one-way rules - replaceEqRules :: Rule -> [Rule] - replaceEqRules (p1 :=> p2) = [p1 :=> p2] - replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1] - replaceEqRules (r :| cond) = map (:| cond) $ replaceEqRules r + compiledRules = map (\r -> (r, compileSource r)) rules' - go it sch = do eNodes <- gets _eNodeToEClass - eClasses <- gets _eClass - --createDB -- TODO: partial db is still incomplete - --db <- gets (_patDB . _eDB) -- createDB -- creates the DB + go it sch compiled = + do -- reset dirty flag before processing this iteration + modify' $ over (eDB . changed) (const False) - -- step 1: match the rules - let matchSch = matchWithScheduler it - matchAll = zipWithM matchSch [0..] - (rules, sch') = runState (matchAll rules') sch + -- step 1: match the rules using cached compiled queries + let matchSch = matchWithScheduler it + adapted i (r, cq) = map (,cq) <$> matchSch i r + matchAll = zipWithM adapted [0..] + (filtered, sch') = runState (matchAll compiled) sch - -- step 2: apply matches and rebuild - matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rules - mapM_ (uncurry (applyMatch costFun)) $ concat matches - rebuild costFun + -- step 2: apply matches and rebuild + matches <- mapM (\(rule, cq) -> map (rule,) <$> case cq of + Just q -> do paged <- isPagedGraph + if paged + then matchStreamCached (Just (show (source rule))) (source rule) + else matchCachedWith (Just (show (source rule))) q + Nothing -> matchSaturated (source rule)) $ concat filtered + -- bound the total number of matches applied per iteration so a + -- single iteration's apply/rebuild work stays bounded on huge + -- graphs (genuine matches; we just process them over more iters). + mapM_ (uncurry (applyMatch costFun)) (take iterMatchBudget (concat matches)) + rebuild costFun - -- recalculate heights - --calculateHeights - eNodes' <- gets _eNodeToEClass - eClasses' <- gets _eClass + -- check dirty flag: if no modifications occurred, we've saturated + changed <- gets (_changed . _eDB) + if it == 1 || not changed + then pure (True, it) + else + do eClasses <- gets _eClass + if IntMap.size eClasses > 1500 + then throttle it sch' compiled + else go (it-1) sch' compiled - -- if nothing changed, return - if it == 1 || (eNodes' == eNodes && eClasses' == eClasses) - then pure (True, it) - else if IntMap.size eClasses' > 1500 -- maximum allowed number of e-classes. TODO: customize - then pure (False, it) - else go (it-1) sch' + throttle it sch compiled = do + cleanMaps + eClasses <- gets _eClass + if IntMap.size eClasses <= 1500 + then go (it-1) sch compiled + else do applySingleMergeOnlyEqSat costFun rules + changed <- gets (_changed . _eDB) + if it <= 1 || not changed + then pure (False, it) -- give up and return early stop + else throttle (it-1) sch compiled -- | apply a single step of merge-only equality saturation -applySingleMergeOnlyEqSat :: Monad m => CostFun -> [Rule] -> EGraphST m () +applySingleMergeOnlyEqSat :: ClassStore m => CostFun -> [Rule] -> EGraphST m () applySingleMergeOnlyEqSat costFun rules = - do db <- gets (_patDB . _eDB) -- createDB - let matchSch = matchWithScheduler 10 + do let matchSch = matchWithScheduler 10 matchAll = zipWithM matchSch [0..] - (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 + (rls, _) = runState (matchAll rules') IntMap.empty matches <- getNMatches 500 rls rebuild costFun - -- recalculate heights - --calculateHeights where rules' = concatMap replaceEqRules rules - -- replaces the equality rules with two one-way rules - replaceEqRules :: Rule -> [Rule] - 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 + getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> matchSaturated (source r) + let (x, _) = splitAt n matches m = length x if m == n then pure matches @@ -179,7 +373,7 @@ matchWithScheduler :: Int -> Int -> Rule -> Scheduler [Rule] -- [(Rule, (Map ClassOrVar ClassOrVar, ClassOrVar))] matchWithScheduler it ruleNumber rule = do mbBan <- gets (IntMap.!? ruleNumber) - if mbBan /= Nothing && fromJust mbBan <= it -- check if the rule is banned + if maybe False (<= it) mbBan -- check if the rule is banned then pure [] else do -- let matches = match db (source rule) modify (IntMap.insert ruleNumber (it+5)) diff --git a/src/Algorithm/EqSat/Build.hs b/src/Algorithm/EqSat/Build.hs --- a/src/Algorithm/EqSat/Build.hs +++ b/src/Algorithm/EqSat/Build.hs @@ -20,76 +20,161 @@ import System.Random (Random (randomR), StdGen) import Control.Lens ( over ) import Control.Monad ( forM_, when, foldM, forM ) -import Data.Maybe ( fromMaybe, catMaybes ) +import Data.Maybe import Data.SRTree import Algorithm.EqSat.Egraph ---import Algorithm.EqSat.Info import Algorithm.EqSat.DB import qualified Data.IntMap.Strict as IntMap +import Data.IntMap.Strict (IntMap) import Data.Map.Strict ( Map ) import qualified Data.Map.Strict as Map +import qualified Data.HashMap.Strict as HashMap import qualified Data.HashSet as Set import Control.Monad.State.Strict import Control.Monad.Identity +import GHC.Stack (HasCallStack) + import Data.SRTree.Recursion (cataM) +import Data.List (sort) import Algorithm.EqSat.Info import qualified Data.IntSet as IntSet -import Data.Maybe -import Data.Sequence (Seq(..), (><)) -import Data.List ( nub ) -import Debug.Trace (trace, traceShow) +import qualified Data.Set as RangeSet + + -- | adds a new or existing e-node (merging if necessary) -add :: Monad m => CostFun -> ENode -> EGraphST m EClassId -add costFun enode = - do enode'' <- canonize enode -- canonize e-node - constEnode <- calculateConsts 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'' +add :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m EClassId +add costFun enode = do + enode'' <- canonize enode + enode''' <- foldConsts costFun enode'' - maybeEid <- gets ((Map.!? enode') . _eNodeToEClass) -- check if canonical e-node exists - case maybeEid of + maybeEid <- lookupNode enode''' + case maybeEid of Just eid -> pure eid Nothing -> do curId <- gets (_nextId . _eDB) -- get the next available e-class id - modify' $ over canonicalMap (IntMap.insert curId curId) -- insert e-class id into canon map - . over eNodeToEClass (Map.insert enode' curId) -- associate new e-node with id - . over (eDB . nextId) (+1) -- update next id - . over (eDB . worklist) (Set.insert (curId, enode')) -- add e-node and id into worklist - forM_ (childrenOf enode') (addParents curId enode') -- update the children's parent list - info <- makeAnalysis costFun enode' - h <- getChildrenMinHeight enode' - let newClass = createEClass curId enode' info h -- create e-class - modify' $ over eClass (IntMap.insert curId newClass) -- insert new e-class into e-graph + insertCanonical curId curId -- register the class as its own representative + insertNode enode''' curId -- associate new e-node with id (bounded on paged graphs) + modify' $ over (eDB . nextId) (+1) -- update next id + . over (eDB . worklist) (Set.insert (curId, enode''')) -- add e-node and id into worklist + forM_ (eChildren enode''') (addParents curId enode''') -- update the children's parent list + info <- makeAnalysis costFun enode''' + h <- getChildrenMinHeight enode''' + let newClass = createEClass curId enode''' info h -- create e-class + -- insert via 'insertClass' so a paged (DB-backed) class store also + -- persists the new class's page; for a pure graph this is identical + -- to inserting into @_eClass@ directly. + insertClass newClass --modifyEClass costFun curId -- simplify eclass if it evaluates to a number -- update database - addToDB enode' curId -- add new node to db - modify' $ over (eDB . sizeDB) - $ IntMap.insertWith (IntSet.union) (_size info) (IntSet.singleton curId) + addToDB enode''' curId -- add new node to db + tracking <- gets (_trackDBs . _eDB) + when tracking $ + modify' $ over (eDB . sizeDB) + $ IntMap.insertWith (IntSet.union) (_size info) (IntSet.singleton curId) modify' $ over (eDB . unevaluated) (IntSet.insert curId) + . over (eDB . changed) (const True) pure curId where - addParents :: Monad m => EClassId -> ENode -> EClassId -> EGraphST m () + addParents :: ClassStore m => EClassId -> ENode -> EClassId -> EGraphST m () addParents cId node c = do ec <- getEClass c let ec' = ec{ _parents = Set.insert (cId, node) (_parents ec) } - modify' $ over eClass (IntMap.insert c ec') + -- write through 'insertClass' so a paged store keeps the updated parents + insertClass ec' +-- | Add a binary (SRTree-based) node, converting it to a flattened ENode. +-- Sub and Div are canonicalized away at insertion: `x - y` becomes +-- `x + (-1)*y` and `x / y` becomes `x * recip y`, so no Sub/Div e-node ever +-- enters the e-graph and the Sub/Div-aware rules become redundant. +addTree :: (ClassStore m, HasCallStack) => CostFun -> SRTree EClassId -> EGraphST m EClassId +addTree costFun (Bin Sub l r) = do + neg <- addNegate costFun r + add costFun =<< mkENary EAdd [l, neg] +addTree costFun (Bin Div l r) = do + rec <- add costFun (EUni Recip r) + add costFun =<< mkENary EMul [l, rec] +addTree costFun t = toENode t >>= add costFun +{-# INLINE addTree #-} + +-- | builds the e-class for the negation of the e-class `t`, represented as +-- `(-1) * t` (matching the pattern-level `negate` encoding in Algorithm.EqSat.DB). +addNegate :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId +addNegate costFun t = do + negOne <- add costFun (EConst (-1)) + add costFun =<< mkENary EMul [negOne, t] + +-- | Fused 'calculateConsts' + 'foldConstants': fetches each child's constant +-- info a single time, detects fully-constant nodes (replaced by EConst/EParam) +-- and folds together all-but-one constant children of an ENAry +-- (e.g. 2+3+x becomes 5+x). Constants that are already folded single subtrees +-- are handled by the same child-constant walk. +foldConsts :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode +foldConsts _ en@(ENAry _ m) | IntMap.null m = pure en +foldConsts costFun en@(ENAry op m) = do + let xs = expandedList m + infos <- mapM (fmap (_consts . _info) . getEClass) xs + case foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) infos of + ConstVal x -> pure (EConst x) + ParamIx x -> pure (EParam x) + _ -> foldENary costFun op m infos +foldConsts _ en = do + infos <- mapM (fmap (_consts . _info) . getEClass) (eChildren en) + case combineConsts (replaceChildren infos (fromENode en)) of + ConstVal x -> pure (EConst x) + ParamIx x -> pure (EParam x) + _ -> pure en +{-# INLINE foldConsts #-} + +-- | Fold together all-but-one constant children of an ENAry multiset. +foldENary :: (ClassStore m, HasCallStack) => CostFun -> NOp -> IntMap Int -> [Consts] -> EGraphST m ENode +foldENary costFun op m infos = do + let xs = expandedList m + (consts, rest) = foldr step ([], []) (zip xs infos) + step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs) + step (x, _) (cs, rs) = (cs, x:rs) + if length consts >= 2 + then do + let folded = case op of + EAdd -> sum consts + EMul -> product consts + if isNaN folded || isInfinite folded + then pure (ENAry op m) + else do + cid <- add costFun (EConst folded) + pure (ENAry op (imFromList (cid : rest))) + else pure (ENAry op m) +{-# INLINE foldENary #-} + +-- | Fold together all-but-one constant children of an ENAry at insertion +-- time (e.g. 2+3+x becomes 5+x). Constants that are already folded +-- single subtrees are handled by 'calculateConsts' above; this handles the +-- flattened case where several constant terms land in one multiset. +foldConstants :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode +foldConstants _ en@(ENAry _ m) | IntMap.size m < 2 = pure en +foldConstants costFun en@(ENAry op m) = do + let xs = expandedList m + infos <- mapM (fmap (_consts . _info) . getEClass) xs + let (consts, rest) = foldr step ([], []) (zip xs infos) + step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs) + step (x, _) (cs, rs) = (cs, x:rs) + if length consts >= 2 + then do + let folded = case op of + EAdd -> sum consts + EMul -> product consts + if isNaN folded || isInfinite folded + then pure en + else do + cid <- add costFun (EConst folded) + pure (ENAry op (imFromList (cid : rest))) + else pure en +foldConstants _ en = pure en + -- | rebuilds the e-graph after inserting or merging -- e-classes -rebuild :: Monad m => CostFun -> EGraphST m () +rebuild :: (ClassStore m, HasCallStack) => CostFun -> EGraphST m () rebuild costFun = do wl <- gets (_worklist . _eDB) al <- gets (_analysis . _eDB) @@ -102,21 +187,23 @@ -- | repairs e-node by canonizing its children -- if the canonized e-node already exists in -- e-graph, merge the e-classes -repair :: Monad m => CostFun -> EClassId -> ENode -> EGraphST m () +repair :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> ENode -> EGraphST m () repair costFun ecId enode = - do modify' $ over eNodeToEClass (Map.delete enode) + do modify' $ over eNodeToEClass (HashMap.delete enode) enode' <- canonize enode ecId' <- canonical ecId - doExist <- gets ((Map.!? enode') . _eNodeToEClass) + doExist <- lookupNode enode' case doExist of Just ecIdCanon -> do mergedId <- merge costFun ecIdCanon ecId' - modify' $ over eNodeToEClass (Map.insert enode' mergedId) - Nothing -> modify' $ over eNodeToEClass (Map.insert enode' ecId') + insertNode enode' mergedId + addToDB enode' mergedId + Nothing -> do insertNode enode' ecId' + addToDB enode' ecId' {-# INLINE repair #-} -- | repair the analysis of the e-class -- considering the new added e-node -repairAnalysis :: Monad m => CostFun -> EClassId -> ENode -> EGraphST m () +repairAnalysis :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> ENode -> EGraphST m () repairAnalysis costFun ecId enode = do ecId' <- canonical ecId enode' <- canonize enode @@ -125,15 +212,17 @@ let newData = joinData (_info eclass) info eclass' = eclass { _info = newData } when (_info eclass /= newData) $ - do modify' $ over (eDB . analysis) (_parents eclass <>) - . over eClass (IntMap.insert ecId' eclass') - . over (eDB . refits) (Set.insert ecId') + do let bestChanged = _best (_info eclass) /= _best newData + modify' $ over (eDB . analysis) (_parents eclass <>) + . (if bestChanged && isJust (_fitness (_info eclass)) then over (eDB . refits) (IntSet.insert ecId') else id) + -- write through 'insertClass' so a paged store keeps the updated body + insertClass eclass' _ <- modifyEClass costFun ecId' pure () {-# INLINE repairAnalysis #-} -- | merge to equivalent e-classes -merge :: Monad m => CostFun -> EClassId -> EClassId -> EGraphST m EClassId +merge :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EClassId -> EGraphST m EClassId merge costFun c1 c2 = do c1' <- canonical c1 c2' <- canonical c2 @@ -142,38 +231,43 @@ else do (led, ledC, ledOrig, sub, subC, subOrig) <- getLeaderSub c1' c1 c2' c2 -- the leader will be the e-class with more parents mergeClasses led ledC ledOrig sub subC subOrig -- merge sub into leader where - mergeClasses :: Monad m => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId + mergeClasses :: (ClassStore m, HasCallStack) => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId mergeClasses led ledC ledO sub subC subO = - 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) - (_parents ledC <> _parents subC) - (min (_height ledC) (_height subC)) - (joinData (_info ledC) (_info subC)) - - modify' $ over eClass (IntMap.insert led newC . IntMap.delete sub) -- delete sub e-class and replace leader - . 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) + do insertCanonical sub led -- persist/register the canonical merges + insertCanonical subO led + let newC = EClass led + (_eNodes ledC `Set.union` _eNodes subC) + (_parents ledC <> _parents subC) + (min (_height ledC) (_height subC)) + (joinData (_info ledC) (_info subC)) + forM_ (Set.toList (_eNodes subC)) $ \en -> insertNode en led + -- write the merged body through the class store (a paged store keeps the + -- authoritative page) and drop the absorbed class + insertClass newC + deleteClass sub + modify' $ over (eDB . worklist) (_parents subC <>) + when (_info newC /= _info ledC) + $ do let bestChanged = _best (_info newC) /= _best (_info ledC) + modify' $ over (eDB . analysis) (_parents ledC <>) + . (if bestChanged && isJust (_fitness (_info ledC)) then over (eDB . refits) (IntSet.insert led) else id) when (_info newC /= _info subC) $ modify' $ over (eDB . analysis) (_parents subC <>) - updateDBs newC led ledC ledO sub subC subO + tracking <- gets (_trackDBs . _eDB) + when tracking $ updateDBs newC led ledC ledO sub subC subO modifyEClass costFun led - --forM_ (_eNodes newC) $ \en -> addToDB (decodeEnode en) led + modify' $ over (eDB . changed) (const True) pure led getLeaderSub c1 c1O c2 c2O = do ec1 <- getEClass c1 ec2 <- getEClass c2 - let n1 = length (_parents ec1) - n2 = length (_parents ec2) + let n1 = Set.size (_parents ec1) + n2 = Set.size (_parents ec2) pure $ if n1 >= n2 then (c1, ec1, c1O, c2, ec2, c2O) else (c2, ec2, c2O, c1, ec1, c1O) - updateDBs :: Monad m => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m () + updateDBs :: (ClassStore m, HasCallStack) => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m () updateDBs newC led ledC ledO sub subC subO = do updateFitnessDB newC led ledC ledO sub subC subO updateSizeDB newC led ledC ledO sub subC subO @@ -188,20 +282,20 @@ updateFitnessDB :: Monad m => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m () updateFitnessDB newC led ledC ledO sub subC subO = - if (isJust fitNew) - then do - when (fitNew /= fitLed) $ do - if isNothing fitLed - then modify' $ over (eDB . unevaluated) (IntSet.delete led . IntSet.delete ledO) - else modify' $ over (eDB . fitRangeDB) (removeRange led (fromJust fitLed) . removeRange ledO (fromJust fitLed)) - . over (eDB . sizeFitDB) (IntMap.adjust (removeRange ledO (fromJust fitLed) . removeRange led (fromJust fitLed)) szLed) - modify' $ over (eDB . fitRangeDB) (insertRange led (fromJust fitNew)) - . over (eDB . sizeFitDB) (IntMap.adjust (insertRange led (fromJust fitNew)) szNew . IntMap.insertWith (><) szNew Empty) - if isNothing fitSub - then modify' $ over (eDB . unevaluated) (IntSet.delete sub . IntSet.delete subO) - else modify' $ over (eDB . fitRangeDB) (removeRange sub (fromJust fitSub) . removeRange subO (fromJust fitSub)) - . over (eDB . sizeFitDB) (IntMap.adjust (removeRange subO (fromJust fitSub) . removeRange sub (fromJust fitSub)) szSub) - else modify' $ over (eDB . unevaluated) (IntSet.insert led . IntSet.delete ledO . IntSet.delete sub . IntSet.delete subO) + case fitNew of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.insert led . IntSet.delete ledO . IntSet.delete sub . IntSet.delete subO) + Just fn -> do + when (fitNew /= fitLed) $ do + modify' $ case fitLed of + Nothing -> over (eDB . unevaluated) (IntSet.delete led . IntSet.delete ledO) + Just fl -> over (eDB . fitRangeDB) (removeRange led fl . removeRange ledO fl) + . over (eDB . sizeFitDB) (IntMap.adjust (removeRange ledO fl . removeRange led fl) szLed) + modify' $ over (eDB . fitRangeDB) (insertRange led fn) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange led fn) szNew . IntMap.insertWith RangeSet.union szNew RangeSet.empty) + modify' $ case fitSub of + Nothing -> over (eDB . unevaluated) (IntSet.delete sub . IntSet.delete subO) + Just fs -> over (eDB . fitRangeDB) (removeRange sub fs . removeRange subO fs) + . over (eDB . sizeFitDB) (IntMap.adjust (removeRange subO fs . removeRange sub fs) szSub) where fitNew = (_fitness . _info) newC fitLed = (_fitness . _info) ledC @@ -211,113 +305,93 @@ szSub = (_size . _info) subC -- | modify an e-class, e.g., add constant e-node and prune non-leaves -modifyEClass :: Monad m => CostFun -> EClassId -> EGraphST m EClassId +modifyEClass :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId modifyEClass costFun ecId = do ec <- getEClass ecId - -- let term = filter isTerm (Set.toList $ _eNodes ec) case (_consts . _info) ec of - ConstVal x -> do - let en = Const x - c <- calculateCost costFun en - 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 + ConstVal x -> + do let en = EConst x + c <- calculateCost costFun en + let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } + maybeEid <- lookupNode en + -- write through 'insertClass' (a paged store keeps the authoritative page) + insertClass ec{ _eNodes = Set.singleton en, _info = infoEc } + when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.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 - ens <- gets (_eNodes . (IntMap.! ecId) . _eClass) - 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' -> merge costFun eid' ecId + ParamIx x -> + do let en = EParam x + c <- calculateCost costFun en + let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } + maybeEid <- lookupNode en + insertClass ec{ _eNodes = Set.insert en (_eNodes ec), _info = infoEc } + when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.insert ecId) + case maybeEid of + Nothing -> pure ecId + Just eid' -> merge costFun eid' ecId _ -> pure ecId where - isTerm (Var _) = True - isTerm (Const _) = True - isTerm (Param _) = True - isTerm _ = False + isTerm (EVar _) = True + isTerm (EConst _) = True + isTerm (EParam _) = True + isTerm _ = False - toConst (Param ix) = ParamIx ix - toConst (Const x) = ConstVal x - toConst _ = NotConst + toConst (EParam ix) = ParamIx ix + toConst (EConst x) = ConstVal x + toConst _ = NotConst -- * DB --- | `createDB` creates a database of patterns from an e-graph --- it simply calls addToDB for every pair (e-node, e-class id) from --- the e-graph. -createDB :: Monad m => EGraphST m DB -createDB = do modify' $ over (eDB . patDB) (const Map.empty) - ecls <- gets (Map.toList . _eNodeToEClass) - mapM_ (uncurry addToDB) ecls - gets (_patDB . _eDB) -{-# INLINE createDB #-} - -createDBBest :: Monad m => EGraphST m DB -createDBBest = do modify' $ over (eDB . patDB) (const Map.empty) - ecls <- gets (Prelude.map (\(eId, ec) -> (_best (_info ec), eId)) . IntMap.toList . _eClass) - mapM_ (uncurry addToDB) ecls - gets (_patDB . _eDB) - -- | `addToDB` adds an e-node and e-class id to the database -addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB () +addToDB :: (ClassStore m, HasCallStack) => ENode -> EClassId -> EGraphST m () -- State DB () addToDB enode' eid = do eid' <- canonical eid - isConst <- gets (_consts . _info . (IntMap.! eid') . _eClass) + ec <- getEClass eid' + let isConst = _consts . _info $ ec let enode = case isConst of - ConstVal x -> Const x - ParamIx x -> Param x + ConstVal x -> EConst x + ParamIx x -> EParam x _ -> enode' - let ids = eid : childrenOf enode -- we will add the e-class id and the children ids - op = getOperator enode -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB - trie <- gets ((Map.!? op) . _patDB . _eDB) -- gets the entry for op, if it exists + let ids = eid : eChildren enode -- we will add the e-class id and the children ids + op = eOpKey enode -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB + trie <- gets (Map.lookup op . _patDB . _eDB) 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 + recordNode enode eid -- register the node for the streaming matcher's source {-# INLINE addToDB #-} -- | Populates an IntTrie with a sequence of e-class ids populate :: Maybe IntTrie -> [EClassId] -> Maybe IntTrie populate _ [] = Nothing --- if it is a new entry, simply add the ids sequentially populate Nothing eids = foldr f Nothing eids where f :: EClassId -> Maybe IntTrie -> Maybe IntTrie - f eid (Just t) = Just $ trie eid (IntMap.singleton eid t) - f eid Nothing = Just $ trie eid IntMap.empty --- if the entry already exists, insert the new key --- and populate the next child entry recursivelly -populate (Just tId) (eid:eids) = let keys = Set.insert eid (_keys tId) - nextTrie = _trie tId IntMap.!? eid - val = fromMaybe (trie eid IntMap.empty) $ populate nextTrie eids - in Just $ IntTrie keys (IntMap.insert eid val (_trie tId)) + f eid (Just t) = Just $ IntTrie (IntMap.singleton eid t) + f eid Nothing = Just $ IntTrie (IntMap.singleton eid (IntTrie IntMap.empty)) +populate (Just tId) (eid:eids) = let nextTrie = IntMap.lookup eid (_trie tId) + val = fromMaybe (IntTrie IntMap.empty) $ populate nextTrie eids + in Just $ IntTrie (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) +canonizeMap :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m (Subst, ClassOrVar) +canonizeMap (subst, cv) = (,cv) <$> traverse g subst where - g :: Monad m => ClassOrVar -> EGraphST m ClassOrVar - g (Left e2) = Left <$> canonical e2 - g e2 = pure e2 - - f :: Monad m => (ClassOrVar, ClassOrVar) -> EGraphST m (ClassOrVar, ClassOrVar) - f (e1, Left e2) = do e2' <- canonical e2 - pure (e1, Left e2') - f (e1, e2) = pure (e1, e2) + g :: ClassStore m => SubVal -> EGraphST m SubVal + g (SVOne e2) = SVOne <$> canonOne e2 + g (SVMap m) = SVMap . IntMap.fromListWith (+) <$> mapM (\(e2, n) -> do + e2' <- canonOne (Left e2) + pure (getInt e2', n)) (IntMap.toList m) + canonOne :: ClassStore m => ClassOrVar -> EGraphST m ClassOrVar + canonOne (Left e2) = Left <$> canonical e2 + canonOne e2 = pure e2 {-# INLINE canonizeMap #-} -applyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m () +applyMatch :: (ClassStore m, HasCallStack) => CostFun -> Rule -> (Subst, ClassOrVar) -> EGraphST m () applyMatch costFun rule match' = do let conds = getConditions rule match <- canonizeMap match' @@ -329,28 +403,16 @@ pure () {-# INLINE applyMatch #-} -applyMergeOnlyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m () -applyMergeOnlyMatch costFun rule match' = - do let conds = getConditions rule - match <- canonizeMap match' - validHeight <- isValidHeight match - validConds <- mapM (`isValidConditions` match) conds - when (validHeight && and validConds) $ - do maybe_eid <- classOfENode costFun (fst match) (target rule) - case maybe_eid of - 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 -classOfENode :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m (Maybe EClassId) -classOfENode costFun subst (VarPat c) = do let maybeEid = getInt <$> subst Map.!? Right (fromEnum c) +classOfENode :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m (Maybe EClassId) +classOfENode costFun subst (VarPat c) = do let maybeEid = case Map.lookup (Right (fromEnum c)) subst of + Just (SVOne v) -> Just v + _ -> Nothing case maybeEid of Nothing -> pure Nothing - Just eid -> Just <$> canonical eid -classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (Const x) + Just eid -> Just <$> canonical (getInt eid) +classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (EConst x) classOfENode costFun subst (Fixed target) = do newChildren <- mapM (classOfENode costFun subst) (getElems target) case sequence newChildren of Nothing -> pure Nothing @@ -358,42 +420,111 @@ cs' <- mapM canonical cs areConsts <- mapM isConst cs' if and areConsts - then do eid <- add costFun new_enode + then do eid <- addTree costFun new_enode rebuild costFun -- eid new_enode pure (Just eid) - else gets ((Map.!? new_enode) . _eNodeToEClass) + else do en <- toENode new_enode + en' <- canonize en + gets (HashMap.lookup en' . _eNodeToEClass) +classOfENode _ _ (NAry _ _) = error "classOfENode: n-ary pattern unsupported" +classOfENode _ _ Hole = error "classOfENode: Hole is only valid in MapP targets" {-# 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 :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m EClassId +reprPrat costFun subst (VarPat c) = do + let k = Right (fromEnum c) + v <- case Map.lookup k subst of + Nothing -> error $ "REPRPRAT_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst) + Just (SVOne x) -> pure x + Just (SVMap _) -> error $ "REPRPRAT_REST_AS_SINGLE var=" <> show (fromEnum c) + canonical $ getInt v reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target) - add costFun (replaceChildren newChildren target) + addTree costFun (replaceChildren newChildren target) +reprPrat costFun subst Hole = error "REPRPRAT_HOLE: Hole must be filled by MapP" +reprPrat costFun subst (NAry op ncs) = do + m <- IntMap.unionsWith (+) <$> mapM (childEidM costFun subst) ncs + case IntMap.toList m of + [] -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1))) + [(c, 1)] -> canonical c + _ -> do en <- mkENaryM op m + add costFun en {-# INLINE reprPrat #-} -isValidHeight :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool +-- | Adds a single child of an n-ary target pattern to the e-graph, returning +-- its contribution as a canonical multiset (so 'Rest' children carry their +-- 'IntMap' straight through without expansion). +childEidM :: (ClassStore m, HasCallStack) => CostFun -> Subst -> NChild -> EGraphST m (IntMap Int) +childEidM costFun subst (Ch p) = (`IntMap.singleton` 1) <$> reprPrat costFun subst p +childEidM costFun subst (Rest c) = restEidsM subst c +childEidM costFun subst (MapP p c) = do + es <- restEids subst c + ms <- forM es $ \e -> reprMapP costFun subst e p + pure (imFromList ms) +{-# INLINE childEidM #-} + +-- | The e-class ids bound to a rest variable, as a canonical multiset. +restEidsM :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m (IntMap Int) +restEidsM subst c = do + let k = Right (fromEnum c) + case Map.lookup k subst of + Just (SVMap m) -> pure m + Just (SVOne _) -> error $ "REPRPRAT_SINGLE_AS_REST var=" <> show (fromEnum c) + Nothing -> error $ "REPRPRAT_MISSING_REST var=" <> show (fromEnum c) +{-# INLINE restEidsM #-} + +-- | The e-class ids bound to a rest variable, expanded one entry per +-- occurrence (used by 'MapP', which needs to instantiate per child). +restEids :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m [EClassId] +restEids subst c = expandedList <$> restEidsM subst c +{-# INLINE restEids #-} + +-- | Build the target of a pattern where every `Hole` is filled with the +-- e-class `e` (used by 'MapP'). +reprMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> Pattern -> EGraphST m EClassId +reprMapP costFun subst e Hole = canonical e +reprMapP costFun subst e (VarPat c) = reprPrat costFun subst (VarPat c) +reprMapP costFun subst e (Fixed target) = do + newChildren <- mapM (reprMapP costFun subst e) (getElems target) + addTree costFun (replaceChildren newChildren target) +reprMapP costFun subst e (NAry op ncs) = do + m <- IntMap.unionsWith (+) <$> mapM (childMapP costFun subst e) ncs + case IntMap.toList m of + [] -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1))) + [(c, 1)] -> canonical c + _ -> do en <- mkENaryM op m + add costFun en +{-# INLINE reprMapP #-} + +-- | A single child of an n-ary pattern inside a 'MapP' function. +childMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> NChild -> EGraphST m (IntMap Int) +childMapP costFun subst e (Ch p) = (`IntMap.singleton` 1) <$> reprMapP costFun subst e p +childMapP costFun subst e (Rest c) = restEidsM subst c +childMapP costFun subst e (MapP _ _) = error "nested MapP unsupported" +{-# INLINE childMapP #-} + +isValidHeight :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m Bool isValidHeight match = do - h <- case snd match of - Left ec -> do ec' <- canonical ec - gets (_height . (IntMap.! ec') . _eClass) - Right _ -> pure 0 - pure $ h < 15 + h <- case snd match of + Left ec -> _height <$> getEClass ec + 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) +isValidConditions :: ClassStore m => Condition -> (Subst, ClassOrVar) -> EGraphST m Bool +isValidConditions (Condition f) match = f (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) +fromTree :: (ClassStore m, HasCallStack) => CostFun -> Fix SRTree -> EGraphST m EClassId +fromTree costFun = cataM sequence (addTree costFun) {-# INLINE fromTree #-} -- | Builds an e-graph from multiple independent trees -fromTrees :: Monad m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId] +fromTrees :: ClassStore m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId] fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) [] {-# INLINE fromTrees #-} @@ -403,119 +534,93 @@ countParamsUniqEg eg rt = countParamsUniq . runIdentity $ getBestExpr rt `evalStateT` eg --- | gets the best expression given the default cost function -getBestExpr :: Monad m => EClassId -> EGraphST m (Fix SRTree) -getBestExpr eid = do eid' <- canonical eid - 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) +getBestENode eid = (_best . _info) <$> getEClass eid {-# INLINE getBestENode #-} -- | returns one expression rooted at e-class `eId` -- TODO: avoid loopings -getExpressionFrom :: Monad m => EClassId -> EGraphST m (Fix SRTree) +getExpressionFrom :: ClassStore m => EClassId -> EGraphST m (Fix SRTree) getExpressionFrom eId' = do - eId <- canonical eId' - nodes <- gets (Set.map decodeEnode . _eNodes . (IntMap.! eId) . _eClass) - let hasTerm = any isTerm nodes - cands = if hasTerm then filter isTerm (Set.toList nodes) else Set.toList nodes - - Fix <$> case head $ Set.toList nodes of - Bin op l r -> Bin op <$> getExpressionFrom l <*> getExpressionFrom r - Uni f t -> Uni f <$> getExpressionFrom t - Var ix -> pure $ Var ix - Const x -> pure $ Const x - Param ix -> pure $ Param ix - where - isTerm (Var _) = True - isTerm (Const _) = True - isTerm (Param _) = True - isTerm _ = False + nodes <- _eNodes <$> getEClass eId' + case Set.toList nodes of + (n:_) -> case n of + EVar ix -> pure $ Fix $ Var ix + EParam ix -> pure $ Fix $ Param ix + EConst x -> pure $ Fix $ Const x + EUni f t -> Fix . Uni f <$> getExpressionFrom t + EBin op l r -> Fix <$> (Bin op <$> getExpressionFrom l <*> getExpressionFrom r) + ENAry op xs -> naryTree op <$> mapM getExpressionFrom (expandedList xs) + [] -> error "getExpressionFrom: empty eclass" {-# INLINE getExpressionFrom #-} -- | returns all expressions rooted at e-class `eId` -- TODO: check for infinite list -getAllExpressionsFrom :: Monad m => EClassId -> EGraphST m [Fix SRTree] +getAllExpressionsFrom :: ClassStore m => EClassId -> EGraphST m [Fix SRTree] getAllExpressionsFrom eId' = do - eId <- canonical eId' - nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass) - let cands = filter isTerm nodes - concat <$> go nodes - --if null cands - -- then concat <$> go nodes - -- else pure [toTree $ head cands] + nodes <- Set.toList . _eNodes <$> getEClass eId' + go nodes where - isTerm (Var _) = True - isTerm (Const _) = True - isTerm (Param _) = True - isTerm _ = False - toTree (Var ix) = Fix $ Var ix - toTree (Const x) = Fix $ Const x - toTree (Param ix) = Fix $ Param ix - toTree _ = undefined - go [] = pure [] go (n:ns) = do - t <- Prelude.map Fix <$> case n of - Bin op l r -> do l' <- getAllExpressionsFrom l - r' <- getAllExpressionsFrom r - pure $ [Bin op li ri | li <- l', ri <- r'] - Uni f t -> Prelude.map (Uni f) <$> getAllExpressionsFrom t - Var ix -> pure [Var ix] - Const x -> pure [Const x] - Param ix -> pure [Param ix] + t <- case n of + EVar ix -> pure [Fix $ Var ix] + EParam ix -> pure [Fix $ Param ix] + EConst x -> pure [Fix $ Const x] + EUni f t -> Prelude.map (Fix . Uni f) <$> getAllExpressionsFrom t + EBin op l r -> do l' <- getAllExpressionsFrom l + r' <- getAllExpressionsFrom r + pure $ [Fix $ Bin op li ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM getAllExpressionsFrom (expandedList xs) + pure [ naryTree op comb | comb <- sequence ts ] ts <- go ns - pure (t:ts) + pure (t ++ ts) {-# INLINE getAllExpressionsFrom #-} -getNExpressionsFrom :: Monad m => Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom :: ClassStore m => Int -> EClassId -> EGraphST m [Fix SRTree] getNExpressionsFrom n eId' = getNExpressionsFrom' n 15 eId' -getNExpressionsFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree] getNExpressionsFrom' _ 0 _ = pure [] getNExpressionsFrom' n d eId' = do - eId <- canonical eId' - nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass) + nodes <- Set.toList . _eNodes <$> getEClass eId' (concat <$> go n d nodes) where - isTerm (Var _) = True - isTerm (Const _) = True - isTerm (Param _) = True + isTerm (EVar _) = True + isTerm (EConst _) = True + isTerm (EParam _) = True isTerm _ = False - toTree (Var ix) = Fix $ Var ix - toTree (Const x) = Fix $ Const x - toTree (Param ix) = Fix $ Param ix + toTree (EVar ix) = Fix $ Var ix + toTree (EConst x) = Fix $ Const x + toTree (EParam ix) = Fix $ Param ix toTree _ = undefined go n' _ [] = pure [] go n' 0 ts = pure [] go n' d (node:ns) = do - tt <- Prelude.map Fix <$> case node of - Bin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l - r' <- getNExpressionsFrom' n' (d-1) r - pure $ Prelude.take n [Bin op li ri | li <- l', ri <- r'] - Uni f t -> Prelude.map (Uni f) <$> getNExpressionsFrom' n' (d-1) t - Var ix -> pure [Var ix] - Const x -> pure [Const x] - Param ix -> pure [Param ix] + tt <- case node of + EVar ix -> pure [Fix $ Var ix] + EParam ix -> pure [Fix $ Param ix] + EConst x -> pure [Fix $ Const x] + EUni f t -> Prelude.map (Fix . Uni f) <$> getNExpressionsFrom' n' (d-1) t + EBin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l + r' <- getNExpressionsFrom' n' (d-1) r + pure $ Prelude.take n [Fix $ Bin op li ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM (getNExpressionsFrom' n' (d-1)) (expandedList xs) + pure $ Prelude.take n [ naryTree op comb | comb <- sequence ts ] let n'' = n' - length tt if n'' <= 0 then pure [tt] else do ts <- go n'' (d-1) ns pure (tt:ts) -getNEclassFrom :: Monad m => Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom :: ClassStore m => Int -> EClassId -> EGraphST m [[EClassId]] getNEclassFrom n eid = getNEclassFrom' n 15 eid -getNEclassFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [[EClassId]] getNEclassFrom' _ 0 _ = pure [] getNEclassFrom' n d eId' = do eId <- canonical eId' - nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass) + nodes <- Set.toList . _eNodes <$> getEClass eId' (Prelude.map (eId:) <$> go n d nodes) where --go :: Int -> Int -> [ENode] -> EGraphST m [[EClassId]] @@ -523,13 +628,15 @@ go n' 0 ts = pure [] go n' d (node:ns) = do tt <- case node of - Bin op l r -> do l' <- getNEclassFrom' n' (d-1) l - r' <- getNEclassFrom' n' (d-1) r - pure $ Prelude.take n [li <> ri | li <- l', ri <- r'] - Uni f t -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]] - Var ix -> pure [[]] - Const x -> pure [[]] - Param ix -> pure [[]] + EBin op l r -> do l' <- getNEclassFrom' n' (d-1) l + r' <- getNEclassFrom' n' (d-1) r + pure $ Prelude.take n [li <> ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM (getNEclassFrom' n' (d-1)) xs + pure $ Prelude.take n [ concat comb | comb <- sequence ts ] + EUni f t -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]] + EVar ix -> pure [[]] + EConst x -> pure [[]] + EParam ix -> pure [[]] pure tt --let n'' = n' - length tt --if n'' <= 0 @@ -537,21 +644,21 @@ -- else do ts <- go n'' (d-1) ns -- pure (tt:ts) -getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId] +getAllChildEClasses :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildEClasses eId' = do eId <- canonical eId' IntSet.toList <$> go [eId] IntSet.empty where hasNoTerminal :: [ENode] -> Bool - hasNoTerminal = all (not . null . childrenOf) - getNodes :: Monad m => EClassId -> EGraphST m [ENode] - getNodes n = gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass) + hasNoTerminal = all (not . null . eChildren) + getNodes :: ClassStore m => EClassId -> EGraphST m [ENode] + getNodes n = Set.toList . _eNodes <$> getEClass n - go :: Monad m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet + go :: ClassStore m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet go [] visited = pure visited go queue visited = do - nodes <- concatMap childrenOf . concat . filter hasNoTerminal <$> mapM getNodes queue + nodes <- concatMap eChildren . concat . filter hasNoTerminal <$> mapM getNodes queue eids <- filter (\e -> e `IntSet.notMember` visited) <$> (mapM canonical nodes) go eids (visited `IntSet.union` IntSet.fromList queue) {- @@ -565,31 +672,27 @@ -} {-# INLINE getAllChildEClasses #-} -getAllChildBestEClasses :: Monad m => EClassId -> EGraphST m [EClassId] +getAllChildBestEClasses :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildBestEClasses eId' = do - eId <- canonical eId' - nub <$> go eId - + IntSet.toList <$> go IntSet.empty eId' where - go :: Monad m => Int -> EGraphST m [Int] - go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass) - let hasTerminal = (null . childrenOf) node - eids <- mapM canonical $ childrenOf node - if hasTerminal - then pure [n] - else do eids' <- mapM go eids - pure ((n : eids) <> concat eids') + go :: ClassStore m => IntSet.IntSet -> EClassId -> EGraphST m IntSet.IntSet + go acc n + | IntSet.member n acc = pure acc + | otherwise = do + let acc' = IntSet.insert n acc + node <- (_best . _info) <$> getEClass n + eids <- mapM canonical $ eChildren node + foldM go acc' eids -getAllChildBestEClassesRep :: Monad m => EClassId -> EGraphST m [EClassId] +getAllChildBestEClassesRep :: ClassStore m => EClassId -> EGraphST m [EClassId] getAllChildBestEClassesRep eId' = do - eId <- canonical eId' - go eId - + go eId' where - go :: Monad m => Int -> EGraphST m [Int] - go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass) - let hasTerminal = (null . childrenOf) node - eids <- mapM canonical $ childrenOf node + go :: ClassStore m => EClassId -> EGraphST m [EClassId] + go n = do node <- (_best . _info) <$> getEClass n + let hasTerminal = (null . eChildren) node + eids <- mapM canonical $ eChildren node if hasTerminal then pure [n] else do eids' <- mapM go eids @@ -598,43 +701,43 @@ -- | returns a random expression rooted at e-class `eId` getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree) getRndExpressionFrom eId' = do - eId <- canonical eId' - nodes <- gets (Set.toList . _eNodes . (IntMap.! eId) . _eClass) + nodes <- Set.toList . _eNodes <$> getEClass eId' n <- lift $ randomFrom nodes - Fix <$> case decodeEnode n of - Bin op l r -> Bin op <$> getRndExpressionFrom l <*> getRndExpressionFrom r - Uni f t -> Uni f <$> getRndExpressionFrom t - Var ix -> pure $ Var ix - Const x -> pure $ Const x - Param ix -> pure $ Param ix + case n of + EUni f t -> Fix . Uni f <$> getRndExpressionFrom t + EBin op l r -> Fix <$> (Bin op <$> getRndExpressionFrom l <*> getRndExpressionFrom r) + ENAry op xs -> naryTree op <$> mapM getRndExpressionFrom (expandedList xs) + EVar ix -> pure $ Fix $ Var ix + EConst x -> pure $ Fix $ Const x + EParam ix -> pure $ Fix $ Param ix where 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 :: ClassStore m => EGraphST m () cleanMaps = do - enode2eclass <- gets _eNodeToEClass - entries <- forM (Map.toList enode2eclass) $ \(k,v) -> do - k' <- canonize k - v' <- canonical v - pure (k',v') - let enode2eclass' = Map.fromList entries - eclassMap <- gets _eClass - entries' <- forM (IntMap.toList eclassMap) $ \(k,v) -> do - k' <- canonical k - pure $ if k==k' then (Just (k,v)) else Nothing - let eclassMap' = IntMap.fromList (catMaybes entries') - canon <- gets _canonicalMap - entries'' <- forM (IntMap.toList canon) $ \(k,v) -> do - pure $ if k==v then Just (k,v) else Nothing - let canon' = IntMap.fromList (catMaybes entries'') - eDB' <- gets _eDB - put $ EGraph canon enode2eclass' eclassMap' eDB' - forceState + hasStore <- gets (isJust . _classStore) + if hasStore + -- the paged store is authoritative for both node->class and canonical + -- lookups, so the bounded resident caches are simply reset (an O(n) rebuild + -- of an unbounded map would defeat the out-of-core goal). + then modify' $ \eg -> eg { _eNodeToEClass = HashMap.empty + , _canonicalMap = IntMap.empty + , _eClass = IntMap.empty } + else do + enode2eclass <- gets _eNodeToEClass + entries <- forM (HashMap.toList enode2eclass) $ \(k,v) -> do + k' <- canonize k + v' <- canonical v + pure (k',v') + let enode2eclass' = HashMap.fromList entries + eclassMap <- gets _eClass + entries' <- forM (IntMap.toList eclassMap) $ \(k,v) -> do + k' <- canonical k + pure $ if k==k' then (Just (k,v)) else Nothing + let eclassMap' = IntMap.fromList (catMaybes entries') + modify' $ \eg -> eg { _eNodeToEClass = enode2eclass' + , _eClass = eclassMap' } {-# INLINE cleanMaps #-} - -forceState :: Monad m => StateT s m () -forceState = get >>= \ !_ -> return () -{-# INLINE forceState #-} diff --git a/src/Algorithm/EqSat/DB.hs b/src/Algorithm/EqSat/DB.hs --- a/src/Algorithm/EqSat/DB.hs +++ b/src/Algorithm/EqSat/DB.hs @@ -1,5 +1,6 @@ {-# LANGUAGE TupleSections #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} ----------------------------------------------------------------------------- -- | -- Module : Algorithm.EqSat.EqSatDB @@ -19,26 +20,36 @@ import Control.Lens ( over ) import Control.Monad (when, foldM, forM) import Control.Monad.State -import Data.IntMap (IntMap) -import qualified Data.IntMap as IntMap -import Data.List (intercalate, nub, sortBy) +import GHC.Stack (HasCallStack) +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IntMap import Data.Map (Map) import qualified Data.Map as Map +import Data.List (sortBy) import Data.Maybe (fromMaybe) import Data.Ord (comparing) import Data.SRTree ---import Data.Set (Set) import Data.HashSet (HashSet) import qualified Data.HashSet as Set +import qualified Data.Set as RangeSet import Data.String (IsString (..)) import Data.SRTree.Recursion (cata) +import Text.Read (readMaybe) -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, Eq, Ord) -- Fixed structure of a pattern or a variable that matches anything +-- A Pattern is either a fixed-point of a tree, an index to a pattern variable +-- (which matches anything), a hole (only used inside a 'MapP' target function), +-- or an n-ary Add/Mul pattern whose children are matched as a multiset. +data Pattern = Fixed (SRTree Pattern) | VarPat Char | Hole | NAry NOp [NChild] + deriving (Show, Eq, Ord) +-- | A child of an n-ary pattern: a single child pattern ('Ch'), a rest +-- variable binding every remaining child of the node ('Rest'), or a +-- target-side map that splices one instantiation of a pattern (with its 'Hole' +-- filled) per child bound to a rest variable ('MapP'). +data NChild = Ch Pattern | Rest Char | MapP Pattern Char + deriving (Show, Eq, Ord) + -- The instance for `IsString` for a `Pattern` is -- valid only for a single letter char from a-zA-Z. -- The patterns can be written as "x" + "y", for example, @@ -54,6 +65,8 @@ 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 Add l r) = NAry EAdd [Ch l, Ch r] + alg (Bin Mul l r) = NAry EMul [Ch l, Ch r] 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 @@ -73,32 +86,45 @@ -- A Query is a list of Atoms type Query = [Atom] --- A `Condition` is a function that takes a substution map, --- an e-graph and returns whether the pattern attends the condition. -type Condition = Map ClassOrVar ClassOrVar -> EGraph -> Bool +-- | A `Condition` is a predicate over a match's substitution that runs inside +-- the e-graph monad so it can fetch e-class data through 'ClassStore' (which +-- streams from a paged store when the graph is out-of-core). The quantification +-- over the monad is intentional: the same condition works for any 'ClassStore' +-- instance, including the IO-backed paged store. +newtype Condition = Condition (forall m. ClassStore m => Subst -> EGraphST m Bool) -- An Atom is composed of either an e-class id or pattern variable id -- and the tree that generated that pattern. Left is e-class id and Right is a VarPat. type ClassOrVar = Either EClassId Int data Atom = Atom ClassOrVar (SRTree ClassOrVar) deriving Show +-- | A substitution value: a single e-class (a matched pattern variable) or the +-- canonical multiset of e-class ids (a matched rest variable). +data SubVal = SVOne ClassOrVar | SVMap (IntMap Int) deriving Show + +-- | Substitution map produced by matching a pattern. +type Subst = Map ClassOrVar SubVal + unFixPat :: Pattern -> SRTree Pattern unFixPat (Fixed p) = p +unFixPat (VarPat _) = error "unFixPat: VarPat is not a fixed pattern" +unFixPat Hole = error "unFixPat: Hole is not a fixed pattern" +unFixPat (NAry _ _) = error "unFixPat: NAry is not a fixed pattern" {-# INLINE unFixPat #-} instance Num Pattern where - l + r = Fixed $ Bin Add l r + l + r = NAry EAdd [Ch l, Ch r] {-# INLINE (+) #-} - l - r = Fixed $ Bin Sub l r + l - r = NAry EAdd [Ch l, Ch (negate r)] {-# INLINE (-) #-} - l * r = Fixed $ Bin Mul l r + l * r = NAry EMul [Ch l, Ch r] {-# INLINE (*) #-} abs = Fixed . Uni Abs {-# INLINE abs #-} - negate t = Fixed (Const (-1)) * t + negate t = NAry EMul [Ch (Fixed (Const (-1))), Ch t] {-# INLINE negate #-} signum t = case t of @@ -108,7 +134,7 @@ {-# INLINE fromInteger #-} instance Fractional Pattern where - l / r = Fixed $ Bin Div l r + l / r = NAry EMul [Ch l, Ch (Fixed (Uni Recip r))] {-# INLINE (/) #-} fromRational = Fixed . Const . fromRational @@ -176,38 +202,347 @@ {-# INLINE cleanDB #-} -- | Returns the substitution rules --- for every match of the pattern `source` inside the e-graph. -match :: Monad m => Pattern -> EGraphST m [(Map ClassOrVar ClassOrVar, ClassOrVar)] -match src = do - 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] +-- for every match of the pattern `source` inside the e-graph. This is the pure +-- matcher (no seen-set) used by user pattern queries; saturation uses +-- 'matchSaturated'. +match :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)] +match src = if hasNAry src + then matchNAryWith Nothing src + else do + paged <- isPagedGraph + if paged + then matchStreamCached Nothing src + else matchCachedWith Nothing (compileToQuery src) {-# INLINE match #-} --- | Returns a Query (list of atoms) of a pattern -compileToQuery :: Pattern -> (Query, ClassOrVar) -compileToQuery pat = evalState (processPat pat) 256 -- returns (atoms, root) +-- | Non-n-ary matching. The match's root e-class anchors it the same way the +-- n-ary matcher anchors one match per trie root, so it shares the same cheap +-- persistent mark-on-attempt seen-set ('_seenMatches', keyed by rule source -> +-- root class id): already-processed roots are skipped so the per-rule budget +-- advances to new matches across the scheduler's ban/unban cycles. Keying by +-- the root (an @O(1)@ class id) avoids serializing every substitution, which +-- would dominate on rules whose @genericJoin@ yields many matches. 'Nothing' +-- disables the seen-set (pure queries). +matchCachedWith :: ClassStore m => Maybe String -> (Query, [ClassOrVar], ClassOrVar) -> EGraphST m [(Subst, ClassOrVar)] +matchCachedWith mSk (q, vars, root) = do + ss <- genericJoin q vars root + seenSk <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + let rootOf s = case Map.lookup root s of + Just (SVOne (Left eid)) -> eid + _ -> 0 + fresh = [ s | s <- ss + , Map.size s > 0 + , maybe True (\_ -> not (RangeSet.member (show (rootOf s)) seenSk)) mSk ] + taken = take ruleMatchBudget fresh + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.fromList (map (show . rootOf) taken))) + Nothing -> pure () + pure [ (s, case Map.lookup root s of + Nothing -> error $ "MATCHCACHED_MISSING root=" <> show (getInt root) <> " substSize=" <> show (Map.size s) + Just v -> fromSVOne v) + | s <- taken ] +{-# INLINE matchCachedWith #-} + +-- | Saturation matching: consults/marks the persistent seen-set so each rule's +-- per-iteration budget advances to genuinely new matches across ban/unban. +matchSaturated :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchSaturated src = if hasNAry src + then matchNAryWith (Just (show src)) src + else do + paged <- isPagedGraph + if paged + then matchStreamCached (Just (show src)) src + else matchCachedWith (Just (show src)) (compileToQuery src) +{-# INLINE matchSaturated #-} + +-- | True if the pattern (or a nested child) is an n-ary Add/Mul pattern. +hasNAry :: Pattern -> Bool +hasNAry (NAry _ _) = True +hasNAry (Fixed t) = any hasNAry (getElems t) +hasNAry _ = False +{-# INLINE hasNAry #-} + +-- | The operator trie key of the top-level pattern. +opOf :: Pattern -> SRTree () +opOf (NAry EAdd _) = Bin Add () () +opOf (NAry EMul _) = Bin Mul () () +opOf (Fixed t) = getOperator t +opOf _ = error "opOf: pattern has no operator" +{-# INLINE opOf #-} + +-- | Matches an n-ary pattern against every root e-node of the operator trie. +-- A per-rule result budget ('ruleBudget') bounds the total number of matches +-- returned for one rule against one individual's nodes, and only the first +-- match per root e-class is kept, taming the O(k^2*m^2) backtracking of +-- Rest/Ch rules (e.g. factoring a common term out of a sum of products). +-- Keeping one match per root is sound: every returned match is genuine, and +-- the egraph merges the equivalent rewrites that further matches would apply, +-- so the rest of the root's matches are redundant work. +ruleBudget :: Int +ruleBudget = 64 + +-- | Cap on how many operator-trie root e-classes a single rule may visit per +-- match. 'ruleBudget' bounds the number of *results* returned, but a rule whose +-- matches are rare would otherwise still scan every root e-class in the trie +-- (every @+@/@*@ class in the graph), doing an expensive 'recursiveMatch' per +-- root -- which blows up on large graphs even though few matches result. +-- Capping root visits bounds the *search work* independently of the result +-- count. Sound: we only stop enumerating (fewer) genuine matches early. +ruleRootVisit :: Int +ruleRootVisit = 512 + +-- | Cap on how many matches a non-n-ary rule (the cached @genericJoin@ path) +-- may return per match. The n-ary matcher has 'ruleBudget'; give the cached +-- path a separate (larger) budget so a single rule cannot flood the iteration. +ruleMatchBudget :: Int +ruleMatchBudget = 1024 + +-- | Cap on how many operator-root e-classes the streaming cached matcher visits +-- per match, bounding the search work (and the page reads) independently of the +-- result count, exactly as 'ruleRootVisit' does for the n-ary matcher. +ruleMatchRootVisit :: Int +ruleMatchRootVisit = 2048 + +-- | Match an n-ary pattern against every root e-class of its operator trie. +-- +-- A persistent per-source set of already-attempted roots ('_seenMatches') lets +-- the matcher skip roots it has already tried, so the per-rule result/search +-- budgets keep advancing to *new* roots across the scheduler's ban/unban cycles +-- instead of re-enumerating the same head of the trie (which starves the tail). +-- Roots are marked as attempted on the first try ('mark-on-attempt'), whether or +-- not they yielded a match, so a match that fails 'applyMatch' conditions is not +-- re-attempted every cycle. +matchNAryWith :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchNAryWith mSk src = do + seen <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + -- skip already-attempted roots so the per-rule budget advances to new roots + -- across the scheduler's ban/unban cycles (matches the trie path's semantics). + let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ] + roots <- streamRoots (opOf src) ruleRootVisit exclude + go roots 0 0 [] where + go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)] + go [] _ _ acc = pure (reverse acc) + go _ n _ acc | n >= ruleBudget = pure (reverse acc) + go (_ : _) _ r acc | r >= ruleRootVisit = pure (reverse acc) + go (eid : eids) n r acc = do + -- mark-on-attempt: remember this root as tried for this rule source + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid))) + Nothing -> pure () + substs <- recursiveMatch src eid Map.empty + let newMs = take 1 [ (s, Left eid) | s <- substs ] + go eids (n + length newMs) (r + 1) (foldr (:) acc newMs) +{-# INLINE matchNAryWith #-} + +-- | Streaming matcher for the cached (non-n-ary @genericJoin@) path on a paged +-- graph. Instead of enumerating candidates from the in-RAM @_patDB@ trie, it +-- streams the candidate root e-classes of the pattern's operator through +-- 'streamRoots' (bounded, skipping the already-attempted seen-set) and matches +-- each root incrementally with 'recursiveMatch' (which reads e-classes through +-- the paged store). This is the out-of-core analogue of 'matchCachedWith': the +-- resident/pure path keeps the optimized trie 'genericJoin', and only a paged +-- graph takes this route, so the matcher never builds an O(nodes) structure. +-- +-- 'ruleMatchBudget' bounds the results and 'ruleMatchRootVisit' bounds the root +-- visits; the persistent mark-on-attempt seen-set makes each rule's budgets +-- advance to new roots across the scheduler's ban/unban cycles. +matchStreamCached :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchStreamCached mSk src = do + seen <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ] + roots <- case opOfMay src of + Just op -> streamRoots op ruleMatchRootVisit exclude + Nothing -> pure [] + go roots 0 0 [] + where + go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)] + go [] _ _ acc = pure (reverse acc) + go _ n _ acc | n >= ruleMatchBudget = pure (reverse acc) + go (_ : _) _ r acc | r >= ruleMatchRootVisit = pure (reverse acc) + go (eid : eids) n r acc = do + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid))) + Nothing -> pure () + substs <- recursiveMatch src eid Map.empty + let newMs = take (ruleMatchBudget - n) [ (s, Left eid) | s <- substs ] + go eids (n + length newMs) (r + 1) (foldr (:) acc newMs) +{-# INLINE matchStreamCached #-} + +-- | The operator trie key of the top-level pattern, or @Nothing@ for a pattern +-- with no operator (e.g. a bare variable), which the streaming matcher treats +-- as matching nothing. +opOfMay :: Pattern -> Maybe (SRTree ()) +opOfMay (NAry EAdd _) = Just (Bin Add () ()) +opOfMay (NAry EMul _) = Just (Bin Mul () ()) +opOfMay (Fixed t) = Just (getOperator t) +opOfMay _ = Nothing +{-# INLINE opOfMay #-} + +-- | Recursively match a pattern against the e-class `eid`, threading a +-- substitution map, returning every substitution that completes the match. +recursiveMatch :: ClassStore m => Pattern -> EClassId -> Subst -> EGraphST m [Subst] +recursiveMatch (VarPat c) eid subst = + pure (bindVar subst (Right (fromEnum c)) eid) +recursiveMatch Hole _ subst = pure [subst] +recursiveMatch (Fixed t) eid subst = matchFixed t eid subst +recursiveMatch (NAry op ncs) eid subst = matchNAryNode op ncs eid subst +{-# INLINE recursiveMatch #-} + +-- | Bind `v` to the e-class `eid`, enforcing that re-occurrences of `v` are +-- consistent. +bindVar :: Subst -> ClassOrVar -> EClassId -> [Subst] +bindVar subst v eid = + case Map.lookup v subst of + Just (SVOne e) | e == Left eid -> [subst] + Just _ -> [] + Nothing -> [Map.insert v (SVOne (Left eid)) subst] +{-# INLINE bindVar #-} + +-- | Match a fixed tree pattern against the e-nodes of the e-class `eid`, +-- returning every substitution that completes the match across all candidate +-- e-nodes. +matchFixed :: ClassStore m => SRTree Pattern -> EClassId -> Subst -> EGraphST m [Subst] +matchFixed t eid subst = do + ec <- getEClass eid + let cands = [n | n <- Set.toList (_eNodes ec), eOpKey n == getOperator t] + fmap concat $ forM cands $ \n -> matchChildren t subst n + where + matchChildren t s n = go (zip (getElems t) (enodeChildren n)) [s] + go [] ss = pure ss + go ((p, c) : ps) ss = do + ms <- concat <$> mapM (\s -> recursiveMatch p c s) ss + go ps ms +{-# INLINE matchFixed #-} + +-- | The child e-class ids of an e-node, in canonical (sorted for ENAry) order. +enodeChildren :: ENode -> [EClassId] +enodeChildren (EUni _ t) = [t] +enodeChildren (EBin _ l r) = [l, r] +enodeChildren (ENAry _ m) = expandedList m +enodeChildren _ = [] +{-# INLINE enodeChildren #-} + +-- | Match an n-ary pattern node against the e-class `eid`: it must contain an +-- ENAry node of the given op, whose children are matched as a multiset. Every +-- ENAry node in the class is tried. +matchNAryNode :: ClassStore m => NOp -> [NChild] -> EClassId -> Subst -> EGraphST m [Subst] +matchNAryNode op ncs eid subst = do + ec <- getEClass eid + let nodes = [m | ENAry op' m <- Set.toList (_eNodes ec), op' == op] + fmap concat $ forM nodes $ \m -> + matchNChildren ncs m subst +{-# INLINE matchNAryNode #-} + +-- | Match a sequence of n-ary children against a multiset of e-class ids. +-- Each 'Ch' consumes one matched child; a 'Rest' child consumes all remaining +-- children. Every multiset assignment is returned. Iterating over the distinct +-- child ids (the multiset's keys) is sound (duplicate copies only differ by +-- position, which 'decChild' already resolves) and avoids duplicate result +-- sets. +-- +-- A per-call result budget ('matchCap') caps the number of substitutions +-- returned, bounding the O(k^2*m^2) backtracking of Rest/Ch rules such as +-- factoring a common term out of a sum of products. Sound: each result is a +-- genuine match; we merely stop enumerating once the budget is exhausted. +matchCap :: Int +matchCap = 64 + +matchNChildren :: ClassStore m => [NChild] -> IntMap Int -> Subst -> EGraphST m [Subst] +matchNChildren ncs children subst = reverse <$> goB ncs children subst matchCap + where + goB :: ClassStore m => [NChild] -> IntMap Int -> Subst -> Int -> EGraphST m [Subst] + goB [] m s _ + | IntMap.null m = pure [s] + | otherwise = pure [] + goB (Rest c : ps) m s b = do + let v = Right (fromEnum c) + case Map.lookup v s of + Just _ -> pure [] -- rest variable already bound + Nothing -> goB ps IntMap.empty (Map.insert v (SVMap m) s) b + goB (Ch p : ps) m s b + | multiplicity m <= nCh ps = pure [] -- not enough children left + | otherwise = goC (IntMap.keys m) 0 [] + where + goC :: ClassStore m => [EClassId] -> Int -> [Subst] -> EGraphST m [Subst] + goC [] _ acc = pure acc + goC _ n acc | n >= b = pure acc + goC (c : cs) n acc = do + ms <- recursiveMatch p c s + goMs c ms cs n acc + goMs :: ClassStore m => EClassId -> [Subst] -> [EClassId] -> Int -> [Subst] -> EGraphST m [Subst] + goMs c [] cs n acc = goC cs n acc + goMs c (s' : ms) cs n acc + | n >= b = pure acc + | otherwise = do + r <- goB ps (decChild c m) s' (b - n) + let r' = take (b - n) r + n' = n + length r' + goMs c ms cs n' (foldr (:) acc r') + goB (MapP _ _ : _) _ _ _ = error "matchNChildren: MapP is only valid in targets" +{-# INLINE matchNChildren #-} + +-- | Total number of children (counting multiplicities) in a multiset. +multiplicity :: IntMap Int -> Int +multiplicity = IntMap.foldr' (+) 0 +{-# INLINE multiplicity #-} + +-- | Remove one occurrence of `c` from the multiset (decrementing its +-- multiplicity, or dropping the key entirely when it reaches zero). +decChild :: Int -> IntMap Int -> IntMap Int +decChild c = IntMap.update (\n -> if n > 1 then Just (n - 1) else Nothing) c +{-# INLINE decChild #-} + +-- | Number of 'Ch' patterns in a child pattern sequence (each consumes one +-- child, so at least this many children must remain). +nCh :: [NChild] -> Int +nCh = length . filter isCh + where + isCh (Ch _) = True + isCh _ = False +{-# INLINE nCh #-} + +-- | Unwrap a single-e-class substitution value. +fromSVOne :: SubVal -> ClassOrVar +fromSVOne (SVOne v) = v +fromSVOne (SVMap _) = error "fromSVOne: expected a single e-class" +{-# INLINE fromSVOne #-} + +-- | Returns a Query (list of atoms) of a pattern with pre-computed ordered vars +compileToQuery :: Pattern -> (Query, [ClassOrVar], ClassOrVar) +compileToQuery pat = (atoms, orderedVars atoms, root) + where (atoms, root) = evalState (processPat pat) 256 -- creates the atoms of a pattern - processPat :: Pattern -> State Int (Query, ClassOrVar) - processPat (VarPat x) = pure ([], Right $ fromEnum x) - processPat (Fixed pat) = do - -- get the next available var id and add as root - v <- get - let root = Right v - -- updates the next available id - modify (+1) - -- recursivelly process the children of the pattern - patChilds <- mapM processPat (getElems pat) - -- create an atom composed of the - -- root and the tree with the children - -- replaced by the childs roots - -- add the child atoms to the list - let atoms = concatMap fst patChilds - roots = map snd patChilds - atom = Atom root (replaceChildren roots pat) - atoms' = atom:atoms - pure (atoms', root) + processPat :: Pattern -> State Int (Query, ClassOrVar) + processPat (VarPat x) = pure ([], Right $ fromEnum x) + processPat (NAry _ _) = error "compileToQuery: n-ary pattern (use matchNAry instead)" + processPat Hole = error "compileToQuery: Hole is only valid in MapP targets" + processPat (Fixed pat) = do + -- get the next available var id and add as root + v <- get + let root = Right v + -- updates the next available id + modify (+1) + -- recursivelly process the children of the pattern + patChilds <- mapM processPat (getElems pat) + -- create an atom composed of the + -- root and the tree with the children + -- replaced by the childs roots + -- add the child atoms to the list + let atoms = concatMap fst patChilds + roots = map snd patChilds + atom = Atom root (replaceChildren roots pat) + atoms' = atom:atoms + pure (atoms', root) {-# INLINE compileToQuery #-} -- get the value from the Either Int Int @@ -226,63 +561,51 @@ -- | Creates the substituion map for -- the pattern variables for each one of the -- matched subgraph -genericJoin :: Monad m => Query -> ClassOrVar -> EGraphST m [Map ClassOrVar ClassOrVar] -genericJoin atoms root = do - let vars = orderedVars atoms -- order the vars, starting with the most frequently occuring - go atoms vars -- TODO: investigate why we need nub +genericJoin :: (ClassStore m, HasCallStack) => Query -> [ClassOrVar] -> ClassOrVar -> EGraphST m [Subst] +genericJoin atoms vars root = go atoms vars where -- for each variable -- for each possible e-class id for that variable -- replace the var id with this e-class id, and -- recurse to find the possible matches for the next atom - go :: Monad m => Query -> [ClassOrVar] -> EGraphST m [Map ClassOrVar ClassOrVar] + go :: ClassStore m => Query -> [ClassOrVar] -> EGraphST m [Subst] go atoms [] = pure [Map.empty] -- | _ <- atoms] go atoms (x:vars) = do cIds1 <- domainX x atoms root maps <- forM cIds1 $ \classId -> do - map (Map.insert x classId) <$> go (updateVar x classId atoms) vars + map (Map.insert x (SVOne 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] -- | returns the e-class id for a certain variable that -- matches the pattern described by the atoms -domainX :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar] +domainX :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar] 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 - -- $ - -- in ss -- | returns all e-class id that can matches this sequence of atoms -intersectAtoms :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId] +intersectAtoms :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId] intersectAtoms _ [] root = pure [] intersectAtoms var (a:atoms) root = do - a0 <- go a - Set.toList <$> (foldM (\acc atom -> Set.intersection acc <$> go atom) a0 atoms) + a0 <- toCanon =<< go a + Set.toList <$> (foldM (\acc atom -> do + res <- go atom + Set.intersection acc <$> toCanon res) a0 atoms) where - -- canonize everything except the root for consistency - -- doing this here prevents traversing the map again toCanon x = if var==root then pure x else Set.fromList <$> (mapM canonical $ Set.toList x) - go (Atom r t) = do - let op = getOperator t - mTrie <- gets ((Map.!? op) . _patDB . _eDB) - case mTrie of - Just trie -> pure (fromMaybe Set.empty $ intersectTries var Map.empty trie (r:getElems t)) - Nothing -> pure Set.empty - -- TODO: remove FlexibleContexts - --if op `Map.member` db -- if the e-graph contains the operator - -- try to find an intersection of the tries that matches each atom of the pattern - -- then - -- else pure Set.empty + go (Atom r t) = + do let op = getOperator t + mTrie <- gets ((Map.!? op) . _patDB . _eDB) + case mTrie of + Just trie -> pure (fromMaybe Set.empty $ intersectTries var IntMap.empty trie (r:getElems t)) + Nothing -> pure Set.empty + {-# INLINE intersectAtoms #-} -- | searches for the intersection of e-class ids that @@ -294,40 +617,26 @@ -- trie is the current trie of the pattern -- (i:ids) sequence of root : children of the atom to investigate -- NOTE: it must be Maybe Set to differentiate between empty set and no answer -intersectTries :: ClassOrVar -> Map ClassOrVar EClassId -> IntTrie -> [ClassOrVar] -> Maybe (HashSet EClassId) +intersectTries :: ClassOrVar -> IntMap EClassId -> IntTrie -> [ClassOrVar] -> Maybe (HashSet EClassId) intersectTries var xs trie [] = Just Set.empty intersectTries var xs trie (i:ids) = case i of - Left x -> if x `Set.member` _keys trie - -- if the current investigated id is an e-class id and - -- it is one of the keys of the trie... - -- ..try to match the next id with the next trie - then intersectTries var xs (_trie trie IntMap.! x) ids - else Nothing - Right x -> if i `Map.member` xs - -- if it is a pattern variable under investigation - -- and the e-class id is part of the trie - then if xs Map.! i `Set.member` _keys trie - -- match the next id with the next trie - then intersectTries var xs (_trie trie IntMap.! (xs Map.! i)) ids - else Nothing + Left x -> case IntMap.lookup x (_trie trie) of + Just subtrie -> intersectTries var xs subtrie ids + Nothing -> Nothing + Right x -> if IntMap.member x xs + then case IntMap.lookup (xs IntMap.! x) (_trie trie) of + Just subtrie -> intersectTries var xs subtrie ids + Nothing -> Nothing else if Right x == var - -- not under investigation and is the var of interest then if all (isDiffFrom x) ids - -- if there are no other occurrence of x in the next vars, - -- the keys of the trie are all possible candidates - then Just $ _keys trie - -- oterwise, put i under investigation and check the next occurrences - -- returning the intersection + then Just $ Set.fromList (IntMap.keys (_trie trie)) else Just $ IntMap.foldrWithKey (\k v acc -> - case intersectTries var (Map.insert i k xs) v ids of + case intersectTries var (IntMap.insert x k xs) v ids of Nothing -> acc _ -> Set.insert k acc) Set.empty (_trie trie) - -- if it is not the var of interest - -- assign and test all possible e-class ids to it - -- and move forward else Just $ IntMap.foldrWithKey (\k v acc -> - case intersectTries var (Map.insert i k xs) v ids of + case intersectTries var (IntMap.insert x k xs) v ids of Nothing -> acc Just s -> Set.union acc s ) Set.empty (_trie trie) @@ -359,15 +668,28 @@ {-# INLINE elemOfAtom #-} -- | sorts the variables in a query by the most frequently occurring +-- Ties are broken by putting an atom ROOT first. The root indexes the +-- operator trie directly, so matching it first replaces repeated whole-trie +-- folds (O(candidates x nodes)) with direct per-node trie descents. The old +-- tie-break (by id) put low-id pattern leaves before the high-id fresh root, +-- which made the root's domain include every operator node regardless of the +-- already-bound children (over-enumeration and O(n^2) folds). +-- Measured on the user config: 33s -> 19s (MT -N8), best loss unchanged. orderedVars :: Query -> [ClassOrVar] -orderedVars atoms = sortBy (comparing varCost) $ nub [a | atom <- atoms, a <- getIdsFrom atom, isRight a] +orderedVars atoms = sortBy (comparing key) $ RangeSet.toList $ RangeSet.fromList [a | atom <- atoms, a <- getIdsFrom atom, isRight a] where getIdsFrom (Atom r t) = r : getElems t isRight (Right _) = True isRight _ = False + -- is the variable the ROOT of some atom (an index into the operator trie)? + isHeader v = any (\a -> case a of Atom r _ -> r == v) atoms + varCost :: ClassOrVar -> Int varCost var = foldr (\a acc -> if elemOfAtom var a then acc - 100 + atomLen a else acc) 0 atoms + + key :: ClassOrVar -> (Int, Int) + key v = (varCost v, if isHeader v then 0 else 1) atomLen (Atom _ t) = 1 + length (getElems t) {-# INLINE orderedVars #-} diff --git a/src/Algorithm/EqSat/Egraph.hs b/src/Algorithm/EqSat/Egraph.hs --- a/src/Algorithm/EqSat/Egraph.hs +++ b/src/Algorithm/EqSat/Egraph.hs @@ -1,9 +1,10 @@ {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE StrictData #-} -{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} ----------------------------------------------------------------------------- -- | -- Module : Algorithm.EqSat.Egraph @@ -21,115 +22,100 @@ module Algorithm.EqSat.Egraph where import Control.Lens (element, makeLenses, view, over, (&), (+~), (-~), (.~), (^.)) ---import Control.Monad (forM, forM_, when, foldM, void) -import Data.List ( intercalate ) +--import Control.Monad (forM_, when, foldM, void) +import Data.List ( intercalate, foldl' ) +import Control.Monad (forM) import Control.Monad.State.Strict hiding ( get, put ) +import Control.Monad.IO.Class (MonadIO(..)) +import Data.Functor.Identity (Identity) +import GHC.Stack (HasCallStack) +import System.Random (StdGen) import Data.IntMap.Strict (IntMap) import qualified Data.IntMap.Strict as IntMap import Data.Map.Strict (Map) import qualified Data.Map.Strict as Map +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap import Data.HashSet (HashSet) import qualified Data.HashSet as Set import Data.IntSet (IntSet) import qualified Data.IntSet as IntSet -import Data.Sequence ( Seq(..), (><) ) -import qualified Data.Sequence as FingerTree -import Data.Foldable ( toList ) +import qualified Data.Set as RangeSet import Data.SRTree import Data.SRTree.Eval +import Data.SRTree.Recursion (cata) import Data.Hashable import Data.Binary import qualified Data.Binary as Bin -import qualified Data.Massiv.Array as MA +import qualified Data.Vector.Unboxed as VU +import Control.DeepSeq (NFData) import GHC.Generics -import Debug.Trace type EClassId = Int -- NOTE: DO NOT CHANGE THIS, this will break the use of IntMap and IntSet type ClassIdMap = IntMap -type ENode = SRTree EClassId -type ENodeEnc = (Int, Int, Int, Double) + +-- | N-ary operators represented as flattened multisets inside the e-graph. +-- Only Add and Mul are associative-commutative in this library; the remaining +-- ops (Sub, Div, Power, PowerAbs, AQ) stay binary and live in 'EBin'. +data NOp = EAdd | EMul deriving (Show, Eq, Ord, Enum, Generic, NFData) + +-- | The e-graph's node language. +-- +-- 'ENAry' stores Add/Mul as a canonical multiset of e-class ids: children are +-- path-compressed, keys sorted by canonical 'EClassId' (commutativity), and +-- nested same-op ENAry children are absorbed at insertion time +-- (associativity), so no commutativity/associativity rewrite rules are needed +-- for Add/Mul. The children are an 'IntMap' of e-class id to multiplicity. +data ENode + = EVar {-# UNPACK #-} !Int + | EParam {-# UNPACK #-} !Int + | EConst {-# UNPACK #-} !Double + | EUni Function EClassId + | EBin Op EClassId EClassId -- Sub | Div | Power | PowerAbs | AQ + | ENAry NOp (IntMap Int) -- canonical multiset: eclass -> multiplicity + deriving (Show, Eq, Generic, NFData) + type EGraphST m a = StateT EGraph m a type Cost = Int type CostFun = SRTree Cost -> Cost -type ECache = IntMap.IntMap PVector +type ECache = IntMap.IntMap Target +instance Hashable NOp where + hashWithSalt n EAdd = n `hashWithSalt` (0 :: Int) + hashWithSalt n EMul = n `hashWithSalt` (1 :: Int) + instance Hashable ENode where - hashWithSalt n enode = hashWithSalt n (encodeEnode enode) + hashWithSalt n (EVar ix) = n `hashWithSalt` (0 :: Int) `hashWithSalt` ix + hashWithSalt n (EParam ix) = n `hashWithSalt` (1 :: Int) `hashWithSalt` ix + hashWithSalt n (EConst x) = n `hashWithSalt` (2 :: Int) `hashWithSalt` x + hashWithSalt n (EUni f t) = n `hashWithSalt` (3 :: Int) `hashWithSalt` (fromEnum f) `hashWithSalt` t + hashWithSalt n (EBin op l r) = n `hashWithSalt` (4 :: Int) `hashWithSalt` (fromEnum op) `hashWithSalt` l `hashWithSalt` r + hashWithSalt n (ENAry op m) = n `hashWithSalt` (5 :: Int) `hashWithSalt` op `hashWithSalt` m -type RangeTree a = Seq (a, EClassId) +type RangeTree a = RangeSet.Set (a, EClassId) --- | this assumes up to 999 variables and params -encodeEnode :: ENode -> ENodeEnc ---encodeEnode = id -{--} -encodeEnode (Var ix) = (0, ix, -1, 0) -encodeEnode (Param ix) = (1, ix, -1, 0) -encodeEnode (Const x) = (2, -1, -1, x) -encodeEnode (Uni f ed) = (300 + fromEnum f, ed, -1, 0) -encodeEnode (Bin op ed1 ed2) = (400 + fromEnum op, ed1, ed2, 0) -{--} -{-# INLINE encodeEnode #-} +-- | Expand a canonical multiset back to the equivalent (multi-)set of child +-- e-class ids, one entry per occurrence. +expandedList :: IntMap Int -> [EClassId] +expandedList = concatMap (\(k, n) -> replicate n k) . IntMap.toAscList +{-# INLINE expandedList #-} -decodeEnode :: ENodeEnc -> ENode ---decodeEnode = id -{--} -decodeEnode (0, ix, _, _) = Var ix -decodeEnode (1, ix, _, _) = Param ix -decodeEnode (2, _, _, x) = Const x -decodeEnode (opCode, arg1, arg2, arg3) - | opCode < 400 = Uni (toEnum $ opCode-300) arg1 - | otherwise = Bin (toEnum $ opCode-400) arg1 arg2 - {--} -{-# INLINE decodeEnode #-} +-- | Build a canonical multiset from a list of child ids (duplicates allowed). +imFromList :: [EClassId] -> IntMap Int +imFromList = IntMap.fromListWith (+) . map (, 1) +{-# INLINE imFromList #-} + + insertRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a -insertRange eid x Empty = FingerTree.singleton (x, eid) -insertRange eid x (y :<| _xs) | (x, eid) < y = (x, eid) :<| y :<| _xs -insertRange eid x (_xs :|> y) | (x, eid) > y = _xs :|> y :|> (x, eid) -insertRange eid x rt = go rt - where - entry = (x, eid) - go root = case FingerTree.splitAt (n `div` 2) root of - (Empty, Empty) -> FingerTree.singleton entry - (Empty, z :<| zs) | entry < z -> entry :<| z :<| zs - | otherwise -> z :<| (go zs) - (ys :|> y, Empty) | entry > y -> ys :|> y :|> entry - | otherwise -> (go ys) :|> y - (ys :|> y, z :<| zs) - | entry > y && entry < z -> (ys :|> y :|> entry) >< (z :<| zs) - | entry > z -> (ys :|> y) >< go (z :<| zs) - | entry < y -> go (ys :|> y) >< (z :<| zs) - | otherwise -> root - where - n = FingerTree.length root +insertRange eid x = RangeSet.insert (x, eid) +{-# INLINE insertRange #-} removeRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a -removeRange eid x Empty = Empty -removeRange eid x (y :<| _xs) | (x, eid) < y = (y :<| _xs) -removeRange eid x (_xs :|> y) | (x, eid) > y = (_xs :|> y) -removeRange eid x rt = go rt - where - entry = (x, eid) - go root = case FingerTree.splitAt (n `div` 2) root of - (Empty, Empty) -> root - (Empty, z :<| zs) - | entry < z -> z :<| zs - | entry == z -> zs - | otherwise -> z :<| (go zs) - (ys :|> y, Empty) - | entry > y -> ys :|> y - | entry == y -> ys - | otherwise -> (go ys) :|> y - (ys :|> y, z :<| zs) - | entry > y && entry < z -> root - | entry > z -> (ys :|> y) >< go (z :<| zs) - | entry < y -> go (ys :|> y) >< (z :<| zs) - | otherwise -> root - - where - n = FingerTree.length root +removeRange eid x = RangeSet.delete (x, eid) +{-# INLINE removeRange #-} @@ -137,46 +123,50 @@ -- TODO: check this \/ getWithinRange :: Ord a => a -> a -> RangeTree a -> [EClassId] -getWithinRange lb ub rt = map snd . toList $ go rt - where - go Empty = Empty - go root = case FingerTree.splitAt (n `div` 2) root of - (Empty, Empty) -> Empty - (ys :|> y, Empty) - | fst y < lb -> Empty - | otherwise -> go (ys :|> y) - (Empty, z :<| zs) - | fst z > ub -> Empty - | otherwise -> go (z :<| zs) - (ys :|> y, z :<| zs) - | fst y < lb -> go (z :<| zs) - | fst z > ub -> go (ys :|> y) - | otherwise -> go (ys :|> y) >< go (z :<| zs) - where - n = FingerTree.length root - +getWithinRange lb ub rt = + let (_, ge) = RangeSet.split (lb, minBound) rt + (inR, _) = RangeSet.split (ub, maxBound) ge + in map snd (RangeSet.toList inR) -getSmallest :: Ord a => RangeTree a -> (a, EClassId) -getSmallest rt = case rt of - Empty -> error "empty finger" - x :<| t -> x +getSmallest :: Ord a => RangeTree a -> Maybe (a, EClassId) +getSmallest = RangeSet.lookupMin {-# INLINE getSmallest #-} -getGreatest :: Ord a => RangeTree a -> (a, EClassId) -getGreatest rt = case rt of - Empty -> error "empty finger" - t :|> x -> x +getGreatest :: Ord a => RangeTree a -> Maybe (a, EClassId) +getGreatest = RangeSet.lookupMax {-# INLINE getGreatest #-} +-- | Handle to an external, lazily paged e-class store (provided by the +-- storage layer, e.g. srtree-db's 'PageStore'). An 'EGraph' carries one when +-- e-classes are backed by a database; the IO actions fetch / persist / +-- evict a single e-class page. 'Nothing' keeps the classic fully-resident +-- behaviour. +data EClassPageStore = EClassPageStore + { cpsLookup :: EClassId -> IO (Maybe EClass) + , cpsInsert :: EClass -> IO () + , cpsDelete :: EClassId -> IO () + , cpsFlush :: IO () -- ^ write back all pending dirty pages + , cpsAll :: IO [EClass] -- ^ all e-classes currently in the store + , cpsKeys :: IO [EClassId] -- ^ all e-class ids currently in the store + , cpsStreamRoots :: SRTree () -> Int -> [EClassId] -> IO [EClassId] -- ^ bounded candidate roots for an operator, skipping an attempted set + , cpsRecordNode :: ENode -> EClassId -> IO () -- ^ register a newly-created node for write-back + , cpsNodeToClass :: ENode -> IO (Maybe EClassId) -- ^ content-address node -> class lookup (live) + , cpsCanonicalOf :: EClassId -> IO (Maybe EClassId) -- ^ e-class -> canonical representative (live) + , cpsRecordCanonical :: EClassId -> EClassId -> IO () -- ^ persist a canonical mapping (write-back) + , cpsBeginFrontier :: IO () -- ^ start a frontier re-saturation (restrict matcher to changed classes) + , cpsEndFrontier :: IO () -- ^ end it: clear the frontier (a pass re-saturated everything) + } + 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 + , _eNodeToEClass :: HashMap ENode EClassId -- maps an e-node to its e-class id + , _eClass :: ClassIdMap EClass -- maps an e-class id to its e-class data (resident cache) , _eDB :: EGraphDB - } deriving (Show, Generic) + , _classStore :: Maybe EClassPageStore -- optional lazily paged store for _eClass + } 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 + , _refits :: IntSet , _patDB :: DB -- database of patterns , _fitRangeDB :: RangeTree Double -- database of valid fitness , _dlRangeDB :: RangeTree Double @@ -184,26 +174,29 @@ , _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, Generic) + , _nextId :: Int -- next available id + , _changed :: !Bool -- dirty flag: true if modified since last check + , _trackDBs :: !Bool -- maintain range DBs (False during pure simplify) + , _seenMatches :: Map String (RangeSet.Set String) -- persistent (rule source -> attempted match keys) + } 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 +data EClass = EClass { _eClassId :: {-# UNPACK #-} !Int -- e-class id (maybe we don't need that here) + , _eNodes :: HashSet ENode -- set of e-nodes inside this e-class , _parents :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s - , _height :: Int -- height + , _height :: {-# UNPACK #-} !Int -- height , _info :: EClassData -- data } deriving (Show, Eq, Generic) -data Consts = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq, Generic) +data Consts = NotConst | ParamIx {-# UNPACK #-} !Int | ConstVal {-# UNPACK #-} !Double deriving (Show, Eq, Generic) data Property = Positive | Negative | NonZero | Real deriving (Show, Eq, Generic) -- TODO: incorporate properties -data EClassData = EData { _cost :: Cost +data EClassData = EData { _cost :: {-# UNPACK #-} !Cost , _best :: ENode , _consts :: Consts , _fitness :: Maybe Double -- NOTE: this cannot be NaN , _dl :: Maybe Double - , _theta :: [PVector] - , _size :: Int + , _theta :: [Target] + , _size :: {-# UNPACK #-} !Int -- , _properties :: Property -- TODO: include evaluation of expression from this e-class } deriving (Show, Generic) @@ -211,21 +204,32 @@ -- * 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 +instance Binary NOp where + put EAdd = put (0 :: Word8) + put EMul = put (1 :: Word8) 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 + 0 -> pure EAdd + 1 -> pure EMul +instance Binary ENode where + put (EVar ix) = put (0 :: Word8) >> put ix + put (EParam ix) = put (1 :: Word8) >> put ix + put (EConst x) = put (2 :: Word8) >> put x + put (EUni f t) = put (3 :: Word8) >> put (fromEnum f) >> put t + put (EBin op l r) = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r + put (ENAry op m) = put (5 :: Word8) >> put op >> put (expandedList m) + + get = do t <- get :: Get Word8 + case t of + 0 -> EVar <$> get + 1 -> EParam <$> get + 2 -> EConst <$> get + 3 -> EUni <$> (toEnum <$> get) <*> get + 4 -> EBin <$> (toEnum <$> get) <*> get <*> get + 5 -> ENAry <$> get <*> (imFromList <$> get) + instance Binary (SRTree ()) where put (Var ix) = put (0 :: Word8) >> put ix put (Param ix) = put (1 :: Word8) >> put ix @@ -245,17 +249,30 @@ 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 k, Binary v, Hashable k, Eq k) => Binary (HashMap k v) where + put hm = put (HashMap.toList hm) + get = HashMap.fromList <$> get +instance Binary Target where + put xs = put (VU.toList xs) + get = VU.fromList <$> get + instance Binary IntTrie instance Binary EClass instance Binary Consts instance Binary Property instance Binary EClassData -instance Binary EGraphDB -instance Binary EGraph +-- Custom: keep `_trackDBs` out of the wire format so on-disk EGraphDB data +-- (written before the flag existed) decodes unchanged; it defaults to True. +instance Binary EGraphDB where + put (EDB w a r p f d s sf sdl u n c _ _) = + put w >> put a >> put r >> put p >> put f >> put d >> put s >> put sf >> put sdl >> put u >> put n >> put c + get = EDB <$> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> pure True <*> pure Map.empty +-- Custom: the wire format omits `_classStore` (a runtime handle to the paged +-- store, never serialized); it decodes to Nothing. +instance Binary EGraph where + put (EGraph c n e d _) = put c >> put n >> put e >> put d + get = EGraph <$> get <*> get <*> get <*> get <*> pure Nothing instance Eq EClassData where 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 @@ -267,92 +284,642 @@ -- 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 (Generic) +newtype IntTrie = IntTrie { _trie :: IntMap IntTrie } deriving (Generic) --- Shows the IntTrie as {keys} -> {show IntTries} instance Show IntTrie where - show (IntTrie k t) = let keys = intercalate "," (map show $ Set.toList k) - tries = intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t) - in "{" <> keys <> "} - {" <> tries <> "}" + show (IntTrie t) = "{" <> intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t) <> "}" makeLenses ''EGraph makeLenses ''EClass makeLenses ''EClassData makeLenses ''EGraphDB +-- * Paged e-class access + +-- | A monad that can serve e-class data. +-- +-- The pure instances ('Identity', 'State StdGen') serve classes from the +-- resident @_eClass@ map; the 'MonadIO' instance consults the optional +-- 'EClassPageStore' when the graph carries one, falling back to the resident +-- map otherwise. All e-class read/write goes through these accessors, which +-- are the single choke point for a paged (out-of-core) e-graph. +class Monad m => ClassStore m where + lookupClass :: EClassId -> EGraphST m (Maybe EClass) + getClass :: HasCallStack => EClassId -> EGraphST m EClass + insertClass :: EClass -> EGraphST m () + deleteClass :: EClassId -> EGraphST m () + adjustClass :: EClassId -> (EClass -> EClass) -> EGraphST m () + -- | Enumerate every e-class (ids / values) in the graph. Paged graphs stream + -- from the store; resident graphs read the full @_eClass@ map. + allClasses :: EGraphST m [EClass] + allKeys :: EGraphST m [EClassId] + -- | Read/write a class directly from/to the backing store, bypassing the + -- resident LRU cache (and its O(n) 'trimResidentCache'). Bulk single-pass + -- traversals such as 'recalculateBestAllStream' must use these: routing every + -- one of ~n classes through 'lookupClass'/'insertClass' inserts each into the + -- resident map and calls 'trimResidentCache' (a full O(n) rebuild) after each + -- write, degenerating to O(n^2) and never terminating at scale. + readDirect :: EClassId -> EGraphST m (Maybe EClass) + writeDirect :: EClass -> EGraphST m () + allClasses = gets (IntMap.elems . _eClass) + allKeys = gets (IntMap.keys . _eClass) + readDirect = lookupClass + writeDirect = insertClass + -- | Enumerate (bounded) candidate e-class ids that contain a node of the + -- given operator, to drive the streaming matcher, skipping any ids in + -- @exclude@ (the already-attempted seen-set, so the per-rule budget advances + -- to new roots across scheduler cycles). The default reads the resident + -- @_patDB@ trie (the fully-in-RAM path); a paged graph streams the candidates + -- from its backing store instead, so the matcher never builds an O(nodes) + -- structure. + streamRoots :: SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId] + streamRoots = streamRootsFromDB + -- | Record a newly-created e-node (and its e-class) so a streaming matcher's + -- candidate source can see it. The default (fully resident graph) is a no-op: + -- the resident @_patDB@ is already updated by 'addToDB'. + recordNode :: ENode -> EClassId -> EGraphST m () + recordNode _ _ = pure () + -- | Content-address node -> class lookup. The default reads the resident + -- @_eNodeToEClass@ map (complete for a resident graph); a paged graph bounds + -- that map and falls back to the backing store on a miss. + lookupNode :: ENode -> EGraphST m (Maybe EClassId) + lookupNode en = gets (HashMap.lookup en . _eNodeToEClass) + -- | Record a node -> class mapping. The default keeps the resident (full) + -- map; a paged graph bounds it (evicting, since the store is authoritative). + insertNode :: ENode -> EClassId -> EGraphST m () + insertNode en eid = modify' $ over eNodeToEClass (HashMap.insert en eid) + -- | Record a canonical mapping (e-class -> representative), persisting it on a + -- paged graph so the store-backed canonical lookup sees merges/new classes. + insertCanonical :: EClassId -> EClassId -> EGraphST m () + insertCanonical eid canon = modify' $ over canonicalMap (IntMap.insert eid canon) + -- | The canonical representative of an e-class, or @Nothing@ when unknown. The + -- default reads the resident @_canonicalMap@; a paged graph bounds it and + -- falls back to the store. + canonicalOf :: EClassId -> EGraphST m (Maybe EClassId) + canonicalOf eid = gets (IntMap.lookup eid . _canonicalMap) + +-- | Default candidate-root enumeration from the resident @_patDB@ trie, capped +-- at @budget@ after skipping @exclude@ (used by the pure instances and as the +-- no-store fallback for a @MonadIO@ graph). +streamRootsFromDB :: Monad m => SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId] +streamRootsFromDB op budget exclude = do + db <- gets (_patDB . _eDB) + let ex = IntSet.fromList exclude + case Map.lookup op db of + Nothing -> pure [] + Just trie -> pure (take budget [ e | e <- IntMap.keys (_trie trie), not (IntSet.member e ex) ]) +{-# INLINE streamRootsFromDB #-} + +-- | Whether the graph is backed by a lazily paged e-class store. Streaming +-- matchers dispatch on this: a paged graph enumerates candidates from the +-- backing store (bounded memory), a resident graph from @_patDB@. +isPagedGraph :: Monad m => EGraphST m Bool +isPagedGraph = gets (maybe False (const True) . _classStore) +{-# INLINE isPagedGraph #-} + +-- Resident-map implementations (used by every pure monad) ------------------ + +pureLookupClass :: Monad m => EClassId -> EGraphST m (Maybe EClass) +pureLookupClass cid = gets (IntMap.lookup cid . _eClass) +{-# INLINE pureLookupClass #-} + +pureGetClass :: (Monad m, HasCallStack) => EClassId -> EGraphST m EClass +pureGetClass cid = do + m <- pureLookupClass cid + case m of + Just ec -> pure ec + Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid +{-# INLINE pureGetClass #-} + +pureInsertClass :: Monad m => EClass -> EGraphST m () +pureInsertClass ec = modify' $ over eClass (IntMap.insert (_eClassId ec) ec) +{-# INLINE pureInsertClass #-} + +pureDeleteClass :: Monad m => EClassId -> EGraphST m () +pureDeleteClass cid = modify' $ over eClass (IntMap.delete cid) +{-# INLINE pureDeleteClass #-} + +pureAdjustClass :: Monad m => EClassId -> (EClass -> EClass) -> EGraphST m () +pureAdjustClass cid f = modify' $ over eClass (IntMap.adjust f cid) +{-# INLINE pureAdjustClass #-} + +-- | Maximum number of e-classes kept in the resident @_eClass@ cache when the +-- graph is backed by a paged store. When exceeded, the largest-id classes are +-- retained and the rest evicted from the resident map. The store remains +-- authoritative (and Little-data reads fall back to it), so eviction only +-- bounds memory, never correctness. +residentClassCap :: Int +residentClassCap = 50000 + +-- | Trim the resident @_eClass@ cache to at most 'residentClassCap' entries +-- by keeping the largest ids. No-op for graphs without a paged store (their +-- resident map must stay complete for the pure instances). Halving on 2x keeps +-- steady churn from triggering an O(n) rebuild on every insert. +trimResidentCache :: Monad m => EGraphST m () +trimResidentCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _eClass eg + n = IntMap.size m + in if n <= 2 * residentClassCap + then eg + else over eClass (const (IntMap.fromList (Prelude.drop (n - residentClassCap) (IntMap.toAscList m)))) eg + +-- | Bound on the resident @_eNodeToEClass@ cache on a paged graph. Beyond the +-- cap (checked at 2x, halved back to cap) the map is pruned; the backing store +-- is authoritative, so eviction only trades a little dedup accuracy for bounded +-- memory, never correctness. +nodeCacheCap :: Int +nodeCacheCap = 100000 + +-- | Bound on the resident @_canonicalMap@ cache on a paged graph (same +-- halve-on-2x policy; evicted entries are re-read from the store). +canonicalCacheCap :: Int +canonicalCacheCap = 100000 +{-# INLINE nodeCacheCap #-} +{-# INLINE canonicalCacheCap #-} + +trimNodeCache :: Monad m => EGraphST m () +trimNodeCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _eNodeToEClass eg + n = HashMap.size m + in if n <= 2 * nodeCacheCap + then eg + else over eNodeToEClass (const (HashMap.fromList (Prelude.take nodeCacheCap (HashMap.toList m)))) eg +{-# INLINE trimNodeCache #-} + +trimCanonicalCache :: Monad m => EGraphST m () +trimCanonicalCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _canonicalMap eg + n = IntMap.size m + in if n <= 2 * canonicalCacheCap + then eg + else over canonicalMap (const (IntMap.fromList (Prelude.take canonicalCacheCap (IntMap.toAscList m)))) eg +{-# INLINE trimCanonicalCache #-} + +instance ClassStore Identity where + lookupClass = pureLookupClass + getClass = pureGetClass + insertClass = pureInsertClass + deleteClass = pureDeleteClass + adjustClass = pureAdjustClass + +instance ClassStore (State StdGen) where + lookupClass = pureLookupClass + getClass = pureGetClass + insertClass = pureInsertClass + deleteClass = pureDeleteClass + adjustClass = pureAdjustClass + +-- Any monad that can run IO is potentially paged: the graph's optional +-- store, when present, is authoritative; otherwise classes come from the +-- resident map. +instance {-# OVERLAPPABLE #-} (Monad m, MonadIO m) => ClassStore m where + -- The resident map is kept in sync by 'insertClass'/'deleteClass', so it is + -- consulted first: repeated reads never touch the store, and a class that + -- was evicted from the store's LRU while still dirty is never served stale. + lookupClass cid = do + eg <- gets id + case IntMap.lookup cid (_eClass eg) of + Just ec -> pure (Just ec) + Nothing -> case _classStore eg of + Nothing -> pure Nothing + Just h -> liftIO (cpsLookup h cid) + getClass cid = do + eg <- gets id + case IntMap.lookup cid (_eClass eg) of + Just ec -> pure ec + Nothing -> case _classStore eg of + Nothing -> pureGetClass cid + Just h -> do + m <- liftIO (cpsLookup h cid) + case m of + Just ec -> do + modify' (over eClass (IntMap.insert cid ec)) + trimResidentCache + pure ec + Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid + insertClass ec = do + eg <- gets id + case _classStore eg of + Nothing -> pureInsertClass ec + Just h -> do liftIO (cpsInsert h ec) + pureInsertClass ec + trimResidentCache + deleteClass cid = do + eg <- gets id + case _classStore eg of + Nothing -> pureDeleteClass cid + Just h -> do liftIO (cpsDelete h cid) + pureDeleteClass cid + adjustClass cid f = do + eg <- gets id + case _classStore eg of + Nothing -> pureAdjustClass cid f + Just _ -> do + m <- lookupClass cid + case m of + Nothing -> pure () + Just ec -> insertClass (f ec) + allClasses = do + eg <- gets id + case _classStore eg of + Nothing -> pure (IntMap.elems (_eClass eg)) + Just h -> liftIO (cpsAll h) + allKeys = do + eg <- gets id + case _classStore eg of + Nothing -> pure (IntMap.keys (_eClass eg)) + Just h -> liftIO (cpsKeys h) + -- Bypass the resident cache entirely: read the page straight from the store + -- and never insert into the (bounded) resident map, so a bulk traversal over + -- every class stays O(n) instead of O(n^2). + readDirect cid = do + eg <- gets id + case _classStore eg of + Nothing -> pureLookupClass cid + Just h -> liftIO (cpsLookup h cid) + writeDirect ec = do + eg <- gets id + case _classStore eg of + Nothing -> pureInsertClass ec + Just h -> liftIO (cpsInsert h ec) + streamRoots op budget exclude = do + eg <- gets id + case _classStore eg of + Nothing -> streamRootsFromDB op budget exclude + Just h -> liftIO (cpsStreamRoots h op budget exclude) + recordNode en eid = do + eg <- gets id + case _classStore eg of + Nothing -> pure () + Just h -> liftIO (cpsRecordNode h en eid) + lookupNode en = do + eg <- gets id + case _classStore eg of + Nothing -> gets (HashMap.lookup en . _eNodeToEClass) + Just h -> do + m <- gets (HashMap.lookup en . _eNodeToEClass) + case m of + Just eid -> pure (Just eid) + Nothing -> do + r <- liftIO (cpsNodeToClass h en) + case r of + Just eid -> do insertNode en eid + pure (Just eid) + Nothing -> pure Nothing + insertNode en eid = do + eg <- gets id + case _classStore eg of + Nothing -> modify' $ over eNodeToEClass (HashMap.insert en eid) + Just _ -> do modify' $ over eNodeToEClass (HashMap.insert en eid) + trimNodeCache + insertCanonical eid canon = do + eg <- gets id + case _classStore eg of + Nothing -> modify' $ over canonicalMap (IntMap.insert eid canon) + Just h -> do modify' $ over canonicalMap (IntMap.insert eid canon) + trimCanonicalCache + liftIO (cpsRecordCanonical h eid canon) + canonicalOf eid = do + eg <- gets id + case _classStore eg of + Nothing -> gets (IntMap.lookup eid . _canonicalMap) + Just h -> do + m <- gets (IntMap.lookup eid . _canonicalMap) + case m of + Just c -> pure (Just c) + Nothing -> do + r <- liftIO (cpsCanonicalOf h eid) + case r of + Just c -> do modify' $ over canonicalMap (IntMap.insert eid c) + trimCanonicalCache + pure (Just c) + Nothing -> pure Nothing + -- * E-Graph basic supporting functions -- | returns an empty e-graph emptyGraph :: EGraph -emptyGraph = EGraph IntMap.empty Map.empty IntMap.empty emptyDB +emptyGraph = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDB Nothing {-# INLINE emptyGraph #-} -- | returns an empty e-graph DB emptyDB :: EGraphDB -emptyDB = EDB Set.empty Set.empty Set.empty Map.empty FingerTree.empty FingerTree.empty IntMap.empty IntMap.empty IntMap.empty IntSet.empty 0 +emptyDB = EDB + Set.empty + Set.empty + IntSet.empty + Map.empty + RangeSet.empty + RangeSet.empty + IntMap.empty + IntMap.empty + IntMap.empty + IntSet.empty + 0 + False + True + Map.empty {-# INLINE emptyDB #-} +-- | like 'emptyDB' but skips range-DB maintenance (pure simplify mode) +emptyDBNoTrack :: EGraphDB +emptyDBNoTrack = emptyDB{ _trackDBs = False } +{-# INLINE emptyDBNoTrack #-} + +-- | an empty e-graph that skips range-DB maintenance (pure simplify mode) +emptyGraphNoTrack :: EGraph +emptyGraphNoTrack = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDBNoTrack Nothing +{-# INLINE emptyGraphNoTrack #-} + -- | Creates a new e-class from an e-class id, a new e-node, -- and the info of this e-class createEClass :: EClassId -> ENode -> EClassData -> Int -> EClass -createEClass cId enode' info h = EClass cId (Set.singleton $ encodeEnode enode') Set.empty h info +createEClass cId enode' info h = EClass cId (Set.singleton enode') Set.empty h info {-# INLINE createEClass #-} --- | gets the canonical id of an e-class -canonical :: Monad m => EClassId -> EGraphST m EClassId -canonical eclassId = - do m <- gets _canonicalMap - let oneStep = m IntMap.! eclassId - if oneStep == eclassId - then pure eclassId - else go m oneStep - where - go :: Monad m => IntMap EClassId -> EClassId -> EGraphST m EClassId - go m ecId - | m IntMap.! ecId == ecId = do modify' $ over canonicalMap (IntMap.insert eclassId ecId) -- creates a shortcut for next time - pure ecId -- if the e-class id is mapped to itself, it's canonical - | otherwise = go m (m IntMap.! ecId) -- otherwise, check the next id in the sequence +-- | gets the canonical id of an e-class with full path compression +canonical :: ClassStore m => EClassId -> EGraphST m EClassId +canonical eclassId = do + mStep <- canonicalOf eclassId + case mStep of + Nothing -> canonError eclassId + Just oneStep + | oneStep == eclassId -> pure eclassId + | otherwise -> do + (root, chain) <- walk [eclassId] oneStep + -- compress the chain in the resident cache (cache-only: the store + -- keeps the authoritative semantic mappings recorded at insert + -- time, so eviction just loses the shortcut, never correctness). + modify' $ \eg -> eg{ _canonicalMap = + foldl' (\m' k -> IntMap.insert k root m') (_canonicalMap eg) chain } + pure root + where + walk :: ClassStore m => [EClassId] -> EClassId -> EGraphST m (EClassId, [EClassId]) + walk chain ecId = do + mNext <- canonicalOf ecId + case mNext of + Nothing -> canonError ecId + Just n + | n == ecId -> pure (ecId, chain) + | otherwise -> walk (ecId : chain) n + + canonError :: ClassStore m => EClassId -> EGraphST m a + canonError eid = do + m <- gets _canonicalMap + error $ "CANON_MISSING eid=" <> show eid <> " mapSize=" <> show (IntMap.size m) {-# INLINE canonical #-} -- | canonize the e-node children -canonize :: Monad m => ENode -> EGraphST m ENode -canonize = mapM canonical -- applies canonical to the children +canonize :: (ClassStore m, HasCallStack) => ENode -> EGraphST m ENode +canonize (EVar ix) = pure (EVar ix) +canonize (EParam ix) = pure (EParam ix) +canonize (EConst x) = pure (EConst x) +canonize (EUni f t) = EUni f <$> canonical t +canonize (EBin op l r) = EBin op <$> canonical l <*> canonical r +-- re-map children to their canonical ids; IntMap keeps keys sorted, so +-- commutativity is structural, no rewrite rule required. +canonize (ENAry op m) = do + m' <- IntMap.fromListWith (+) <$> forM (IntMap.toList m) (\(c, n) -> do + c' <- canonical c + pure (c', n)) + pure (ENAry op m') {-# INLINE canonize #-} --- | gets an e-class with id `c` -getEClass :: Monad m => EClassId -> EGraphST m EClass -getEClass c = gets ((IntMap.! c) . _eClass) +-- | The children e-class ids of an e-node. +eChildren :: ENode -> [EClassId] +eChildren (EVar _) = [] +eChildren (EParam _) = [] +eChildren (EConst _) = [] +eChildren (EUni _ t) = [t] +eChildren (EBin _ l r) = [l, r] +eChildren (ENAry _ m) = expandedList m +{-# INLINE eChildren #-} + +toOp :: NOp -> Op +toOp EAdd = Add +toOp EMul = Mul +{-# INLINE toOp #-} + +-- | Operator shape key used to index the pattern database. ENAry maps back to +-- the corresponding binary operator shape so existing (binary) Add/Mul +-- patterns address the same trie. +eOpKey :: ENode -> SRTree () +eOpKey (EVar ix) = Var ix +eOpKey (EParam ix) = Param ix +eOpKey (EConst x) = Const x +eOpKey (EUni f _) = Uni f () +eOpKey (EBin op _ _) = Bin op () () +eOpKey (ENAry EAdd _) = Bin Add () () +eOpKey (ENAry EMul _) = Bin Mul () () +{-# INLINE eOpKey #-} + +-- | Convert an e-node (children still as e-class ids) into the equivalent +-- binary SRTree shape. NOTE: only called on non-ENary nodes; flattened +-- ENAry nodes have no binary skeleton (see 'naryTree' / the explicit ENAry +-- cases in the analyses). +fromENode :: ENode -> SRTree EClassId +fromENode (EVar ix) = Var ix +fromENode (EParam ix) = Param ix +fromENode (EConst x) = Const x +fromENode (EUni f t) = Uni f t +fromENode (EBin op l r) = Bin op l r +fromENode (ENAry _ _) = error "fromENode: ENAry has no binary skeleton" +{-# INLINE fromENode #-} + +-- | Right-fold a list of e-class child expressions into a binary Fix SRTree +-- for a flattened ENAry multiset (extraction). +naryTree :: NOp -> [Fix SRTree] -> Fix SRTree +naryTree op ts = normalizeSubDiv (foldr1 (\a b -> Fix (Bin (toOp op) a b)) ts) +{-# INLINE naryTree #-} + +-- | Re-render the internal negate/recip canonical forms back as Sub/Div so +-- extraction output keeps the familiar shape: `x + (-1)*y` -> `x - y`, +-- `x + (-3)` -> `x - 3` and `x * recip y` -> `x / y`. Sub and Div never +-- appear as e-nodes; they only reappear here during reconstruction. +normalizeSubDiv :: Fix SRTree -> Fix SRTree +normalizeSubDiv = cata alg + where + alg :: SRTree (Fix SRTree) -> Fix SRTree + alg (Bin Add l r) = case pick l r of + Just (pos, neg) -> Fix (Bin Sub pos neg) + Nothing -> Fix (Bin Add l r) + where + pick a b = case negated a of + Just t -> Just (b, t) + Nothing -> case negated b of + Just t -> Just (a, t) + Nothing -> Nothing + negated (Fix (Bin Mul (Fix (Const c)) t)) | c == -1 = Just t + negated (Fix (Bin Mul t (Fix (Const c)))) | c == -1 = Just t + negated (Fix (Const c)) | c < 0 = Just (Fix (Const (-c))) + negated _ = Nothing + alg (Bin Mul l r) = case pick l r of + Just (num, den) -> Fix (Bin Div num den) + Nothing -> Fix (Bin Mul l r) + where + pick a b = case a of + Fix (Uni Recip t) -> Just (b, t) + _ -> case b of + Fix (Uni Recip t) -> Just (a, t) + _ -> Nothing + alg t = Fix t + +-- | Convert a binary SRTree (children as e-class ids) into an e-node, +-- flattening Add/Mul into canonical ENAry multisets. +toENode :: (ClassStore m, HasCallStack) => SRTree EClassId -> EGraphST m ENode +toENode (Var ix) = pure (EVar ix) +toENode (Param ix) = pure (EParam ix) +toENode (Const x) = pure (EConst x) +toENode (Uni f t) = EUni f <$> canonical t +toENode (Bin Add l r) = mkENary EAdd [l, r] +toENode (Bin Mul l r) = mkENary EMul [l, r] +toENode (Bin op l r) = EBin op <$> canonical l <*> canonical r +toENode n = error $ "toENode: unsupported node " <> show n +{-# INLINE toENode #-} + +-- | Build a canonical ENAry from child ids: canonicalize children, absorb +-- nested same-op ENAry children (associativity), sort by key (commutativity). +mkENary :: (ClassStore m, HasCallStack) => NOp -> [EClassId] -> EGraphST m ENode +mkENary op cids = mkENaryM op (imFromList cids) + +-- | Build a canonical ENAry from a canonical multiset of child ids. +mkENaryM :: (ClassStore m, HasCallStack) => NOp -> IntMap Int -> EGraphST m ENode +mkENaryM op m = do + flat <- IntMap.unionsWith (+) <$> mapM (expandM op) (IntMap.toList m) + pure (ENAry op flat) + +-- | If the e-class of `cid` holds exactly one e-node and that node is an ENAry +-- of the same op, return its children scaled by `n` (flattening `n` +-- occurrences); otherwise return `n` copies of `cid`. Flattening is only sound +-- through a class with a single node: if the class were merged with other +-- nodes (e.g. `{Add[a,b], Mul[x,c]}`) flattening would silently pick one +-- representative and change the meaning of the term. +expandM :: (ClassStore m, HasCallStack) => NOp -> (EClassId, Int) -> EGraphST m (IntMap Int) +expandM op (cid, n) = do + ec <- getEClass cid + case Set.toList (_eNodes ec) of + [ENAry op' m'] | op' == op -> pure (IntMap.map (* n) m') + _ -> pure (IntMap.singleton cid n) + +-- | Reconstruct a binary Fix SRTree from an e-node, right-folding ENAry +-- into nested Bin Add/Mul. +enodeToTree :: (ClassStore m, HasCallStack) => ENode -> EGraphST m (Fix SRTree) +enodeToTree (EVar ix) = pure (Fix (Var ix)) +enodeToTree (EParam ix) = pure (Fix (Param ix)) +enodeToTree (EConst x) = pure (Fix (Const x)) +enodeToTree (EUni f t) = Fix . Uni f <$> getBestExpr t +enodeToTree (EBin op l r) = do + tl <- getBestExpr l + tr <- getBestExpr r + pure (Fix (Bin op tl tr)) +enodeToTree (ENAry op m) = do + ts <- mapM getBestExpr (expandedList m) + pure (naryTree op ts) +{-# INLINE enodeToTree #-} + +-- | gets an e-class with id `c` (auto-canonizes) +getEClass :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m EClass +getEClass c = do c' <- canonical c; getClass c' {-# INLINE getEClass #-} +-- | gets the best expression given the default cost function. Cycle-safe and +-- budgeted: see 'getBestExprBounded'. +getBestExpr :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree) +getBestExpr eid = getBestExprBounded eid + +-- | Like 'getBestExpr' but terminates on pathological graphs: a visited set +-- stops the expansion from re-entering an already-expanded class (a @_best@ +-- cycle arising from supersaturation/merges), and a node budget caps the total +-- expanded size (so an exponentially-shared DAG is truncated rather than +-- exploded). Both guards substitute a @Var 0@ placeholder for the part that +-- would otherwise blow up. On well-formed acyclic graphs with small bests +-- neither guard triggers, so the result is identical to the unbounded version. +-- This keeps out-of-core extraction (e.g. 'dbTop') bounded in memory. +getBestExprBounded :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree) +getBestExprBounded eid = fst <$> expand Set.empty 0 eid + where + budget :: Int + budget = 200 + -- expand returns the tree and the running count of expanded nodes, so the + -- budget bounds the TOTAL size (not just the depth): an exponentially-shared + -- DAG is truncated instead of exploded. A revisited (cyclic) class or a + -- full budget yields a @Var 0@ placeholder. + expand :: ClassStore m => HashSet EClassId -> Int -> EClassId -> EGraphST m (Fix SRTree, Int) + expand _ n _ | n >= budget = pure (Fix (Var 0), n) + expand seen n eid + | Set.member eid seen = pure (Fix (Var 0), n) + | otherwise = do + best <- (_best . _info) <$> getEClass eid + let seen' = Set.insert eid seen + n0 = n + 1 + case best of + EVar ix -> pure (Fix (Var ix), n0) + EParam ix -> pure (Fix (Param ix), n0) + EConst x -> pure (Fix (Const x), n0) + EUni f t -> do (tt, n1) <- expand seen' n0 t + pure (Fix (Uni f tt), n1) + EBin op l r -> do + (tl, n1) <- expand seen' n0 l + (tr, n2) <- expand seen' n1 r + pure (Fix (Bin op tl tr), n2) + ENAry op m -> do + (xs, nEnd) <- goNary seen' n0 (IntMap.toAscList m) [] + pure (if null xs then (Fix (Var 0), nEnd) else (naryTree op xs, nEnd)) + -- build the ENAry children from the multiset WITHOUT materialising the + -- expanded multiplicity list: an enormous count (a pathological supersaturated + -- class) is capped per-child and by the total budget, so each copy counts + -- toward the budget and no giant list is ever allocated. + goNary seen n es acc + | n >= budget = pure (reverse acc, n) + | otherwise = case es of + [] -> pure (reverse acc, n) + ((c, cnt) : rest) -> do + (t, n1) <- expand seen n c + let take = min cnt (budget - n1 + 1) + n2 = n1 + (take - 1) + acc' = Prelude.replicate take t ++ acc + goNary seen n2 rest acc' + -- | Creates a singleton trie from an e-class id trie :: EClassId -> IntMap IntTrie -> IntTrie -trie eid = IntTrie (Set.singleton eid) +trie eid = IntTrie {-# INLINE trie #-} -- | Check whether an e-class is a constant value -isConst :: Monad m => EClassId -> EGraphST m Bool -isConst eid = do ec <- gets ((IntMap.! eid) . _eClass) +isConst :: ClassStore m => EClassId -> EGraphST m Bool +isConst eid = do ec <- getEClass eid case (_consts . _info) ec of ConstVal _ -> pure True _ -> pure False {-# INLINE isConst #-} -getFitness :: Monad m => EClassId -> EGraphST m (Maybe Double) -getFitness c = gets (_fitness . _info . (IntMap.! c) . _eClass) +getFitness :: ClassStore m => EClassId -> EGraphST m (Maybe Double) +getFitness c = (_fitness . _info) <$> getEClass c {-# INLINE getFitness #-} -getTheta :: Monad m => EClassId -> EGraphST m ([PVector]) -getTheta c = gets (_theta . _info . (IntMap.! c) . _eClass) +getTheta :: ClassStore m => EClassId -> EGraphST m ([Target]) +getTheta c = (_theta . _info) <$> getEClass c {-# INLINE getTheta #-} -getSize :: Monad m => EClassId -> EGraphST m Int -getSize c = gets (_size . _info . (IntMap.! c) . _eClass) +getSize :: ClassStore m => EClassId -> EGraphST m Int +getSize c = (_size . _info) <$> getEClass c {-# INLINE getSize #-} isSizeOf :: (Int -> Bool) -> EClass -> Bool isSizeOf p = p . _size . _info {-# INLINE isSizeOf #-} -getBestFitness :: Monad m => EGraphST m (Maybe Double) +getBestFitness :: ClassStore m => EGraphST m (Maybe Double) getBestFitness = do - bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical) - gets (_fitness . _info . (IntMap.! bec) . _eClass) -getDL :: Monad m => EClassId -> EGraphST m (Maybe Double) -getDL c = gets (_dl . _info . (IntMap.! c) . _eClass) + mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB) + case mbec of + Just bec -> (_fitness . _info) <$> getEClass bec + Nothing -> pure Nothing +getDL :: ClassStore m => EClassId -> EGraphST m (Maybe Double) +getDL c = (_dl . _info) <$> getEClass c {-# INLINE getDL #-} diff --git a/src/Algorithm/EqSat/Info.hs b/src/Algorithm/EqSat/Info.hs --- a/src/Algorithm/EqSat/Info.hs +++ b/src/Algorithm/EqSat/Info.hs @@ -15,28 +15,24 @@ module Algorithm.EqSat.Info where import Control.Lens ( over ) -import Control.Monad --(forM, forM_, when, foldM, void) +import Control.Monad import Control.Monad.State import Data.AEq (AEq ((~==))) -import Data.IntMap (IntMap) -- , delete, empty, insert, toList) +import Data.IntMap (IntMap) import qualified Data.IntMap as IntMap import Data.Map (Map) import qualified Data.Map as Map import Data.SRTree -import Data.SRTree.Eval (evalFun, evalOp, PVector) +import Data.SRTree.Eval (evalFun, evalOp, Target) import Data.HashSet (HashSet) import qualified Data.HashSet as Set +import qualified Data.Set as RangeSet import qualified Data.IntSet as IntSet import Algorithm.EqSat.Egraph -import Data.AEq (AEq ((~==))) import Algorithm.EqSat.Queries -import Data.Maybe import qualified Data.Set as TrueSet -import Data.Sequence (Seq(..), (><)) -import Debug.Trace - -- * Data related functions -- | join data from two e-classes @@ -84,51 +80,71 @@ combineConsts (ConstVal x) (ParamIx ix) = ConstVal x -- p - p = 0 combineConsts x y = error (show x <> " " <> show y) +-- | Fetch consts, cost, and size for all children in a single state traversal +getChildrenData :: ClassStore m => [EClassId] -> EGraphST m [(Consts, Cost, Int)] +getChildrenData ids = do + ids' <- mapM canonical ids + mapM (\cid -> do + ec <- getEClass cid + let d = _info ec + pure (_consts d, _cost d, _size d)) ids' +{-# INLINE getChildrenData #-} + -- | Calculate e-node data (constant values and cost) -makeAnalysis :: Monad m => CostFun -> ENode -> EGraphST m EClassData +makeAnalysis :: ClassStore m => CostFun -> ENode -> EGraphST m EClassData makeAnalysis costFun enode = - do consts <- calculateConsts enode + do let cs = eChildren enode + childData <- getChildrenData cs + let (consts', costs', sizes) = unzip3 childData + consts = combineNode enode consts' + cost = costNode enode costs' + sz = sum sizes 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 [] (sz + 1) + where + -- ENAry folds children pairwise (constant folding over a multiset); the + -- binary skeleton cannot represent n children. + combineNode (ENAry op _) cs = foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) cs + combineNode _ cs = combineConsts (replaceChildren cs (fromENode enode)) + -- ENAry is a single flattened op node: op cost + sum of child costs. + costNode (ENAry op _) cs = costFun (Bin (toOp op) 0 0) + sum cs + costNode _ cs = costFun (replaceChildren cs (fromENode enode)) -getChildrenMinHeight :: Monad m => ENode -> EGraphST m Int +getChildrenMinHeight :: ClassStore m => ENode -> EGraphST m Int getChildrenMinHeight enode = do - let children = childrenOf enode - minimum' [] = 0 - minimum' xs = minimum xs - minimum' <$> mapM (\ec -> gets (_height . (IntMap.! ec) . _eClass)) children + let children = eChildren enode + if null children then pure 0 else do + children' <- mapM canonical children + hs <- mapM (fmap _height . getEClass) children' + pure (minimum hs) -- | update the heights of each e-class -- won't work if there's no root -calculateHeights :: Monad m => EGraphST m () +calculateHeights :: ClassStore m => EGraphST m () calculateHeights = do queue <- findRootClasses - classes <- gets (Prelude.map fst . IntMap.toList . _eClass) + classes <- allKeys let nClasses = length classes forM_ classes (setHeight nClasses) -- set all heights to max possible height (number of e-classes) forM_ queue (setHeight 0) -- set root e-classes height to zero go queue (TrueSet.fromList queue) 1 -- next height is 1 where + setHeight :: ClassStore m => Int -> EClassId -> EGraphST m () setHeight x eId' = do eId <- canonical eId' ec <- getEClass eId let ec' = over height (const x) ec - modify' $ over eClass (IntMap.insert eId ec') + insertClass ec' + setMinHeight :: ClassStore m => Int -> EClassId -> EGraphST m () setMinHeight x eId' = -- set height to the minimum between current and x do eId <- canonical eId' h <- _height <$> getEClass eId setHeight (min h x) eId - getChildrenEC :: Monad m => EClassId -> EGraphST m [EClassId] - getChildrenEC ec' = do ec <- canonical ec' - gets (concatMap childrenOf' . _eNodes . (IntMap.! ec) . _eClass) - - childrenOf' (_, -1, -1, _) = [] - childrenOf' (_, e1, -1, _) = [e1] - childrenOf' (_, e1, e2, _) = [e1, e2] + getChildrenEC :: ClassStore m => EClassId -> EGraphST m [EClassId] + getChildrenEC ec' = do ec <- getEClass ec' + pure $ concatMap eChildren (_eNodes ec) go [] _ _ = pure () go qs tabu h = @@ -138,19 +154,23 @@ go childrenL (TrueSet.union tabu childrenOf) (h+1) -- move one breadth search style -- | calculates the cost of a node -calculateCost :: Monad m => CostFun -> SRTree EClassId -> EGraphST m Cost -calculateCost f t = - do let cs = childrenOf t +calculateCost :: ClassStore m => CostFun -> ENode -> EGraphST m Cost +calculateCost f enode = + do let cs = eChildren enode costs <- traverse (fmap (_cost . _info) . getEClass) cs - pure . f $ replaceChildren costs t + pure $ case enode of + ENAry op _ -> f (Bin (toOp op) 0 0) + sum costs + _ -> f (replaceChildren costs (fromENode enode)) -- | check whether an e-node evaluates to a const -calculateConsts :: Monad m => SRTree EClassId -> EGraphST m Consts -calculateConsts t = - do let cs = childrenOf t - eg <- get +calculateConsts :: ClassStore m => ENode -> EGraphST m Consts +calculateConsts enode = + do let cs = eChildren enode consts <- traverse (fmap (_consts . _info) . getEClass) cs - case combineConsts $ replaceChildren consts t of + let c = case enode of + ENAry op _ -> foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) consts + _ -> combineConsts (replaceChildren consts (fromENode enode)) + case c of ConstVal x | isNaN x -> pure (ConstVal x) a -> pure a @@ -168,41 +188,35 @@ evalOp' (ConstVal x) (ConstVal y) = ConstVal $ evalOp op x y evalOp' _ _ = NotConst -insertFitness :: Monad m => EClassId -> Double -> [PVector] -> EGraphST m () -insertFitness eId' fit params = do - eId <- canonical eId' - tree <- getBestExpr' eId - let p = fromIntegral (length params) - let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0 - ec <- gets ((IntMap.! eId) . _eClass) - let oldFit = _fitness . _info $ ec - --when (oldFit < Just fit) $ do - let newInfo = (_info ec){_fitness = Just fit, _theta = params} - newEc = ec{_info = newInfo} - sz = _size newInfo - modify' $ over eClass (IntMap.insert eId newEc) - if (isNothing oldFit) - then modify' $ over (eDB . unevaluated) (IntSet.delete eId) - . over (eDB . fitRangeDB) (insertRange eId fit) - . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty) - . over (eDB . dlRangeDB) (insertRange eId f_compl) - else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit)) +insertFitness :: ClassStore m => EClassId -> Double -> [Target] -> EGraphST m () +insertFitness eId' fit params = + do eId <- canonical eId' + tree <- getBestExpr eId + let p = fromIntegral (length params) + let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0 + ec <- getEClass eId + let oldFit = _fitness . _info $ ec + let newInfo = (_info ec){_fitness = Just fit, _theta = params} + newEc = ec{_info = newInfo} + sz = _size newInfo + insertClass newEc + case oldFit of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.delete eId) + . over (eDB . fitRangeDB) (insertRange eId fit) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty) + . over (eDB . dlRangeDB) (insertRange eId f_compl) + Just oldVal -> modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId oldVal) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit . removeRange eId oldVal) sz) -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) +insertDL :: ClassStore m => EClassId -> Double -> EGraphST m () +insertDL eId fit' = + do let fit = negate fit' + ec <- getEClass eId + let sz = _size . _info $ ec + newInfo = (_info ec){_dl = Just fit'} + newEc = ec{_info=newInfo} + insertClass newEc + modify' $ over (eDB . dlRangeDB) (insertRange eId fit) + . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty) --- | TODO: remove from here gets the best expression given the default cost function -getBestExpr' :: Monad m => EClassId -> EGraphST m (Fix SRTree) -getBestExpr' eid = do eid' <- canonical eid - best <- gets (_best . _info . (IntMap.! eid') . _eClass) - childs <- mapM getBestExpr' $ childrenOf best - pure . Fix $ replaceChildren childs best diff --git a/src/Algorithm/EqSat/Queries.hs b/src/Algorithm/EqSat/Queries.hs --- a/src/Algorithm/EqSat/Queries.hs +++ b/src/Algorithm/EqSat/Queries.hs @@ -21,66 +21,54 @@ import qualified Data.IntMap as IntMap import qualified Data.Map as Map import qualified Data.HashSet as Set +import qualified Data.Set as RangeSet import Control.Monad.State ( gets, modify' ) -import Control.Monad ( filterM ) import Control.Lens ( over ) import Data.Maybe -import Data.Sequence ( Seq(..) ) -import qualified Data.Sequence as FingerTree -import qualified Data.Foldable as Foldable import Data.SRTree (childrenOf) -import Debug.Trace - --- this is too slow for now, it needs a db of its own --- basically a db for each query we need -getEClassesThat :: Monad m => (EClass -> Bool) -> EGraphST m [EClassId] +getEClassesThat :: ClassStore m => (EClass -> Bool) -> EGraphST m [EClassId] getEClassesThat p = do - gets (map fst . filter (\(ecId, ec) -> p ec) . IntMap.toList . _eClass) - --go ecs - where - go :: Monad m => [EClassId] -> EGraphST m [EClassId] - go [] = pure [] - go (ecId:ecs) = do ec <- gets (p . (IntMap.! ecId) . _eClass) - ecs' <- go ecs - if ec - then pure (ecId:ecs') - else pure ecs' + classes <- allClasses + pure [ _eClassId ec | ec <- classes, p ec ] -updateFitness :: Monad m => Double -> EClassId -> EGraphST m () +updateFitness :: ClassStore m => Double -> EClassId -> EGraphST m () updateFitness f ecId = do - ec <- gets ((IntMap.! ecId) . _eClass) + ec <- getEClass ecId let info = _info ec - modify' $ over eClass (IntMap.insert ecId ec{_info=info{_fitness = Just f}}) + insertClass ec{_info=info{_fitness = Just f}} -- | returns all the root e-classes (e-class without parents) -findRootClasses :: Monad m => EGraphST m [EClassId] -findRootClasses = gets (Prelude.map fst . Prelude.filter isParent . IntMap.toList . _eClass) +findRootClasses :: ClassStore m => EGraphST m [EClassId] +findRootClasses = do + classes <- allClasses + pure [ _eClassId ec | ec <- classes, isParent (_eClassId ec, ec) ] where isParent (k, v) = Prelude.null (_parents v) || (k `Set.member` (Set.map fst (_parents v))) -- | returns the e-class id with the best fitness that -- is true to a predicate -getTopECLassThat :: Monad m => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopECLassThat :: ClassStore 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] + go :: ClassStore 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 p ec - then go (m-1) (ecId:bests) t - else go m bests t + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests t + else if p ec + then go (m-1) (ecId:bests) t + else go m bests t -getTopEClassInRange :: Monad m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId] +getTopEClassInRange :: ClassStore m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId] getTopEClassInRange b n p range = do let f = if b then _fitRangeDB else _dlRangeDB gets (f . _eDB) @@ -92,43 +80,45 @@ | v > y = 1 | otherwise = 1 - go :: Monad m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId] + go :: ClassStore m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId] go _ bests [] _ = pure bests go 0 bests (r:rs) rt = go n bests rs rt - go m bests (r:rs) rt = case rt of - Empty -> pure bests - t :|> y -> do let x = snd y - ecId <- canonical x - ec <- gets ((IntMap.! ecId) . _eClass) - if (isInfinite . fromJust . _fitness . _info $ ec) - then go m bests (r:rs) t - else do let v = p ec - case (v `inRange` r) of - 0 -> go (m-1) (ecId:bests) (r:rs) t -- it is in range, go to the next range - -1 -> go n bests rs (t :|> y) -- it is smaller than the range, get the first n of the next range - 1 -> go m bests (r:rs) t -- y is still greater than the range, keep looking in the same range + go m bests (r:rs) rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests (r:rs) t + else do let v = p ec + case (v `inRange` r) of + 0 -> go (m-1) (ecId:bests) (r:rs) t + -1 -> go n bests rs (RangeSet.insert y t) + 1 -> go m bests (r:rs) t -getTopECLassIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopECLassIn :: ClassStore 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 :: ClassStore 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 + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests t + 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 :: ClassStore 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) @@ -136,65 +126,65 @@ where ecs = Set.fromList ecs' - go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go :: ClassStore 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 + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _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 :: ClassStore 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 + go :: ClassStore m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _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 || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t) + go m bests (Just rt) = case RangeSet.maxView rt of + Nothing -> bests + Just ((f, x), t) -> 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 :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopFitEClassThat = getTopECLassThat True -getTopDLEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopDLEClassThat :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] getTopDLEClassThat = getTopECLassThat False -getTopFitEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopFitEClassIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopFitEClassIn = getTopECLassIn True -getTopDLEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopDLEClassIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopDLEClassIn = getTopECLassIn False -getTopFitEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopFitEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] getTopFitEClassNotIn = getTopECLassNotIn True -getTopDLEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] -getTopDLEClassNotIn = getTopECLassNotIn True +getTopDLEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopDLEClassNotIn = getTopECLassNotIn False 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 :: ClassStore m => EGraphST m () rebuildAllRanges = do szF <- gets (_sizeFitDB._eDB) >>= traverse rebuildRange dlF <- gets (_sizeDLDB._eDB) >>= traverse rebuildRange fR <- gets (_fitRangeDB._eDB) >>= rebuildRange @@ -205,17 +195,19 @@ . 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) +canonizeRange :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double) +canonizeRange = fmap RangeSet.fromList . mapM (\(x, eid) -> (x,) <$> canonical eid) . RangeSet.toList -rebuildRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double) -rebuildRange rt = go Set.empty Empty <$> canonizeRange rt +rebuildRange :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double) +rebuildRange rt = do + canonRt <- canonizeRange rt + pure $ snd $ go canonRt 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) + go rt' = case RangeSet.maxView rt' of + Nothing -> (Set.empty, RangeSet.empty) + Just ((x, eid), rest) -> + let (seen, result) = go rest + in if Set.member eid seen + then (seen, result) + else (Set.insert eid seen, RangeSet.insert (x, eid) result) diff --git a/src/Algorithm/EqSat/SearchSR.hs b/src/Algorithm/EqSat/SearchSR.hs --- a/src/Algorithm/EqSat/SearchSR.hs +++ b/src/Algorithm/EqSat/SearchSR.hs @@ -15,26 +15,33 @@ import Data.SRTree import Data.SRTree.Datasets +import Data.SRTree.Eval (compileLoss) import System.Random import Control.Monad.State.Strict +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (mapConcurrently) +import Data.Maybe (catMaybes) +import Control.Exception (evaluate) +import qualified Control.DeepSeq as DeepSeq import Algorithm.EqSat.Egraph import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.AD (ADBackEnd(..)) +import Algorithm.SRTree.AD.Unboxed (setMTPopParallel) import qualified Data.IntMap as IM import qualified Data.IntSet as IntSet import qualified Data.SRTree.Random as Random import Data.Function ( on ) -import Algorithm.SRTree.Likelihoods import Algorithm.SRTree.NonlinearOpt import Control.Monad ( when, replicateM, forM, forM_ ) -import Algorithm.EqSat.Egraph -import Algorithm.SRTree.Opt +import Numeric.Optimization.NLOPT import Algorithm.EqSat.Info import Algorithm.EqSat.Build -import Data.Maybe ( fromJust ) import Data.SRTree.Random import Algorithm.EqSat.Queries import Data.List ( maximumBy ) -import qualified Data.Map.Strict as Map +import qualified Data.List as Data.List +import qualified Data.HashMap.Strict as HashMap +import qualified Data.Vector.Unboxed as V -- Environment of an e-graph with support to random generator and IO type RndEGraph a = EGraphST (StateT StdGen IO) a @@ -46,6 +53,55 @@ rnd = lift {-# INLINE rnd #-} +-- | Run an 'RndEGraph' action against a read-only egraph snapshot with the given +-- generator (for concurrent workers that do not mutate the shared egraph). +runRndEGraph :: EGraph -> StdGen -> RndEGraph a -> IO a +runRndEGraph eg g m = do + ((a, _), _) <- runStateT (runStateT m eg) g + pure a +{-# INLINE runRndEGraph #-} + +-- | Fit a batch of e-classes in parallel, then insert the results serially. +-- Semantics mirror 'updateIfNothing' (skip already-fitted) unless 'force' is +-- True. The shared 'StdGen' is split once; each worker gets its own generator, +-- so the global draw sequence differs from the serial search (acceptable). +-- While the batch runs, the MultiThread backend is switched to single-chunk so +-- cores go to the batch rather than oversubscribing the inner per-tree split. +fitBatch :: Bool + -> (Fix SRTree -> RndEGraph (Double, [Target])) + -> [EClassId] + -> RndEGraph () +fitBatch force fitFun ecs0 = do + ecs <- Prelude.mapM canonical ecs0 + jobs <- fmap catMaybes $ forM ecs $ \ec -> do + mf <- getFitness ec + if force || mf == Nothing + then do tree <- getBestExpr ec + pure (Just (ec, tree)) + else pure Nothing + case jobs of + [] -> pure () + _ -> do + nCaps <- io getNumCapabilities + g0 <- rnd get + let (seed, g1) = random g0 :: (Int, StdGen) + gs = [ mkStdGen (seed + fromIntegral i) | i <- [0 .. length jobs - 1] ] + jobsG = [ (ec, tree, g) | ((ec, tree), g) <- zip jobs gs ] + chunk k xs = [ [ xs !! j | j <- [i, i + k .. length xs - 1] ] | i <- [0 .. k - 1] ] + rnd (put g1) + eg <- get + io (setMTPopParallel True) + results <- io $ fmap concat (mapConcurrently (mapM (runJob eg fitFun)) (chunk nCaps jobsG)) + io (setMTPopParallel False) + forM_ results $ \(ec0, f, p) -> insertFitness ec0 f p + where + runJob :: EGraph -> (Fix SRTree -> RndEGraph (Double, [Target])) -> (EClassId, Fix SRTree, StdGen) -> IO (EClassId, Double, [Target]) + runJob eg fit' (ec, tree, g) = do + (f, p) <- runRndEGraph eg g (fit' tree) + f' <- evaluate (DeepSeq.force f) + p' <- evaluate (DeepSeq.force p) + pure (ec, f', p') + myCost :: SRTree Int -> Int myCost (Var _) = 1 myCost (Const _) = 1 @@ -59,34 +115,34 @@ while p arg' prog else pure arg -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 +fitnessFun :: ADBackEnd -> Bool -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> Target -> (Double, Target) +fitnessFun backend skipVal nIter loss (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig = + if isNaN val + then (-(1/0), theta) else (val, theta) where - --tree = relabelParams _tree - nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 - (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig - evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta - --tr = evalF x y mYErr - val = evalF x_val y_val mYErr_val + nParams = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL Gaussian then 1 else 0 + (theta, lossVal, _) = minimizeNLL' VAR1 backend loss mYErr nIter x y tree thetaOrig + evalF a b c = negate $ compileLoss a (buildLoss loss (fromIntegral (V.length b)) tree) b c $ if nParams == 0 then thetaOrig else theta + -- at folds=1 the validation split is the training data itself, so the + -- train loss returned by minimizeNLL' already is the val loss; skipping + -- the separate compileLoss below avoids re-evaluating every expression. + val = if skipVal then negate lossVal else evalF x_val y_val mYErr_val --{-# INLINE fitnessFun #-} -fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector) -fitnessFunRep nRep nIter distribution dataTrain dataVal tree = do - let nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 +fitnessFunRep :: ADBackEnd -> Bool -> Int -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, Target) +fitnessFunRep backend skipVal nRep nIter loss dataTrain dataVal tree = do + let nParams = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL 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 + pure $ maximumBy (\(x, _) (y, _) -> compare x y) $ Prelude.map (fitnessFun backend skipVal nIter loss dataTrain dataVal tree) thetaOrigs --{-# INLINE fitnessFunRep #-} -fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [PVector]) -fitnessMV shouldReparam nRep nIter distribution dataTrainsVals _tree = do +fitnessMV :: ADBackEnd -> Bool -> Bool -> Int -> Int -> Loss -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [Target]) +fitnessMV backend skipVal shouldReparam nRep nIter loss dataTrainsVals _tree = do let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree - response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv tree + response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep backend skipVal nRep nIter loss dt dv tree pure (minimum (Prelude.map fst response), Prelude.map snd response) @@ -95,7 +151,7 @@ -- RndEGraph utils -- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val -insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId +insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [Target])) -> RndEGraph EClassId insertExpr t fitFun = do ecId <- fromTree myCost t >>= canonical (f, p) <- fitFun t @@ -116,29 +172,28 @@ 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 + then do rndId' <- rnd $ randomFrom ecIds + rndId <- canonical rndId' + constType <- (_consts . _info) <$> getEClass rndId + case constType of + NotConst -> pure $ Just rndId + _ -> pure Nothing + else pure Nothing getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n) getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n) getBestExprWithSize n = do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical - if (not (null ec)) - then do - bestFit <- getFitness $ head ec - bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass) - pure [(head ec, bestFit)] - else pure [] + case ec of + (x:_) -> do bestFit <- getFitness x + bestP <- (_theta . _info) <$> getEClass x + pure [(x, bestFit)] + [] -> pure [] insertRndExpr maxSize rndTerm rndNonTerm = do grow <- rnd toss - n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize]) + n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. max 1 maxSize]) t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow fromTree myCost t >>= canonical @@ -152,11 +207,11 @@ --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 + mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB) + case mbec of + Just bec -> do bestFit <- (_fitness . _info) <$> getEClass bec + printExprFun 0 bec + Nothing -> pure () --paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph () paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0)) @@ -166,18 +221,17 @@ | n > maxSize = pure [] | otherwise = do ecList <- getBestExprWithSize n - if not (null ecList) - then do let (ec, mf) = head ecList - f' = fromJust mf - improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f' - ec' <- canonical ec - if improved - then do refit fitFun ec' - t <- printExprFun ix ec' - ts <- go (n+1) (ix + if improved then 1 else 0) (max f f') - pure (t:ts) - else go (n+1) (ix + if improved then 1 else 0) (max f f') - else go (n+1) ix f + case ecList of + ((ec, Just f'):_) -> do + let improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f' + ec' <- canonical ec + if improved + then do refit fitFun ec' + t <- printExprFun ix ec' + ts <- go (n+1) (ix + if improved then 1 else 0) (max f f') + pure (t:ts) + else go (n+1) (ix + if improved then 1 else 0) (max f f') + _ -> go (n+1) ix f evaluateUnevaluated fitFun = do ec <- gets (IntSet.toList . _unevaluated . _eDB) @@ -196,26 +250,26 @@ -- | check whether an e-node exists or does not exist in the e-graph doesExist, doesNotExist :: ENode -> RndEGraph Bool -doesExist en = gets ((Map.member en) . _eNodeToEClass) -doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass) +doesExist en = gets ((HashMap.member en) . _eNodeToEClass) +doesNotExist en = gets ((not . HashMap.member en) . _eNodeToEClass) -- | check whether the partial tree defined by a list of ancestors will create -- a non-existent expression when combined with a certain e-node. doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool -doesNotExistGens [] en = gets ((Map.notMember en) . _eNodeToEClass) -doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass) +doesNotExistGens [] en = gets ((not . HashMap.member en) . _eNodeToEClass) +doesNotExistGens (mGrand:grands) en = do b <- gets ((not . HashMap.member en) . _eNodeToEClass) if b then pure True else case mGrand of Nothing -> pure False - Just gf -> do ec <- gets ((Map.! en) . _eNodeToEClass) + Just gf -> do ec <- gets ((HashMap.! en) . _eNodeToEClass) en' <- canonize (gf ec) doesNotExistGens grands en' -- | check whether combining a partial tree `parent` with the e-node `en'` -- will create a new expression checkToken parent en' = do en <- canonize en' - mEc <- gets ((Map.!? en) . _eNodeToEClass) + mEc <- gets (HashMap.lookup en . _eNodeToEClass) case mEc of Nothing -> pure True Just ec -> do ec' <- canonical ec diff --git a/src/Algorithm/EqSat/SearchSRCache.hs b/src/Algorithm/EqSat/SearchSRCache.hs deleted file mode 100644 --- a/src/Algorithm/EqSat/SearchSRCache.hs +++ /dev/null @@ -1,244 +0,0 @@ ------------------------------------------------------------------------------ --- | --- Module : Algorithm.EqSat.Search --- Copyright : (c) Fabricio Olivetti 2021 - 2024 --- License : BSD3 --- Maintainer : fabricio.olivetti@gmail.com --- Stability : experimental --- Portability : --- --- Support functions for search symbolic expressions with e-graphs --- ------------------------------------------------------------------------------ - -module Algorithm.EqSat.SearchSRCache where - -import Data.SRTree -import Data.SRTree.Datasets -import System.Random -import Control.Monad.State.Strict -import Algorithm.EqSat.Egraph -import Algorithm.SRTree.Likelihoods -import qualified Data.IntMap as IM -import qualified Data.IntSet as IntSet -import qualified Data.SRTree.Random as Random -import Data.Function ( on ) -import Algorithm.SRTree.Likelihoods -import Algorithm.SRTree.NonlinearOpt -import Control.Monad ( when, replicateM, forM, forM_ ) -import Algorithm.EqSat.Egraph -import Algorithm.SRTree.Opt -import Algorithm.EqSat.Info -import Algorithm.EqSat.Build -import Data.Maybe ( fromJust ) -import Data.SRTree.Random -import Algorithm.EqSat.Queries -import Data.List ( maximumBy ) -import qualified Data.Map.Strict as Map -import Control.Monad.Identity - -import Debug.Trace - --- Environment of an e-graph with support to random generator and IO -type RndEGraph a = EGraphST (StateT StdGen (StateT [ECache] IO)) a - -io :: IO a -> RndEGraph a -io = lift . lift . lift -{-# INLINE io #-} -getCache :: StateT [ECache] IO a -> RndEGraph a -getCache = lift . lift -rnd :: StateT StdGen (StateT [ECache] IO) a -> RndEGraph a -rnd = lift -{-# INLINE rnd #-} - -myCost :: SRTree Int -> Int -myCost (Var _) = 1 -myCost (Const _) = 1 -myCost (Param _) = 1 -myCost (Bin _ l r) = 2 + l + r -myCost (Uni _ t) = 3 + t - -while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f t -while p arg prog = do if (p arg) - then do arg' <- prog arg - while p arg' prog - else pure arg - -fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> EGraph -> EClassId -> ECache -> PVector -> (Double, PVector, ECache) -fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) egraph root cache thetaOrig = - if isNaN val -- || isNaN tr - then (-(1/0), theta,cache') -- infinity - else (val, theta, cache') - where - tree = runIdentity $ getBestExpr root `evalStateT` egraph - nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 - (theta, val, _, cache') = minimizeNLLEGraph VAR1 distribution mYErr nIter x y egraph root cache thetaOrig - evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta - -- val = evalF x_val y_val mYErr_val - ---{-# INLINE fitnessFun #-} - -fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> EClassId -> ECache -> RndEGraph (Double, PVector, ECache) -fitnessFunRep nRep nIter distribution dataTrain dataVal root cache = do - egraph <- get - let nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0 - fst' (a, _, _) = a - thetaOrigs <- replicateM nRep (rnd $ randomVec nParams) - let fits = maximumBy (compare `on` fst') $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal egraph root cache) thetaOrigs - pure fits ---{-# INLINE fitnessFunRep #-} - - -fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector]) -fitnessMV shouldReparam nRep nIter distribution dataTrainsVals root = do - -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree - -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm' - caches <- getCache get - response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache - getCache $ put (Prelude.map trd response) - pure (minimum (Prelude.map fst' response), Prelude.map snd' response) - where fst' (a, _, _) = a - snd' (_, a, _) = a - trd (_, _, a) = a - -fitnessMVNoCache :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector]) -fitnessMVNoCache shouldReparam nRep nIter distribution dataTrainsVals root = do - -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree - -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm' - caches <- getCache get - response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache - pure (minimum (Prelude.map fst' response), Prelude.map snd' response) - where fst' (a, _, _) = a - snd' (_, a, _) = a - trd (_, _, a) = a - - - --- RndEGraph utils --- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val -insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId -insertExpr t fitFun = do - ecId <- fromTree myCost t >>= canonical - (f, p) <- fitFun t - insertFitness ecId f p - pure ecId - where powabs l r = Fix (Bin PowerAbs l r) - -updateIfNothing fitFun ec = do - mf <- getFitness ec - case mf of - Nothing -> do - --t <- getBestExpr ec - (f, p) <- fitFun ec - insertFitness ec f p - pure True - Just _ -> pure False - -pickRndSubTree :: RndEGraph (Maybe EClassId) -pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB) - if not (null ecIds) - then do rndId' <- rnd $ randomFrom ecIds - rndId <- canonical rndId' - constType <- gets (_consts . _info . (IM.! rndId) . _eClass) - case constType of - NotConst -> pure $ Just rndId - _ -> pure Nothing - else pure Nothing - -getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n) -getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n) - -getBestExprWithSize n = - do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical - if (not (null ec)) - then do - bestFit <- getFitness $ head ec - bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass) - pure [(head ec, bestFit)] - else pure [] - -insertRndExpr maxSize rndTerm rndNonTerm = - do grow <- rnd toss - n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize]) - t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow - fromTree myCost t >>= canonical - -refit fitFun ec = do - --t <- getBestExpr ec - (f, p) <- fitFun ec - mf <- getFitness ec - case mf of - Nothing -> insertFitness ec f p - Just f' -> when (f > f') $ insertFitness ec f p - ---printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph () -printBest fitFun printExprFun = do - bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical - bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass) - --refit fitFun bec - --io.print $ "should be " <> show bestFit - printExprFun 0 bec - ---paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph () -paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0)) - where - go :: Int -> Int -> Double -> RndEGraph [[String]] - go n ix f - | n > maxSize = pure [] - | otherwise = do - ecList <- getBestExprWithSize n - if not (null ecList) - then do let (ec, mf) = head ecList - f' = fromJust mf - improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f' - ec' <- canonical ec - if improved - then do refit fitFun ec' - t <- printExprFun ix ec' - ts <- go (n+1) (ix + if improved then 1 else 0) (max f f') - pure (t:ts) - else go (n+1) (ix + if improved then 1 else 0) (max f f') - else go (n+1) ix f - -evaluateUnevaluated fitFun = do - ec <- gets (IntSet.toList . _unevaluated . _eDB) - forM_ ec $ \c -> do - --t <- getBestExpr c - (f, p) <- fitFun c - insertFitness c f p - -evaluateRndUnevaluated fitFun = do - ec <- gets (IntSet.toList . _unevaluated . _eDB) - c <- rnd . randomFrom $ ec - --t <- getBestExpr c - (f, p) <- fitFun c - insertFitness c f p - pure c - --- | check whether an e-node exists or does not exist in the e-graph -doesExist, doesNotExist :: ENode -> RndEGraph Bool -doesExist en = gets ((Map.member en) . _eNodeToEClass) -doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass) - --- | check whether the partial tree defined by a list of ancestors will create --- a non-existent expression when combined with a certain e-node. -doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool -doesNotExistGens [] en = gets ((Map.notMember en) . _eNodeToEClass) -doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass) - if b - then pure True - else case mGrand of - Nothing -> pure False - Just gf -> do ec <- gets ((Map.! en) . _eNodeToEClass) - en' <- canonize (gf ec) - doesNotExistGens grands en' - --- | check whether combining a partial tree `parent` with the e-node `en'` --- will create a new expression -checkToken parent en' = do en <- canonize en' - mEc <- gets ((Map.!? en) . _eNodeToEClass) - case mEc of - Nothing -> pure True - Just ec -> do ec' <- canonical ec - ec'' <- canonize (parent ec') - not <$> doesExist ec'' diff --git a/src/Algorithm/EqSat/Simplify.hs b/src/Algorithm/EqSat/Simplify.hs --- a/src/Algorithm/EqSat/Simplify.hs +++ b/src/Algorithm/EqSat/Simplify.hs @@ -18,24 +18,35 @@ import Algorithm.EqSat.Egraph import Algorithm.EqSat.DB ( ClassOrVar, - Pattern (Fixed, VarPat), + Condition (Condition), + NChild (Ch, MapP, Rest), + Pattern (Fixed, Hole, NAry, VarPat), Rule (..), + Subst, + SubVal (SVMap, SVOne), getInt, ) import Control.Monad.State.Strict (evalState) -import Data.IntMap (IntMap) -import qualified Data.IntMap as IM +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IM import Data.Map (Map) import qualified Data.Map as Map import Data.SRTree -type ConstrFun = Pattern -> Map ClassOrVar ClassOrVar -> EGraph -> Bool +-- | A constraint over a match's substitution: when applied to a substitution it +-- runs in the e-graph monad and fetches e-class data through 'ClassStore', so it +-- works on a paged (out-of-core) graph whose resident cache is bounded/empty. +type ConstrFun = Pattern -> Condition -constrainOnVal :: (Consts -> Bool) -> Pattern -> Map ClassOrVar ClassOrVar -> EGraph -> Bool -constrainOnVal f (VarPat c) subst eg = - let cid = getInt $ subst Map.! Right (fromEnum c) - in f (_consts . _info $ _eClass eg IM.! cid) -constrainOnVal _ _ _ _ = False +constrainOnVal :: (Consts -> Bool) -> Pattern -> Condition +constrainOnVal f (VarPat c) = Condition $ \subst -> do + let cid = getInt $ case Map.lookup (Right (fromEnum c)) subst of + Nothing -> error $ "CONSTRAINVAL_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst) + Just (SVOne v) -> v + Just (SVMap _) -> error $ "CONSTRAINVAL_REST_AS_SINGLE var=" <> show (fromEnum c) + ec <- getEClass cid + pure (f (_consts . _info $ ec)) +constrainOnVal _ _ = Condition $ \_ -> pure False -- TODO: aux functions to avoid repeated pattern in constraint creation -- @@ -94,156 +105,149 @@ ConstVal x -> not (isNaN x || isInfinite x) _ -> True --- basic algebraic rules +-- | e-class ids bound to a rest variable +restEidsOf :: Char -> Subst -> [EClassId] +restEidsOf c subst = case Map.lookup (Right (fromEnum c)) subst of + Just (SVMap m) -> expandedList m + _ -> [] + +-- | every e-class bound to a rest variable holds a valid value +allValidRest :: Char -> Condition +allValidRest c = Condition $ \subst -> do + let eids = restEidsOf c subst + validEid eid = getEClass eid >>= \ec -> + pure $ case _consts . _info $ ec of + ConstVal x -> not (isNaN x || isInfinite x) + _ -> True + and <$> mapM validEid eids + +-- basic algebraic rules rewriteBasic :: [Rule] rewriteBasic = [ - "x" * "y" :=> "y" * "x" - , "x" + "y" :=> "y" + "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") - --, ("x" + "y") - "z" :=> "x" + ("y" - "z") -- TODO: check that I don't need that - , ("x" * "y") * "z" :=> "x" * ("y" * "z") - , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z") - , "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") -- 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" * "x" :=> "x" ** 2 - , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 - , "x" ** 2 + "x" * "y" :=> "x" * ("x" + "y") - -- , "x" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL - -- , "x" + "y" * "z" :=> "y" * ("x" * "y" ** (-1) + "z") :| isNotZero "y" -- GABRIEL + -- B7/B8/C5: factor a common term out of a sum of products, and the + -- reverse (distribute), which make x*(y+z) and x*y+x*z equivalent. + NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1']) + , Ch (NAry EMul [Ch "x", Rest '2']) + , Rest '3' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1', Rest '2']) + ]) + , Rest '3' ] + , NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1']) + ]) + , Rest '2' ] + :=> + NAry EAdd [ MapP (NAry EMul [Ch "x", Ch Hole]) '1' + , Rest '2' ] + -- C5: x*y - z*x = x*(y - z) + , NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1']) + , Ch (NAry EMul [Ch (Fixed (Const (-1))), Ch "x", Ch "z"]) + , Rest '3' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1', Ch (negate (VarPat 'z'))]) + ]) + , Rest '3' ] + -- B1: group duplicate factors into a power (x*x = x^2) + , NAry EMul [Ch "x", Ch "x"] :=> "x" ** 2 + -- C9: binomial expansion of a closed 2-ary square + , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 + -- C10: x^2 + x*y + ... = x*(x + y) + ... + , NAry EAdd [ Ch (Fixed (Bin Power (VarPat 'x') (Fixed (Const 2)))) + , Ch (NAry EMul [Ch "x", Rest '1']) + , Rest '2' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Ch "x", Rest '1']) + ]) + , Rest '2' ] ] -- rules for nonlinear functions rewritesFun :: [Rule] rewritesFun = [ - log (exp "x") :==: exp (log "x") - , log (exp "x") :=> "x" - -- , exp (log "x") :=> "x" -- :| isPositive "x" ??? exp(log(x)), x, log(exp(0)) - , log ("x" * "y") :=> log "x" + log "y" :| isConstPos "x" :| isConstPos "y" - -- , log ("x" / "y") :=> log "x" - log "y" :| isConstPos "x" :| isConstPos "y" + log (exp "x") :=> "x" + -- C11: log(x*y*z*...) = log x + log y + ... + , log (NAry EMul [Rest '1']) :=> NAry EAdd [MapP (Fixed (Uni Log Hole)) '1'] , 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" - , abs ("x" * "y") :=> abs "x" * abs "y" -- :| isConstPt "x" + -- C12: abs(x*y*z*...) = abs x * abs y * ... + , abs (NAry EMul [Rest '1']) :=> NAry EMul [MapP (Fixed (Uni Abs Hole)) '1'] , abs ("x" ** "y") :=> abs "x" ** "y" - , abs ("x" - "y") :=> abs ("y" - "x") - --, sqrt ("z" * ("x" - "y")) :=> sqrt (negate "z") * sqrt ("y" - "x") - --, sqrt ("z" * ("x" + "y")) :=> sqrt "z" * sqrt ("x" + "y") , recip (recip "x") :=> "x" :| isNotZero "x" - , ("x" * "y") ** "z" :==: ("x" ** "z") * ("y" ** "z") -- :| bothSameSign "x" "y" - , ("x" * "y") ** "z" :==: ("x" ** "z") * ("y" ** "z") -- :| isInteger "z" - --, recip "x" :==: "x" ** (-1) -- GABRIEL - --, "x" / "y" :==: "x" * "y" ** (-1) -- GABRIEL + -- C13: (x*y*z*...)^w = x^w * y^w * ... [was disabled: combinatorial blowup on (x*x)^t; the multiset matcher + matchCap bound that] + , (NAry EMul [Rest '1']) ** "z" :=> NAry EMul [MapP (Hole ** VarPat 'z') '1'] , abs "x" ** "y" :=> "x" ** "y" :| isEven "y" - , sqrt ("x" * "x") :=> abs "x" + -- C14: sqrt(x*x) = abs x + , sqrt (NAry EMul [Ch "x", Ch "x"]) :=> abs "x" ] -- Rules that reduces redundant parameters constReduction :: [Rule] constReduction = [ - 0 + "x" :=> "x" - -- , "x" - 0 :=> "x" - --, 1 * "x" :=> "x" - -- , 0 / "x" :=> 0 :| isNotZero "x" - --, "x" - "x" :=> 0 :| isNotParam "x" - --, "x" / "x" :=> 1 :| isNotZero "x" :| isNotParam "x" + -- B3: 0 + rest = rest + NAry EAdd [Ch (Fixed (Const 0)), Rest '1'] :=> NAry EAdd [Rest '1'] , "x" ** 1 :=> "x" , powabs "x" 1 :=> abs "x" - -- , "x" * (1 / "x") :=> 1 :| isNotParam "x" :| isNotZero "x" - -- , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "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" + -- B9: x^y * x^z = x^(y+z) + , NAry EMul [Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'z')))] + :==: + Fixed (Bin Power (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')])) + :| isPositive "x" + -- B10: |x|^y * |x|^z = |x|^(y+z) (fixed: target used "y+x" instead of "y+z") + , NAry EMul [Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'z')))] + :=> + Fixed (Bin PowerAbs (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')])) + -- B11: (x^y)^z = x^(y*z) + , Fixed (Bin Power (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))) (VarPat 'z')) + :==: + Fixed (Bin Power (VarPat 'x') (NAry EMul [Ch (VarPat 'y'), Ch (VarPat '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" ** "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" :=> 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" -- <==> + , "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" + -- B4: 0 * rest = 0 (provided every factor is valid) + , NAry EMul [Ch (Fixed (Const 0)), Rest '1'] :=> 0 :| allValidRest '1' , 0 ** "x" :=> 0 :| isPositive "x" , powabs 0 "x" :=> 0 - , 0 - "x" :=> negate "x" - , "x" + negate "y" :==: "x" - "y" + -- n-ary cancellation: x + y - x = y + , NAry EAdd [ Ch "a" + , Ch (NAry EMul [ Ch (Fixed (Const (-1.0))), Ch "a" ]) + , Rest 'r' ] + :=> NAry EAdd [Rest 'r'] + -- combining like terms: x + x = 2*x + , NAry EAdd [ Ch "a", Ch "a", Rest 'r' ] + :=> NAry EAdd [ Ch (2 * "a"), Rest 'r' ] ] 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" - ] +rewritesSimple = rewriteBasic <> constReduction <> rewritesFun powabs l r = Fixed (Bin PowerAbs l r) -- | default cost function for simplification @@ -267,14 +271,14 @@ rewritesParams :: [Rule] rewritesParams = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithParam --- | simplify using the default parameters +-- | simplify using the default parameters simplifyEqSatDefault :: Fix SRTree -> Fix SRTree -simplifyEqSatDefault t = eqSat t rewrites myCost 30 `evalState` emptyGraph +simplifyEqSatDefault t = eqSat t rewrites myCost 30 `evalState` emptyGraphNoTrack -- | simplifies with custom parameters simplifyEqSat :: [Rule] -> CostFun -> Int -> Fix SRTree -> Fix SRTree simplifyEqSat rwrts costFun it t = eqSat t rwrts costFun it `evalState` emptyGraph -- | apply a single step of merge-only using default rules -applyMergeOnlyDftl :: Monad m => CostFun -> EGraphST m () +applyMergeOnlyDftl :: ClassStore m => CostFun -> EGraphST m () applyMergeOnlyDftl costFun = applySingleMergeOnlyEqSat costFun rewrites diff --git a/src/Algorithm/EqSat/Store.hs b/src/Algorithm/EqSat/Store.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Store.hs @@ -0,0 +1,243 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveDataTypeable #-} + +module Algorithm.EqSat.Store + ( GraphRows(..) + , EClassRow(..) + , exportEGraph + , importEGraph + , mergeEGraph + , rebuildDBs + ) where + +import Control.Lens ( over ) +import Control.Monad ( forM, forM_, foldM ) +import Control.Monad.Identity ( Identity, runIdentity ) +import Control.Monad.State.Strict ( StateT, execStateT, modify', gets ) +import GHC.Generics ( Generic ) +import GHC.Stack ( HasCallStack ) + +import qualified Data.HashMap.Strict as HashMap +import Data.HashMap.Strict ( HashMap ) +import qualified Data.HashSet as Set +import qualified Data.IntMap.Strict as IntMap +import Data.IntMap.Strict ( IntMap ) +import qualified Data.IntSet as IntSet +import qualified Data.Set as RangeSet +import Data.List ( sortOn ) + +import Data.SRTree +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build + +-- | Row representation of the core (structural) state of an e-graph, +-- normalized for external storage (e.g. a relational DB). +data GraphRows = GraphRows + { _grCanonical :: IntMap EClassId -- ^ eid -> canonical representative (self-loop for roots) + , _grENodeToEClass :: HashMap ENode EClassId -- ^ canonical e-node -> its e-class + , _grEClasses :: IntMap EClassRow -- ^ canonical e-class id -> data row + , _grNextId :: Int -- ^ next free e-class id + , _grTrackDBs :: Bool -- ^ whether range DBs are maintained + } deriving (Show, Eq, Generic) + +-- | Per-e-class data row. +data EClassRow = EClassRow + { _rcNodes :: Set.HashSet ENode + , _rcParents :: Set.HashSet (EClassId, ENode) + , _rcHeight :: Int + , _rcInfo :: EClassData + } deriving (Show, Eq, Generic) + +-- | Export the core structural state of an e-graph into a normalised row format. +exportEGraph :: EGraph -> GraphRows +exportEGraph eg = GraphRows + { _grCanonical = _canonicalMap eg + , _grENodeToEClass = _eNodeToEClass eg + , _grEClasses = IntMap.map toRow (_eClass eg) + , _grNextId = _nextId (_eDB eg) + , _grTrackDBs = _trackDBs (_eDB eg) + } + where + toRow ec = EClassRow (_eNodes ec) (_parents ec) (_height ec) (_info ec) + +-- | Reconstruct an e-graph from normalised rows, rebuilding all derived indexes. +-- +-- Real e-graphs may carry stale @_eNodeToEClass@ entries left behind by +-- merges (a node pointing at a class whose canonical representative is +-- another class). Such entries are canonicalized at import: node -> class +-- values are routed through the canonical map and any non-root class rows +-- are dropped. Parent pointers are recomputed from the canonicalized node +-- map so they never reference dead classes. +importEGraph :: GraphRows -> Either String EGraph +importEGraph rows + | not (validate rows) = Left (validationMsg rows) + | otherwise = Right (runIdentity $ execStateT rebuildDBs (buildCore (canonicalize rows))) + +-- | Normalize stale rows: route node->class values through the canonical map +-- and drop non-root class rows. +-- +-- Parent pointers come from the stored @_rcParents@ when a class has any +-- (e.g. after a storage-layer round-trip through the @parent@ table); parent +-- class ids are routed through the canonical map so they never reference dead +-- classes. Classes without stored parents (legacy rows, hand-built rows) fall +-- back to recomputing parents from the canonicalized node map. +canonicalize :: GraphRows -> GraphRows +canonicalize rows = + let canon = _grCanonical rows + rep eid = IntMap.findWithDefault eid eid canon + nodeMap' = HashMap.map rep (_grENodeToEClass rows) + classes' = IntMap.filterWithKey + (\eid _ -> IntMap.lookup eid canon == Just eid) + (_grEClasses rows) + parents' = IntMap.fromListWith Set.union + [ (c, Set.singleton (eid, en)) + | (en, eid) <- HashMap.toList nodeMap' + , c <- eChildren en ] + stored' = IntMap.mapWithKey + (\_ r -> Set.map (\(pEid, pEn) -> (rep pEid, pEn)) (_rcParents r)) + classes' + fixRow eid r = + let stored = IntMap.findWithDefault Set.empty eid stored' + in r { _rcParents = if Set.null stored + then IntMap.findWithDefault Set.empty eid parents' + else stored } + in rows { _grENodeToEClass = nodeMap' + , _grEClasses = IntMap.mapWithKey fixRow classes' } + +buildCore :: GraphRows -> EGraph +buildCore rows = EGraph + { _canonicalMap = _grCanonical rows + , _eNodeToEClass = _grENodeToEClass rows + , _eClass = IntMap.mapWithKey mkEClass (_grEClasses rows) + , _eDB = (emptyDB){ _nextId = _grNextId rows, _trackDBs = _grTrackDBs rows } + , _classStore = Nothing + } + where + mkEClass eid r = EClass eid (_rcNodes r) (_rcParents r) (_rcHeight r) (_rcInfo r) + +rebuildDBs :: EGraphST Identity () +rebuildDBs = do + -- Rebuild the pattern database from the canonical e-node -> class mapping + nodes <- gets _eNodeToEClass + forM_ (HashMap.toList nodes) $ \(en, eid) -> addToDB en eid + + -- Rebuild range/size indexes from class info + classes <- gets _eClass + forM_ (IntMap.toList classes) $ \(eid, ec) -> do + let info = _info ec + sz = _size info + fit = _fitness info + dl = _dl info + modify' $ over (eDB . sizeDB) (IntMap.insertWith IntSet.union sz (IntSet.singleton eid)) + case fit of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.insert eid) + Just fn -> modify' $ over (eDB . fitRangeDB) (insertRange eid fn) + . over (eDB . sizeFitDB) (IntMap.insertWith RangeSet.union sz (RangeSet.singleton (fn, eid))) + case dl of + Nothing -> pure () + Just dn -> modify' $ over (eDB . dlRangeDB) (insertRange eid dn) + . over (eDB . sizeDLDB) (IntMap.insertWith RangeSet.union sz (RangeSet.singleton (dn, eid))) + +-- | Validate that the exported rows form a consistent graph. +-- +-- All referenced ids must be present in the canonical map. Node -> class +-- values and class rows may reference classes that are not their own +-- canonical representative (stale entries left behind by merges); those are +-- repaired by 'canonicalize' during import. +validate :: GraphRows -> Bool +validate rows = + let canon = _grCanonical rows + classes = _grEClasses rows + nodeIds = HashMap.keys (_grENodeToEClass rows) + extraIds = IntMap.keys classes + ++ HashMap.elems (_grENodeToEClass rows) + ++ concatMap eChildren nodeIds + inCanon = all (`IntMap.member` canon) extraIds + nextOk = _grNextId rows >= 0 + in inCanon && nextOk + +validationMsg :: GraphRows -> String +validationMsg rows + | not inCanon = "some e-node/e-class id is not present in the canonical map" + | not nextOk = "next id is negative" + | otherwise = "invalid GraphRows" + where + canon = _grCanonical rows + classes = _grEClasses rows + nodeIds = HashMap.keys (_grENodeToEClass rows) + extraIds = IntMap.keys classes + ++ HashMap.elems (_grENodeToEClass rows) + ++ concatMap eChildren nodeIds + inCanon = all (`IntMap.member` canon) extraIds + nextOk = _grNextId rows >= 0 + +-- | Return canonical e-class ids ordered children-before-parents (ascending height). +classOrder :: GraphRows -> Either String [EClassId] +classOrder rows = + Right $ map fst $ sortOn (_rcHeight . snd) $ IntMap.toAscList (_grEClasses rows) + +-- | Remap a B-e-graph's e-node into A's id-space using the correspondence map. +remapNode + :: GraphRows -- ^ rows of graph B (source) + -> IntMap EClassId -- ^ corr: B canonical id -> A id + -> ENode + -> Either String ENode +remapNode rowsB corr = go + where + canonB :: EClassId -> EClassId + canonB cid = IntMap.findWithDefault cid cid (_grCanonical rowsB) + + toA :: EClassId -> Either String EClassId + toA cid = + case IntMap.lookup (canonB cid) corr of + Just eidA -> Right eidA + Nothing -> Left ("child " <> show cid <> " of graph B not yet merged") + + go (EVar ix) = Right (EVar ix) + go (EParam ix) = Right (EParam ix) + go (EConst x) = Right (EConst x) + go (EUni f t) = EUni f <$> toA t + go (EBin op l r) = EBin op <$> toA l <*> toA r + go (ENAry op m) = do + m' <- foldM step IntMap.empty (IntMap.toList m) + Right (ENAry op m') + where + step acc (cid, n) = do + cidA <- toA cid + pure (IntMap.insertWith (+) cidA n acc) + +-- | Merge class ids by unioning their e-classes under the given cost function. +mergeClass :: HasCallStack => CostFun -> EClassId -> EClassId -> EGraphST Identity EClassId +mergeClass costFun x y = + if x == y then pure x else merge costFun x y + +-- | Structurally merge graph @b@ into a copy of graph @a@. +-- +-- The e-nodes of @b@ are canonicalized under @a@'s id space, deduplicated +-- against @a@'s existing content, and equivalent classes are unioned. Cost and +-- best of newly introduced content are computed with @costFun@ (i.e. merging +-- adopts @a@'s cost function). Dataset-specific values (fitness/DL/theta) are +-- NOT transferred: they are per-dataset data managed by the storage layer. +mergeEGraph :: HasCallStack => CostFun -> EGraph -> EGraph -> Either String EGraph +mergeEGraph costFun a b = + let rowsB = exportEGraph b + in case classOrder rowsB of + Left err -> Left err + Right order -> Right (runIdentity $ execStateT (step IntMap.empty order) a) + where + step :: IntMap EClassId -> [EClassId] -> EGraphST Identity () + step _ [] = rebuild costFun + step corr (bCanon : rest) = do + let ec = _grEClasses rowsB IntMap.! bCanon + resolved <- forM (Set.toList (_rcNodes ec)) $ \en -> + case remapNode rowsB corr en of + Left err -> pure (Left err) + Right enA -> Right <$> add costFun enA + case sequence resolved of + Left err -> error ("mergeEGraph: " <> err) -- pre-validated + Right [] -> step corr rest + Right (x : xs) -> do + rep <- foldM (mergeClass costFun) x xs + step (IntMap.insert bCanon rep corr) rest + rowsB = exportEGraph b diff --git a/src/Algorithm/Massiv/Utils.hs b/src/Algorithm/Massiv/Utils.hs deleted file mode 100644 --- a/src/Algorithm/Massiv/Utils.hs +++ /dev/null @@ -1,278 +0,0 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} -module Algorithm.Massiv.Utils where - -import Data.Massiv.Array hiding ( forM_, unzip, map, init, zipWith, zip, tail, replicate, take ) -import qualified Data.Massiv.Array as A -import qualified Data.Massiv.Array.Unsafe as UMA -import qualified Data.Massiv.Array.Mutable as MMA -import Control.Monad -import Data.Vector.Storable ((//)) -import System.IO.Unsafe - --- taken from https://hackage.haskell.org/package/cubicspline-0.1.2 -import Control.Arrow -import Data.List(unfoldr) - -import Data.SRTree.Eval - -type MMassArray m = MMA.MArray (PrimState m) S Ix2 Double - -getRows :: SRMatrix -> Array B Ix1 PVector -getRows = computeAs B . outerSlices -{-# INLINE getRows #-} -getCols :: SRMatrix -> Array B Ix1 PVector -getCols = computeAs B . A.map (computeAs S) . innerSlices -{-# INLINE getCols #-} - -appendRow :: MonadThrow m => SRMatrix -> PVector -> m SRMatrix -appendRow xs v = computeAs S <$> (stackOuterSlicesM . toList . computeAs B $ snoc (outerSlices xs) v) -{-# INLINE appendRow #-} - -appendCol :: MonadThrow m => SRMatrix -> PVector -> m SRMatrix -appendCol xs v = computeAs S <$> (stackInnerSlicesM . toList . computeAs B $ snoc (A.map (computeAs S) $ innerSlices xs) v) -{-# INLINE appendCol #-} - -updateS :: Array S Ix1 Double -> [(Int, Double)] -> Array S Ix1 Double -updateS vec new = fromStorableVector compMode $ toStorableVector vec // new - -linSpace :: Int -> (Double, Double) -> [Double] -linSpace num (lo, hi) = Prelude.take num $ iterate (\x -> x + step) lo - where - step = (hi - lo) / (fromIntegral num - 1) -{-# INLINE linSpace #-} - -outer :: (MonadThrow m) - => PVector - -> PVector - -> m SRMatrix -outer arr1 arr2 - | isEmpty arr1 || isEmpty arr2 = pure $ setComp comp empty - | otherwise = - pure $ makeArray comp (Sz2 m1 m2) $ \(i :. j) -> - UMA.unsafeIndex arr1 i * UMA.unsafeIndex arr2 j - where - comp = getComp arr1 <> getComp arr2 - Sz1 m1 = size arr1 - Sz1 m2 = size arr2 -{-# INLINE outer #-} - -det :: SRMatrix -> Double -det mtx - | m==0 || n==0 = 1 - | otherwise = (^2) $ Prelude.product [l ! (i :. i) | i <- [0 .. m-1]] - where - Sz (m :. n) = size mtx - (l, _) = unsafePerformIO (lu mtx) - -detChol :: SRMatrix -> Double -detChol mtx - | m==0 || n==0 = 1 - | otherwise = (^2) $ Prelude.product [cho ! (i :. i) | i <- [0 .. m-1]] - where - Sz (m :. n) = size mtx - cho = unsafePerformIO (cholesky mtx) -{-# INLINE det #-} - -rangedLinearDotProd :: PrimMonad m => Int -> Int -> Int -> MMassArray m -> m Double -rangedLinearDotProd r1 r2 len arr = go 0 0 - where - go !acc k - | k < len = do x <- UMA.unsafeLinearRead arr (r1 + k) - y <- UMA.unsafeLinearRead arr (r2 + k) - go (acc + x*y) (k + 1) - | otherwise = pure acc -{-# INLINE rangedLinearDotProd #-} - -data NegDef = NegDef - deriving Show - -instance Exception NegDef - -cholesky :: (PrimMonad m, MonadThrow m, MonadIO m) - => SRMatrix - -> m SRMatrix -cholesky arr - | m /= n = throwM $ SizeMismatchException (size arr) (size arr) - | isEmpty arr = pure $ setComp comp empty - | otherwise = MMA.createArrayS_ (size arr) create - where - comp = getComp arr - (Sz2 m n) = size arr - create l = Prelude.mapM_ (update l) [i :. j | i <- [0..m-1], j <- [0..m-1]] - - update l ix@(i :. j) - | i < j = UMA.unsafeWrite l ix 0 - | otherwise = do let cur = UMA.unsafeIndex arr ix - rowI = i*m - rowJ = j*m - xjj <- UMA.unsafeLinearRead l (rowJ + j) - tot <- rangedLinearDotProd rowI rowJ j l - let delta = cur - tot - if i == j - then if delta <= 0 - then throwM NegDef -- SizeMismatchException (size arr) (size arr) -- look at a better exception - else UMA.unsafeLinearWrite l (rowI + j) (sqrt delta) - else UMA.unsafeLinearWrite l (rowI + j) (delta / xjj) -{-# INLINE cholesky #-} - -invChol :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> m SRMatrix -invChol arr = do l <- cholesky arr -- lower diag - mtx <- thawS l - forM_ [0 .. m-1] $ \i -> do - lII <- UMA.unsafeRead mtx (i :. i) - UMA.unsafeWrite mtx (i :. i) (1 / lII) - forM_ [0 .. i-1] $ \j -> do - tot <- rangedLinearDotProd (i*m + j) (j*m + j) (i-j) mtx - UMA.unsafeWrite mtx (j :. i) ((-tot)/lII) - UMA.unsafeWrite mtx (i :. j) 0 - mm <- newMArray (Sz2 m m) 0 - forM_ [0 .. m-1] $ \i -> do - dii <- rangedLinearDotProd (i*m + i) (i*m + i) (m - i) mtx - UMA.unsafeWrite mm (i :. i) dii - forM_ [i+1 .. m-1] $ \j -> do - dij <- rangedLinearDotProd (i*m + j) (j*m + j) (m - j) mtx - UMA.unsafeWrite mm (i :. j) dij - UMA.unsafeWrite mm (j :. i) dij - freezeS mm - - where - Sz2 m _ = size arr -{-# INLINE invChol #-} - --- LU decomposition and solver taken from https://hackage.haskell.org/package/linear-1.23/docs/src/Linear.Matrix.html -lu :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> m (SRMatrix, SRMatrix) -lu mtx = do - let (Sz2 m n) = size mtx - u <- thawS $ computeAs S $ identityMatrix (Sz m) - l <- thawS $ A.replicate compMode (Sz2 m n) 0 - - let buildLVal !i !j = do - let go !k !s - | k == j = pure s - | otherwise = do lik <- UMA.unsafeRead l (i :. k) - ukj <- UMA.unsafeRead u (k :. j) - go (k+1) ( s + (lik * ukj) ) - s' <- go 0 0 - UMA.unsafeWrite l (i :. j) ((mtx ! (i :. j)) - s') - -- pure l - buildL !i !j - = when (i /= n) $ do buildLVal i j - buildL (i+1) j - buildUVal !i !j = do - let go !k !s - | k == j = pure s - | otherwise = do ljk <- UMA.unsafeRead l (j :. k) - uki <- UMA.unsafeRead u (k :. i) - go (k+1) (s + ljk * uki) - - s' <- go 0 0 - ljj <- UMA.unsafeRead l (j :. j) - UMA.unsafeWrite u (j :. i) (((mtx ! (j :. i)) - s') / (ljj)) - -- pure u - - buildU !i !j - = when (i /= n) $ do buildUVal i j - buildU (i+1) j - buildLU !j - = when (j /= n) $ - do buildL j j - buildU j j - buildLU (j+1) - buildLU 0 - finalL <- freezeS l - finalU <- freezeS u - pure (finalL, finalU) - -forwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector -forwardSub a b = do - let (Sz m) = size b - x <- thawS $ A.replicate compMode (Sz1 m) 0 - let coeff !i !j !s - | j == i = pure s - | otherwise = do let aij = a ! (i :. j) - xj <- UMA.unsafeRead x j - coeff i (j+1) (s + aij * xj) - go !i = when (i/= m) $ - do let bi = b ! i - aii = a ! (i :. i) - c <- coeff i 0 0 - UMA.unsafeWrite x i ((bi - c)/aii) - go (i+1) - go 0 - freezeS x - -backwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector -backwardSub a b = do - let (Sz m) = size b - x <- thawS $ A.replicate compMode (Sz1 m) 0 - let coeff !i !j !s - | j == m = pure s - | otherwise = do let aij = a ! (i :. j) - xj <- UMA.unsafeRead x j - coeff i (j+1) (s + aij * xj) - go !i = when (i >= 0) $ - do let bi = b ! i - aii = a ! (i :. i) - c <- coeff i (i+1) 0 - UMA.unsafeWrite x i ((bi - c)/aii) - go (i-1) - go (m-1) - freezeS x - -luSolve :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector -luSolve a b = do (l, u) <- lu a - forwardSub l b >>= backwardSub u - -type PolyCos = (Double, Double, Double) - --- | Given a list of (x,y) co-ordinates, produces a list of coefficients to cubic equations, with knots at each of the initially provided x co-ordinates. Natural cubic spline interpololation is used. See: . -cubicSplineCoefficients :: [(Double, Double)] -> [PolyCos] -cubicSplineCoefficients xs = Prelude.zip3 x y z' - where - x = map fst xs - y = map snd xs - xdiff = zipWith (-) (tail x) x - xdiff' = fromList compMode xdiff :: Vector S Double - dydx :: Vector S Double - dydx = fromList compMode $ Prelude.zipWith3 (\y0 y1 xd -> (y0-y1)/xd) (tail y) y xdiff - n = length x - - w :: [Double] - w = 0 : nextW 1 w - where - nextW ix (wi : t) - | ix == n-1 = [] - | otherwise = let m = (xdiff' ! (ix-1)) * (2 - wi) + 2 * (xdiff' ! ix) - wn = (xdiff' ! ix) / m - in wn : nextW (ix+1) t - z :: [Double] - z = 0 : nextZ 1 z - where - nextZ ix (zi : t) - | ix == n-1 = [0] - | otherwise = let m = (xdiff' ! (ix-1)) * (2 - (w !! (ix-1))) + 2 * (xdiff' ! ix) - zn = (6*((dydx ! ix) - (dydx ! (ix-1))) - (xdiff' ! (ix-1)) * zi) / m - in zn : nextZ (ix+1) t - - z' :: [Double] - z' = Prelude.reverse $ 0 : [z !! i - w !! i * z !! (i+1) | i <- [n-2,n-3 .. 0]] - -chunkBy :: Int -> [t] -> [[t]] -chunkBy n = unfoldr go - where go [] = Nothing - go x = Just $ splitAt n x - -genSplineFun :: [(Double, Double)] -> Double -> Double -genSplineFun pts x = go xs $ zip coefs (tail coefs) - where - xs = map fst pts - coefs = cubicSplineCoefficients pts - evalAt (a1,b1,c1) (a2,b2,c2) y = let hi1 = a2 - a1 - in c1/(6*hi1)*(a2-y)^3 + c2/(6*hi1)*(y-a1)^3 + (b2/hi1 - c2*hi1/6)*(y-a1) + (b1/hi1 - c1*hi1/6)*(a2-y) - - go [x1,x2] [(c1,c2)] = evalAt c1 c2 x - go (x1:x2:xs) ((c1,c2):cs) - | x < x1 = evalAt c1 c2 x - | x >= x1 && x <= x2 = evalAt c1 c2 x - | otherwise = go (x2:xs) cs diff --git a/src/Algorithm/SRTree/AD.hs b/src/Algorithm/SRTree/AD.hs --- a/src/Algorithm/SRTree/AD.hs +++ b/src/Algorithm/SRTree/AD.hs @@ -1,12 +1,3 @@ -{-# language FlexibleInstances, DeriveFunctor #-} -{-# language ScopedTypeVariables #-} -{-# language RankNTypes #-} -{-# language ViewPatterns #-} -{-# language FlexibleContexts #-} -{-# language BangPatterns #-} -{-# language TypeApplications #-} -{-# language MultiWayIf #-} - ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.AD @@ -21,539 +12,21 @@ ----------------------------------------------------------------------------- module Algorithm.SRTree.AD - ( reverseModeArr - , reverseModeEGraph - , reverseModeGraph - , forwardModeUniqueJac - , evalCache + ( compileFunAndGrad + , ADBackEnd(..) ) where -import Control.Monad (forM_, foldM, when) -import Control.Monad.ST ( runST ) -import Data.Bifunctor (bimap, first, second) -import qualified Data.DList as DL -import Data.Massiv.Array hiding (forM_, map, replicate, zipWith) -import qualified Data.Massiv.Array as M -import qualified Data.Massiv.Array.Unsafe as UMA -import Data.Massiv.Core.Operations (unsafeLiftArray) -import Data.SRTree.Derivative ( derivative ) -import Data.SRTree.Eval - ( SRVector, evalFun, evalOp, SRMatrix, PVector, replicateAs ) -import Data.SRTree.Internal -import Data.SRTree.Print (showExpr) -import Data.SRTree.Recursion ( cataM, cata, accu ) -import qualified Data.Vector as V -import Debug.Trace (trace, traceShow) -import GHC.IO (unsafePerformIO) -import qualified Data.IntMap.Strict as IntMap -import Data.List ( foldl' ) -import qualified Data.Vector.Storable as VS -import Control.Scheduler -import Data.Maybe ( fromJust, isJust ) -import Algorithm.EqSat.Egraph - -import Control.Monad.State.Strict -import Control.Monad.Identity - ---import UnliftIO.Async - -import qualified Data.Map.Strict as Map - -evalCache :: SRMatrix -> EGraph -> ECache -> EClassId -> VS.Vector Double -> ECache -evalCache xss egraph cache root' theta = cache' - where - (Sz2 _ m') = M.size xss - m = Sz1 m' - root = canon root' - p = VS.length theta - comp = M.getComp xss - one :: Array S Ix1 Double - one = M.replicate comp m 1 - - canon rt = case _canonicalMap egraph IntMap.!? rt of - Nothing -> error "wrong canon" - Just rt' -> if rt == rt' then rt else canon rt' - - getNode rt' = let rt = canon rt' - cls = _eClass egraph IntMap.! rt - in (_best . _info) cls - - getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph - in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n' - - ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty) - where - evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - evalCached rt = insertKey rt - - insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - insertKey key' = do - let key = canon key' - isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst) - isCachedLocal <- gets ((key `IntMap.member`) . snd . fst) - when (not isCachedLocal && not isCachedGlobal) $ do - let node = getNode key - (ev, toLocal) <- evalKey node - modify' (insKey node ev toLocal) - getVal key - - evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - evalKey (Var ix) = pure $ (M.computeAs S $ xss evalKey n - Const v -> evalKey n - Param ix -> evalKey n - _ -> getFromCache rt - getFromCache rt = do - global <- gets ((IntMap.!? rt) . fst . fst) - local <- gets ((IntMap.!? rt) . snd . fst) - if | isJust global -> pure (fromJust global, False) - | isJust local -> pure (fromJust local, True) - | otherwise -> insertKey rt - --- reverse mode applied directly on an e-graph. Supports caching. --- assumes root points to the loss function, so for an expression --- f(x) and the loss (y - (f(x))^2), root will point to "^" -reverseModeEGraph :: SRMatrix -> PVector -> Maybe PVector -> EGraph -> ECache -> EClassId -> VS.Vector Double -> (Array D Ix1 Double, VS.Vector Double) -reverseModeEGraph xss ys mYErr egraph cache root' theta = - (delay $ rootVal - , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]] - ) - where - rootVal = extractCache (cache'' IntMap.!? root', localcache' IntMap.!? root') - root = canon root' - yErr = fromJust mYErr - m = M.size ys - p = VS.length theta - comp = M.getComp xss - one :: Array S Ix1 Double - one = M.replicate comp m 1 - - canon rt = case _canonicalMap egraph IntMap.!? rt of - Nothing -> error "wrong canon" - Just rt' -> if rt == rt' then rt else canon rt' - - getNode rt' = let rt = canon rt' - cls = _eClass egraph IntMap.! rt - in (_best . _info) cls - - getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph - in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n' - - ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty) - where - evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - evalCached rt = insertKey rt - - insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - insertKey key' = do - let key = canon key' - isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst) - isCachedLocal <- gets ((key `IntMap.member`) . snd . fst) - when (not isCachedLocal && not isCachedGlobal) $ do - let node = getNode key - (ev, toLocal) <- evalKey node - modify' (insKey node ev toLocal) - getVal key - - evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) - evalKey (Var ix) = pure $ if | ix == -1 -> (ys, False) - | ix == -2 -> (yErr, False) - | otherwise -> (M.computeAs S $ xss evalKey n - Const v -> evalKey n - Param ix -> evalKey n - _ -> getFromCache rt - getFromCache rt = do - global <- gets ((IntMap.!? rt) . fst . fst) - local <- gets ((IntMap.!? rt) . snd . fst) - if | isJust global -> pure (fromJust global, False) - | isJust local -> pure (fromJust local, True) - | otherwise -> insertKey rt - - extractCache (Nothing, Nothing) = error "no root info" - extractCache (Just r, _) = r - extractCache (_, Just r) = r - - ((cache'', localcache'), cachedGrad) = calcGrad root one `execState` ((cache', localcache), Map.empty) - - calcGrad :: Int -> Array S Ix1 Double -> State ((IntMap.IntMap (Array S Ix1 Double), IntMap.IntMap (Array S Ix1 Double)), Map.Map (SRTree Int) (Array S Ix1 Double)) () - calcGrad rt v = do let node = getNode rt - case node of - Bin op l r -> do xl <- fst <$> getVal l - xr <- fst <$> getVal r - (dl, dr) <- diff op v xl xr l r - calcGrad l dl - calcGrad r dr - Uni f t -> do x <- fst <$> getVal t - calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x)) - Param ix -> modify' (insertGrad v (Param ix)) - _ -> pure () - where - insertGrad v k ((a, b), g) = ((a, b), Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g) - - --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double) - diff Add dx fx gy l r = pure (dx, dx) - diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx) - diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx) - diff Div dx fx gy l r = do - let k = getId (Bin Div l r) - v <- fst <$> getVal k - pure (M.computeAs S $ M.zipWith (/) dx gy - , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy)) - diff Power dx fx gy l r = do - let k = getId (Bin Power l r) - v <- fst <$> getVal k - pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v - , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v) - - diff PowerAbs dx fx gy l r = do - let k = getId (Bin PowerAbs l r) - v <- fst <$> getVal k - let v2 = M.map abs fx - v3 = M.computeAs S $ M.zipWith (*) fx gy - pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2 - , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v) - - diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx - dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl - in pure (M.computeAs S $ dxl, M.computeAs S $ dxy) - - fixNaN x = if isNaN x then 0 else x - - -reverseModeGraph :: SRMatrix -> PVector -> Maybe PVector -> VS.Vector Double -> Fix SRTree -> (Array D Ix1 Double, VS.Vector Double) -reverseModeGraph xss ys mYErr theta tree = (delay $ cachedVal' IntMap.! root - , 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) - - 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) - - graph (a, _, _, _) = a - insKey key ev (a, b, c, d) = (Map.insert key d a, IntMap.insert d key b, IntMap.insert d ev c, d+1) - -- get the values from the cache - getVal key (a, b, c, d) = c IntMap.! key - -- maps the the struct to an integer key - getKey key (a, b, c, d) = a Map.! key - - -- this tells the order in which we traverse the tree - leftToRight (Uni f mt) = Uni f <$> mt; - leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr - leftToRight (Var ix) = pure (Var ix) - leftToRight (Param ix) = pure (Param ix) - leftToRight (Const c) = pure (Const c) - - evalKey (Var ix) = pure $ if ix == -1 - then ys - else if ix == -2 - then yErr - else M.computeAs S $ xss gets (getVal t) - evalKey (Bin op l r) = M.computeAs S <$> (M.zipWith (evalOp op) <$> gets (getVal l) <*> gets (getVal r)) - - 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) - - 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) - - diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx - dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl - in pure (M.computeAs S $ dxl, M.computeAs S $ dxy) - - fixNaN x = if isNaN x then 0 else x - - 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 with the tree encoded as an array, that is even faster. -reverseModeArr :: SRMatrix - -> PVector - -> 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) -reverseModeArr xss ys mYErr theta t j2ix = - unsafePerformIO $ do - fwd <- M.newMArray (Sz2 n m) 0 - partial <- M.newMArray (Sz2 n m) 0 - jacob <- M.newMArray (Sz p) 0 - 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 - v <- UMA.unsafeFreeze (getComp xss) val - pure (delay v, j) - - where - (Sz2 m _) = M.size xss - p = VS.length theta - n = length t - toLin i j = i*m + j - yErr = fromJust mYErr - eps = 1e-8 - - 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 - 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 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 - 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) - 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) - 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 #-} - - 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 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 (lj :. i) dxl - UMA.unsafeWrite partial (rj :. i) dxr - makeRev _ = pure () - {-# 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)| - - -- |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 - - fixNaN x | isNaN x = 0 - | otherwise = x - - 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) - - {-# 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 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 () - {-# INLINE combine #-} +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as V +import Data.SRTree +import Algorithm.SRTree.AD.Unboxed --- | 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] -forwardModeUniqueJac xss theta = snd . second (map (M.computeAs M.S) . DL.toList) . cata alg - where - (Sz n) = M.size theta - one = replicateAs xss 1 +data ADBackEnd = SingleThread | MultiThread deriving (Read, Show) - alg (Var ix) = (xss [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double) -> Fix SRTree -> V.Vector Double -> (Double, V.Vector Double) +compileFunAndGrad SingleThread xss ys mYerr tree = + let ct = compileTree xss ys mYerr tree + in \theta -> evalGradVec ct theta +compileFunAndGrad MultiThread xss ys mYerr tree = + let cts = compileTreeMulti xss ys mYerr tree + in \theta -> evalGradMulti cts theta diff --git a/src/Algorithm/SRTree/AD/CompiledAD.hs b/src/Algorithm/SRTree/AD/CompiledAD.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD/CompiledAD.hs @@ -0,0 +1,39 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.AD.CompiledAD +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Automatic Differentiation for Expression trees +-- +----------------------------------------------------------------------------- + +module Algorithm.SRTree.AD.CompiledAD + ( CompiledTree(..) + ) where + +import Data.SRTree.Internal +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector as VB + +-- --------------------------------------------------------------------- +-- Public entry point -- same signature/behaviour as before. +-- --------------------------------------------------------------------- +data CompiledTree = CompiledTree + { ctNodes :: !(VB.Vector (SRTree Int)) -- id -> node, children already resolved to ids + , ctRoot :: !Int + , ctDyn :: !(VU.Vector Bool) -- id -> depends on theta? + , ctStatic :: VU.Vector Double -- flat [staticSlot * m + row]; only static nodes + , ctStaticBase :: !(VU.Vector Int) -- id -> staticSlot * m (0 for dynamic ids and Var leaves) + , ctM :: !Int + , ctNPred :: !Int -- root + 1 (stride for flat static) + , ctKind :: !(VU.Vector Int) -- id -> node kind: 0 Var, 1 Param, 2 Const, 3 Uni, 4 Bin + , ctArg :: !(VU.Vector Int) -- id -> Param: param ix; Var: var ix (-1 = y, -2 = yErr); Uni: child id; Bin: left id + , ctArg2 :: !(VU.Vector Int) -- id -> Bin: right id; else 0 + , ctFcode :: !(VU.Vector Int) -- id -> Uni: fromEnum Function + , ctOcode :: !(VU.Vector Int) -- id -> Bin: fromEnum Op + , ctVars :: !(VB.Vector (VU.Vector Double)) -- leaf source columns xss ++ [y, yErr] (referenced, not copied) + } diff --git a/src/Algorithm/SRTree/AD/Unboxed.hs b/src/Algorithm/SRTree/AD/Unboxed.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD/Unboxed.hs @@ -0,0 +1,965 @@ +{-# language FlexibleInstances, DeriveFunctor #-} +{-# language ScopedTypeVariables #-} +{-# language RankNTypes #-} +{-# language ViewPatterns #-} +{-# language FlexibleContexts #-} +{-# language BangPatterns #-} +{-# language TypeApplications #-} +{-# language MultiWayIf #-} +{-# LANGUAGE LambdaCase #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.AD +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Automatic Differentiation for Expression trees +-- +----------------------------------------------------------------------------- + +module Algorithm.SRTree.AD.Unboxed + ( compileTree + , compileTreeMulti + , evalGradMulti + , evalGrad + , evalGradVec + , evalLossVec + , CompiledTree(..) + , setMTPopParallel + ) where + +import Control.Monad (forM_, foldM, when, unless) +import Control.Monad.ST +import Data.STRef (newSTRef, readSTRef, modifySTRef') +import Data.Bifunctor (bimap, first, second) +import Data.SRTree.Derivative ( derivative ) +import Data.SRTree.Eval + ( Target, Theta, Columns, evalFun, evalOp, replicateAs ) +import Data.SRTree.Internal +import Data.SRTree.Print (showExpr) +import Data.SRTree.Recursion ( cataM, cata, accu ) +import qualified Data.Vector.Storable as V +import qualified Data.Vector.Storable.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VMB +import Debug.Trace (trace, traceShow) +import qualified Data.IntMap.Strict as IntMap +import Data.List ( foldl', foldl1' ) +import Data.Maybe (isJust, fromMaybe) + +import Control.Monad.State.Strict +import Control.Monad.Identity + + +import Data.List (transpose) +import System.IO.Unsafe (unsafePerformIO) +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (forConcurrently) +import Control.Exception (evaluate) +import Data.IORef (IORef, newIORef, writeIORef, readIORef) + +import qualified Data.Map.Strict as Map +import Algorithm.SRTree.AD.CompiledAD + +compileTree :: [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double) -> Fix SRTree -> CompiledTree +compileTree xss ys mYErr tree = + CompiledTree { ctNodes = nodes, ctRoot = root, ctDyn = dynArr, ctStatic = staticArr, ctStaticBase = staticBaseArr, ctM = m, ctNPred = root + 1 + , ctKind = kindArr, ctArg = argArr, ctArg2 = arg2Arr, ctFcode = fcodeArr, ctOcode = ocodeArr, ctVars = vars } + where + -- yErr is only defined when mYErr is present (a tree referencing Var -2 + -- always pairs with mYErr = Just e, see the likelihood loss wrappers). The + -- ctVars list must stay well-defined for every column even when mYErr is + -- Nothing -- the Accelerate leaf array concatenates the whole list -- so a + -- missing yErr is represented by a zero column rather than the bottom + -- `fromJust mYErr` (which the old static-array copy path could keep lazy). + yErr = case mYErr of + Just e -> e + Nothing -> VU.replicate m 0 + m = VU.length ys + -- Leaf source columns, referenced (never copied per tree): a static Var + -- leaf reads feature column ix (arg), y (arg = -1), or yErr (arg = -2) + -- straight from these run-fixed vectors instead of a materialized copy in + -- staticArr. ctVars ix = xss !! ix, ctVars nFeats = y, ctVars (nFeats+1) + -- = yErr. + vars = VB.fromList (xss <> [ys, yErr]) + nFeats = VB.length vars - 2 + + -- Rewrite x ** 2.0 into the unary Square kernel (x*x, fcode 17): + -- the loss wrap ((tree - y) ** 2) / m is the single hottest subgraph in + -- every NLopt call, and replacing the per-element pow with a multiply + -- avoids the slow ** (x**2.0 == x*x exactly, and the derivative 2x + -- matches), so no numerical semantics change. + tree' = rewritePowSq tree + + -- state: (structural CSE map, id -> node, id -> isDynamic, counter) + (_, int2key, dynMap, (subtract 1) -> root) = + cataM leftToRight alg tree' + `execState` (Map.empty, IntMap.empty, IntMap.empty, 0) + + nodes = VB.fromList (IntMap.elems int2key) + dynArr = VU.fromList (IntMap.elems dynMap) + stride = root + 1 + -- static nodes in ascending (topological) id order, so a single bottom-up + -- sweep fills every column before its parent. Dynamic nodes are omitted + -- entirely, and so are Var leaves (their values are read directly from the + -- run-fixed `vars` columns, see the eval kernels): their static slots were + -- zeros that evalGrad*/forwardPassRange never read, so the flat array + -- shrinks from stride * m to #static * m (a handful of feature/const + -- columns per tree instead of all nodes). + staticKeys = [k | k <- [0 .. root], not (VU.unsafeIndex dynArr k), not (isVarLeaf k)] + nStatic = length staticKeys + isVarLeaf k = case VB.unsafeIndex nodes k of { Var _ -> True; _ -> False } + -- id -> static slot base (slot * m); 0 for dynamic ids (never read) + staticBaseArr = VU.create $ do + arr <- VUM.replicate (root + 1) 0 + forM_ (zip staticKeys [0 ..]) $ \(k, slot) -> + VUM.write arr k (slot * m) + pure arr + -- flat [slot * m + row]; computed in a single bottom-up sweep over the + -- static ids (a child always gets a smaller id than its parent, since + -- cataM assigns the id only after both children are built), writing each + -- static node's column directly into the flat array. This fuses the old + -- per-node VU.map/VU.zipWith intermediates into the array. + staticArr = VU.create $ do + arr <- VUM.replicate (nStatic * m) 0 + let slice slot = VUM.slice (slot * m) m arr + slotOf k = VU.unsafeIndex staticBaseArr k `div` m + -- Resolve a static child @c@ to its source column ONCE per + -- column (hoisted out of the row loop): Var leaves are not + -- materialized in staticArr, so their column is the run-fixed + -- `vars` vector; every other static node is a column already + -- written into the (mutable) arr (children always have smaller + -- ids). `Left` = pure vector (Var leaf), `Right` = mutable slice. + staticSrc c + | isVarLeaf c = Left (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex argArr c))) + | otherwise = Right (slice (slotOf c)) + -- Read row @i from a hoisted source (see staticSrc). Called per + -- element, but the Left/Right tag is fixed per column, so GHC + -- keeps the dispatch cheap and no slice/leaf lookup is repeated. + readSrc s i = case s of + Left v -> pure (VU.unsafeIndex v i) + Right m -> VUM.unsafeRead m i + mapStatic f t k = go 0 + where + dst = slice (slotOf k) + src = staticSrc t + go !i | i >= m = pure () + | otherwise = do + x <- readSrc src i + VUM.unsafeWrite dst i (evalFun f x) + go (i + 1) + zipStatic op l r k = go 0 + where + dst = slice (slotOf k) + srcL = staticSrc l + srcR = staticSrc r + go !i | i >= m = pure () + | otherwise = do + xl <- readSrc srcL i + xr <- readSrc srcR i + VUM.unsafeWrite dst i (evalOp op xl xr) + go (i + 1) + forM_ (zip staticKeys [0 ..]) $ \(k, slot) -> + case VB.unsafeIndex nodes k of + -- Var leaves are excluded from staticKeys (their columns live + -- in `vars`), so they never reach this sweep. + Const v -> VUM.set (slice slot) v + Uni f t -> mapStatic f t k + Bin op l r -> zipStatic op l r k + Param _ -> pure () + Var _ -> pure () + pure arr + + -- compact unboxed per-id code arrays (length root+1) so the hot row loop + -- never touches the boxed `nodes` vector nor dispatches through the + -- function-returning evalOp/evalFun + kindArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Var _ -> 0 + Param _ -> 1 + Const _ -> 2 + Uni _ _ -> 3 + Bin _ _ _ -> 4 + argArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Var ix -> ix + Param ix -> ix + Uni _ t -> t + Bin _ l _ -> l + Const _ -> 0 + arg2Arr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Bin _ _ r -> r + _ -> 0 + fcodeArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Uni f _ -> fromEnum f + _ -> 0 + ocodeArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Bin op _ _ -> fromEnum op + _ -> 0 + + 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) + + alg = insertKey + + graph (a, _, _, _) = a + isDynSt k (_, _, d, _) = d IntMap.! k + + insEntry key isD (a, b, d, c) = + ( Map.insert key c a + , IntMap.insert c key b + , IntMap.insert c isD d + , c + 1 ) + + -- a node depends on theta iff it IS a Param, or any child does + nodeIsDynamic (Param _) = pure True + nodeIsDynamic (Var _) = pure False + nodeIsDynamic (Const _) = pure False + nodeIsDynamic (Uni _ t) = gets (isDynSt t) + nodeIsDynamic (Bin _ l r) = (||) <$> gets (isDynSt l) <*> gets (isDynSt r) + + -- Data.Map is unreliable with NaN-valued keys (Ord Double is not a valid + -- total order for NaN: insert(Const NaN) then member/lookup can disagree), + -- and eqsat constant folding can yield Const NaN nodes. So do a single + -- direct lookup; on a miss, return the fresh id that insEntry assigns + -- instead of looking the key back up. Repeated NaN nodes simply get + -- separate ids (no CSE), which is harmless since their static value is + -- recomputed identically. + insertKey key = do + cached <- gets (Map.lookup key . graph) + case cached of + Just v -> pure v + Nothing -> do + d <- nodeIsDynamic key + fresh <- state $ \st@(_, _, _, c) -> let st' = insEntry key d st in (c, st') + pure fresh + +-- Rewrite (a) Bin Power t (Const 2.0) into the unary Square kernel and +-- (b) Bin Div t (Const c) into Bin Mul t (Const (1/c)). Both are exact at +-- the Double level (x ** 2.0 == x * x; x / c == x * (1/c) up to one ulp) +-- and replace the slow per-element pow()/div with a multiply. The loss +-- wrap ((tree - y) ** 2) / m appears in every NLopt objective/gradient +-- call, so these two rewrites are worth a measurable fraction of the AD +-- time. +rewritePowSq :: Fix SRTree -> Fix SRTree +rewritePowSq = cata alg + where + alg :: SRTree (Fix SRTree) -> Fix SRTree + alg (Bin Power t (Fix (Const 2.0))) = Fix (Uni Square t) + alg (Bin Div t (Fix (Const c))) | c /= 0 = Fix (Bin Mul t (Fix (Const (recip c)))) + alg n = Fix n + +-- --------------------------------------------------------------------- +-- Static-child source resolution. A static Var leaf is NOT materialized +-- into ctStatic anymore: its value column lives in ctVars (= xss ++ [y, +-- yErr]) and is read directly at the absolute row (base 0, so the chunk +-- start s0 positions the read). Every other static node is a computed +-- column inside ctStatic at ctStaticBase k. +-- --------------------------------------------------------------------- + +-- | Map a Var leaf's arg (feature ix, or -1 = y, -2 = yErr) to an index +-- into ctVars = xss ++ [y, yErr]. +leafSrcIdx :: Int -> Int -> Int +leafSrcIdx nFeats a | a >= 0 = a + | a == -1 = nFeats + | otherwise = nFeats + 1 +{-# INLINE leafSrcIdx #-} + +-- | Resolve the (source vector, base) of a static child node @k@, where the +-- row value is read at @src (base + i)@. +resolveStatic :: VU.Vector Double + -> VB.Vector (VU.Vector Double) + -> VU.Vector Int + -> VU.Vector Int + -> VU.Vector Int + -> Int -> Int -> Int + -> (VU.Vector Double, Int) +resolveStatic static vars kind arg staticBase nFeats s0 k = + if VU.unsafeIndex kind k == 0 + then (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex arg k)), s0) + else (static, VU.unsafeIndex staticBase k + s0) +{-# INLINE resolveStatic #-} + +-- --------------------------------------------------------------------- +-- Per-theta evaluation: the hot path, called once per NLopt objective/ +-- gradient call. Forward pass only recomputes dynamic nodes (ids are +-- already topologically ordered, so a single left-to-right fold works). +-- Backward pass is the same recursive shape as the original calcGrad, +-- except it stops immediately on any non-dynamic node -- that subtree +-- has no Param in it, so it can never contribute to the gradient. +-- --------------------------------------------------------------------- + +-- Row-fused evaluation: instead of storing one full length-m array per +-- node (which meant ~2 * #nodes large allocations per objective/gradient +-- call), we walk the m data rows one at a time and, for each row, run the +-- forward pass and the reverse-mode backward pass over small per-node +-- scratch arrays of Double (length root+1). This mirrors what the fused +-- Accelerate/LLVM kernel does (one pass per row, no big intermediate +-- arrays) while staying in plain ST: allocation drops from O(nodes * m) +-- to O(nodes + params), and the tight inner loops are all unboxed. +evalGrad :: CompiledTree -> V.Vector Double -> (Double, V.Vector Double) +evalGrad ct theta = runST $ do + fwd <- VUM.new (root + 1) -- node id -> forward value, current row + adj <- VUM.new (root + 1) -- node id -> adjoint (dL/dnode), current row + gradM <- VUM.replicate p 0 -- accumulated per-parameter gradient + objRef <- newSTRef 0 + + let -- forward pass for a single row: fills `fwd` for ids 0..root + forwardLoop !row !key + | key > root = pure () + | otherwise = do + v <- if not (VU.unsafeIndex dyn key) + then if VU.unsafeIndex kind key == 0 + then pure (VU.unsafeIndex (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex arg key))) row) + else pure (VU.unsafeIndex static (VU.unsafeIndex staticBase key + row)) + else case VU.unsafeIndex kind key of + 1 -> pure (V.unsafeIndex theta (VU.unsafeIndex arg key)) + 3 -> do x <- VUM.unsafeRead fwd (VU.unsafeIndex arg key) + pure (evalFunCode (VU.unsafeIndex fcode key) x) + 4 -> do xl <- VUM.unsafeRead fwd (VU.unsafeIndex arg key) + xr <- VUM.unsafeRead fwd (VU.unsafeIndex arg2 key) + pure (evalOpCode (VU.unsafeIndex ocode key) xl xr) + _ -> error "evalGrad: unreachable" + VUM.unsafeWrite fwd key v + forwardLoop row (key + 1) + + -- backward pass for a single row: ids are visited from root down + -- to 0, which is a valid reverse-topological order since every + -- child id is smaller than its parent's id by construction. + backwardLoop !key + | key < 0 = pure () + | otherwise = do + when (VU.unsafeIndex dyn key) $ do + v <- VUM.unsafeRead adj key + case VU.unsafeIndex kind key of + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + xl <- VUM.unsafeRead fwd l + xr <- VUM.unsafeRead fwd r + fg <- VUM.unsafeRead fwd key + let (dl, dr) = diffScalarCode (VU.unsafeIndex ocode key) v xl xr fg + VUM.unsafeModify adj (+ dl) l + VUM.unsafeModify adj (+ dr) r + 3 -> do + let t = VU.unsafeIndex arg key + x <- VUM.unsafeRead fwd t + VUM.unsafeModify adj (+ v * derivFunCode (VU.unsafeIndex fcode key) x) t + 1 -> VUM.unsafeModify gradM (+ v) (VU.unsafeIndex arg key) + _ -> pure () + backwardLoop (key - 1) + + rowLoop !row + | row >= m = pure () + | otherwise = do + forwardLoop row 0 + rootVal <- VUM.unsafeRead fwd root + modifySTRef' objRef (+ rootVal) + when (VU.unsafeIndex dyn root) $ do + VUM.set adj 0 + VUM.unsafeWrite adj root 1 + backwardLoop root + rowLoop (row + 1) + + rowLoop 0 + + obj <- readSTRef objRef + gradFrozen <- VU.unsafeFreeze gradM + pure (obj, V.convert gradFrozen) + where + root = ctRoot ct + m = ctM ct + p = V.length theta + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + +-- --------------------------------------------------------------------- +-- Node-outer (vectorized-over-rows) evaluation: mirrors reverseModeGraph's +-- shape (one full length-m column per node, node-major loops) so the inner +-- loops are fused per node over all rows, with the static/dynamic pattern +-- decided once per node instead of once per row. Uses the same flat +-- [staticSlot * m + row] layout and compact op-code dispatch as `evalGrad`, but +-- trades the O(nodes + params) scratch of the row-fused version for the +-- O(nodes * m) fwd/adj columns of the massiv-style whole-column kernel. +evalGradVec :: CompiledTree -> V.Vector Double -> (Double, V.Vector Double) +evalGradVec ct theta = runST $ do + -- Per-chunk buffers of O(stride * chunk) instead of one O(stride * m) + -- allocation per call: the fwd/adj matrices are streamed one chunk of + -- `chunk` rows at a time, so the per-call allocation drops ~m/chunk x + -- (and the working set stays L3-resident). The chunk partition does not + -- change any value: each row is independent, the objective row sums + -- accumulate in order and the gradient accumulates row-sums per chunk. + fwd <- VUM.new (stride * chunk) -- [node * nb + i]; dynamic columns written before read + adj <- VUM.replicate (stride * chunk) 0 -- [node * nb + i] + gradM <- VUM.replicate p 0 + + let go !start !acc + | start >= m = do + gradFrozen <- VU.unsafeFreeze gradM + pure (acc, V.convert gradFrozen) + | otherwise = do + let nb = min chunk (m - start) + s0 = start + forwardPassRange ct theta fwd s0 nb + -- objective contribution = sum over this chunk's rows of root + s <- if VU.unsafeIndex dyn root + then {-# SCC "objSumFwd" #-} sumCol fwd (root * nb) nb + else {-# SCC "objSumStatic" #-} sumStatic (VU.unsafeIndex staticBase root + s0) nb + -- seed the root adjoint: d(obj)/d(root value) = 1 per row + unless (s0 == 0) $ VUM.set adj 0 -- reuse the buffer; keep it clean + when (VU.unsafeIndex dyn root) $ {-# SCC "seedAdj" #-} VUM.set (VUM.slice (root * nb) nb adj) 1 + -- backward: nodes from root down to 0 (valid reverse-topological order) + let goBwd !key + | key < 0 = pure () + | otherwise = do + bwdNode key + goBwd (key - 1) + + bwdNode key + | not (VU.unsafeIndex dyn key) = pure () -- no Param in subtree + | otherwise = case VU.unsafeIndex kind key of + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + oc = VU.unsafeIndex ocode key + dl = VU.unsafeIndex dyn l + dr = VU.unsafeIndex dyn r + kb = key * nb + lb = l * nb + rb = r * nb + case (dl, dr) of + (True, True) -> {-# SCC "bwdBinTT" #-} bwdBin nb fwd adj 0 oc kb lb rb (static, 0) + (True, False) -> {-# SCC "bwdBinTS" #-} bwdBin nb fwd adj 1 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 r) + (False, True) -> {-# SCC "bwdBinST" #-} bwdBin nb fwd adj 2 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 l) + (False, False) -> pure () -- no dynamic children to propagate to + 3 -> do + let t = VU.unsafeIndex arg key + fc = VU.unsafeIndex fcode key + kb = key * nb + tb = t * nb + if VU.unsafeIndex dyn t + then {-# SCC "bwdUni" #-} bwdUni nb fwd adj fc kb tb + else pure () -- static child: no Param below, nothing to accumulate + 1 -> do + let a = VU.unsafeIndex arg key + kb = key * nb + {-# SCC "bwdParam" #-} do + s' <- sumCol adj kb nb + VUM.unsafeModify gradM (+ s') a + _ -> pure () + goBwd root + go (start + nb) (acc + s) + + go 0 0 + where + root = ctRoot ct + m = ctM ct + p = V.length theta + stride = root + 1 + chunk = 1024 + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + + sumCol v vbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = VUM.unsafeRead v (vbase + i) >>= \vv -> go (i + 1) (acc + vv) + sumStatic sbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = go (i + 1) (acc + VU.unsafeIndex static (sbase + i)) + +-- --------------------------------------------------------------------- +-- Forward-only objective evaluation: runs the forward pass and the row +-- sum but skips the adjoint/backward pass. Used where only the objective +-- value is needed (reporting loss / R2 metrics, the validation fitness in +-- the search), avoiding the ~2/3 of evalGradVec's work that computes the +-- gradient. +-- --------------------------------------------------------------------- +-- Chunked loss evaluation: runs the same node-outer forward pass as +-- `evalGradVec` (static columns precomputed in `ctStatic`, op codes +-- dispatched once per node into INLINE kernels) but only over a chunk of +-- `chunk` rows at a time with a per-call buffer of O(stride * chunk) +-- instead of O(stride * m). The chunk partition does not change any value +-- (each row is computed independently, the row sums accumulate in order), +-- but it cuts the per-call allocation ~30x so this is cheap enough for the +-- val-eval hot path that runs once per explored expression. +evalLossVec :: CompiledTree -> V.Vector Double -> Double +evalLossVec ct theta = runST $ do + buf <- VUM.new (stride * chunk) + go buf 0 0 + where + root = ctRoot ct + m = ctM ct + stride = root + 1 + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + chunk = 4096 + + go :: VUM.MVector s Double -> Int -> Double -> ST s Double + go buf !start !acc + | start >= m = pure acc + | otherwise = do + let nb = min chunk (m - start) + forwardPassRange ct theta buf start nb + s <- if VU.unsafeIndex dyn root + then sumCol buf (root * nb) nb + else sumStatic (VU.unsafeIndex staticBase root + start) nb + go buf (start + nb) (acc + s) + + sumCol buf vbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = VUM.unsafeRead buf (vbase + i) >>= \vv -> go (i + 1) (acc + vv) + sumStatic sbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = go (i + 1) (acc + VU.unsafeIndex static (sbase + i)) + +-- Forward pass shared by evalGradVec and evalLossVec: fills the `fwd` +-- columns of every dynamic node (ids are topologically ordered, so one +-- left-to-right sweep computes all of them; static columns are already in +-- `ctStatic`). The op/function codes are dispatched once per node and the +-- INLINE loop helpers run a tight fused kernel over the rows. +-- +-- `s0`/`nb` select a range of rows [start, start+nb): with nb = m, start = 0 +-- this is the full-matrix pass used by evalGradVec; evalLossVec calls it on +-- chunks of rows with a stride*nb buffer. The fwd buffer is indexed +-- [key * nb + i], static columns are read at [slot(key) * m + s0 + i]. +forwardPassRange :: CompiledTree -> V.Vector Double -> VUM.MVector s Double -> Int -> Int -> ST s () +forwardPassRange ct theta fwd s0 nb = goFwd 0 + where + root = ctRoot ct + m = ctM ct + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + + goFwd !key + | key > root = pure () + | otherwise = do + if VU.unsafeIndex dyn key + then case VU.unsafeIndex kind key of + 1 -> {-# SCC "fwdParam" #-} VUM.set (VUM.slice (key * nb) nb fwd) (V.unsafeIndex theta (VU.unsafeIndex arg key)) + 3 -> do + let t = VU.unsafeIndex arg key + fc = VU.unsafeIndex fcode key + kb = key * nb + tb = t * nb + if VU.unsafeIndex dyn t + then {-# SCC "fwdUniD" #-} fwdUniD nb fwd fc kb tb + else pure () -- a dynamic Uni always has a dynamic child + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + oc = VU.unsafeIndex ocode key + dl = VU.unsafeIndex dyn l + dr = VU.unsafeIndex dyn r + kb = key * nb + lb = l * nb + rb = r * nb + case (dl, dr) of + (True, True) -> {-# SCC "fwdBinTT" #-} fwdBin nb fwd 0 oc kb lb rb (static, 0) + (True, False) -> {-# SCC "fwdBinTS" #-} fwdBin nb fwd 1 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 r) + (False, True) -> {-# SCC "fwdBinST" #-} fwdBin nb fwd 2 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 l) + (False, False) -> pure () -- unreachable: a dynamic Bin always has a dynamic child + _ -> pure () + else pure () -- static node: column already in `static` + goFwd (key + 1) + + -- Forward binary kernels: `combo` 0=TT, 1=TS, 2=ST (SS is unreachable + -- for dynamic nodes). The opcode is dispatched ONCE per node; the loop + -- helpers are INLINE with the literal operator so each row iteration + -- is a tight fused kernel with no per-element `case oc of` dispatch. + -- `stSrc` is the (source vector, base) of the static child (either a + -- run-fixed leaf column from `vars` or a computed column of `static`), + -- used by the TS/ST variants; `nb` is the number of rows in this chunk. +fwdBin :: Int -> VUM.MVector s Double -> Int -> Int -> Int -> Int -> Int -> (VU.Vector Double, Int) -> ST s () +fwdBin nb fwd combo oc kb lb rb stSrc = case (combo, oc) of + (0, 0) -> fwdTT nb fwd (+) kb lb rb + (0, 1) -> fwdTT nb fwd (-) kb lb rb + (0, 2) -> fwdTT nb fwd (*) kb lb rb + (0, 3) -> fwdTT nb fwd (/) kb lb rb + (0, 4) -> fwdTT nb fwd (**) kb lb rb + (0, 5) -> fwdTT nb fwd (\l r -> abs l ** r) kb lb rb + (0, 6) -> fwdTT nb fwd (\l r -> l / sqrt (1 + r * r)) kb lb rb + (1, 0) -> fwdTS nb stSrc fwd (+) kb lb + (1, 1) -> fwdTS nb stSrc fwd (-) kb lb + (1, 2) -> fwdTS nb stSrc fwd (*) kb lb + (1, 3) -> fwdTS nb stSrc fwd (/) kb lb + (1, 4) -> fwdTS nb stSrc fwd (**) kb lb + (1, 5) -> fwdTS nb stSrc fwd (\l r -> abs l ** r) kb lb + (1, 6) -> fwdTS nb stSrc fwd (\l r -> l / sqrt (1 + r * r)) kb lb + (2, 0) -> fwdST nb stSrc fwd (+) kb rb + (2, 1) -> fwdST nb stSrc fwd (-) kb rb + (2, 2) -> fwdST nb stSrc fwd (*) kb rb + (2, 3) -> fwdST nb stSrc fwd (/) kb rb + (2, 4) -> fwdST nb stSrc fwd (**) kb rb + (2, 5) -> fwdST nb stSrc fwd (\l r -> abs l ** r) kb rb + (2, 6) -> fwdST nb stSrc fwd (\l r -> l / sqrt (1 + r * r)) kb rb + _ -> pure () +{-# INLINE fwdBin #-} + +-- Backward binary kernels: same dispatch structure, `diff` is the local +-- (dl/dchild, dr/dchild) rule keyed on the opcode. `nb` is the number of +-- rows in this chunk, `stSrc` is the (source vector, base) of the static +-- child (either a run-fixed leaf column from `vars` or a computed column of +-- `static`), used by the TS/ST variants. +bwdBin :: Int -> VUM.MVector s Double -> VUM.MVector s Double -> Int -> Int -> Int -> Int -> Int -> (VU.Vector Double, Int) -> ST s () +bwdBin nb fwd adj combo oc kb lb rb stSrc = case (combo, oc) of + (0, 0) -> bwdTT nb fwd adj (\dx _ _ _ -> (dx, dx)) kb lb rb + (0, 1) -> bwdTT nb fwd adj (\dx _ _ _ -> (dx, negate dx)) kb lb rb + (0, 2) -> bwdTT nb fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb lb rb + (0, 3) -> bwdTT nb fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb lb rb + (0, 4) -> bwdTT nb fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb lb rb + (0, 5) -> bwdTT nb fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb rb + (0, 6) -> bwdTT nb fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb lb rb + (1, 0) -> bwdTS nb stSrc fwd adj (\dx _ _ _ -> (dx, dx)) kb lb + (1, 1) -> bwdTS nb stSrc fwd adj (\dx _ _ _ -> (dx, negate dx)) kb lb + (1, 2) -> bwdTS nb stSrc fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb lb + (1, 3) -> bwdTS nb stSrc fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb lb + (1, 4) -> bwdTS nb stSrc fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb lb + (1, 5) -> bwdTS nb stSrc fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb + (1, 6) -> bwdTS nb stSrc fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb lb + (2, 0) -> bwdST nb stSrc fwd adj (\dx _ _ _ -> (dx, dx)) kb rb + (2, 1) -> bwdST nb stSrc fwd adj (\dx _ _ _ -> (dx, negate dx)) kb rb + (2, 2) -> bwdST nb stSrc fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb rb + (2, 3) -> bwdST nb stSrc fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb rb + (2, 4) -> bwdST nb stSrc fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb rb + (2, 5) -> bwdST nb stSrc fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb + (2, 6) -> bwdST nb stSrc fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb rb + _ -> pure () +{-# INLINE bwdBin #-} + +fwdTT nb fwd op kb lb rb = forRows nb $ \i -> do + xl <- VUM.unsafeRead fwd (lb + i) + xr <- VUM.unsafeRead fwd (rb + i) + VUM.unsafeWrite fwd (kb + i) (op xl xr) +{-# INLINE fwdTT #-} + +fwdTS nb (src, base) fwd op kb lb = forRows nb $ \i -> do + xl <- VUM.unsafeRead fwd (lb + i) + VUM.unsafeWrite fwd (kb + i) (op xl (VU.unsafeIndex src (base + i))) +{-# INLINE fwdTS #-} + +fwdST nb (src, base) fwd op kb rb = forRows nb $ \i -> do + xr <- VUM.unsafeRead fwd (rb + i) + VUM.unsafeWrite fwd (kb + i) (op (VU.unsafeIndex src (base + i)) xr) +{-# INLINE fwdST #-} + +bwdTT nb fwd adj diff kb lb rb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xl <- VUM.unsafeRead fwd (lb + i) + xr <- VUM.unsafeRead fwd (rb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (gl, gr) = diff v xl xr fg + a <- VUM.unsafeRead adj (lb + i) + VUM.unsafeWrite adj (lb + i) (a + gl) + b <- VUM.unsafeRead adj (rb + i) + VUM.unsafeWrite adj (rb + i) (b + gr) +{-# INLINE bwdTT #-} + +bwdTS nb (src, base) fwd adj diff kb lb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xl <- VUM.unsafeRead fwd (lb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (gl, _) = diff v xl (VU.unsafeIndex src (base + i)) fg + a <- VUM.unsafeRead adj (lb + i) + VUM.unsafeWrite adj (lb + i) (a + gl) +{-# INLINE bwdTS #-} + +bwdST nb (src, base) fwd adj diff kb rb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xr <- VUM.unsafeRead fwd (rb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (_, gr) = diff v (VU.unsafeIndex src (base + i)) xr fg + b <- VUM.unsafeRead adj (rb + i) + VUM.unsafeWrite adj (rb + i) (b + gr) +{-# INLINE bwdST #-} + +-- Forward unary kernels (dynamic child): the function code is dispatched +-- ONCE per node and the loop helper is INLINE with the literal function, +-- so each row iteration is a tight fused kernel with no per-element +-- `case fc of` / closure build (a dynamic Uni node always has a dynamic +-- child, so there is no static-child variant here). +fwdUniD :: Int -> VUM.MVector s Double -> Int -> Int -> Int -> ST s () +fwdUniD nb fwd fc kb tb = case fc of + 0 -> fwdUniD' nb fwd (\x -> x) kb tb + 1 -> fwdUniD' nb fwd abs kb tb + 2 -> fwdUniD' nb fwd sin kb tb + 3 -> fwdUniD' nb fwd cos kb tb + 4 -> fwdUniD' nb fwd tan kb tb + 5 -> fwdUniD' nb fwd sinh kb tb + 6 -> fwdUniD' nb fwd cosh kb tb + 7 -> fwdUniD' nb fwd tanh kb tb + 8 -> fwdUniD' nb fwd asin kb tb + 9 -> fwdUniD' nb fwd acos kb tb + 10 -> fwdUniD' nb fwd atan kb tb + 11 -> fwdUniD' nb fwd asinh kb tb + 12 -> fwdUniD' nb fwd acosh kb tb + 13 -> fwdUniD' nb fwd atanh kb tb + 14 -> fwdUniD' nb fwd sqrt kb tb + 15 -> fwdUniD' nb fwd (\x -> sqrt (abs x)) kb tb + 16 -> fwdUniD' nb fwd (\x -> signum x * abs x ** (1 / 3)) kb tb + 17 -> fwdUniD' nb fwd (\x -> x * x) kb tb + 18 -> fwdUniD' nb fwd log kb tb + 19 -> fwdUniD' nb fwd (\x -> log (abs x)) kb tb + 20 -> fwdUniD' nb fwd exp kb tb + 21 -> fwdUniD' nb fwd recip kb tb + 22 -> fwdUniD' nb fwd (\x -> x * x * x) kb tb + _ -> pure () +{-# INLINE fwdUniD #-} + +fwdUniD' nb fwd f kb tb = forRows nb $ \i -> do + x <- VUM.unsafeRead fwd (tb + i) + VUM.unsafeWrite fwd (kb + i) (f x) +{-# INLINE fwdUniD' #-} + +-- Backward unary kernel: derivative of the function, dispatched once per +-- node and inlined into the accumulation loop. `nb` is the number of rows +-- in the current chunk. +bwdUni :: Int -> VUM.MVector s Double -> VUM.MVector s Double -> Int -> Int -> Int -> ST s () +bwdUni nb fwd adj fc kb tb = case fc of + 0 -> bwdUni' nb fwd adj (\_ -> 1) kb tb + 1 -> bwdUni' nb fwd adj (\x -> x / abs x) kb tb + 2 -> bwdUni' nb fwd adj cos kb tb + 3 -> bwdUni' nb fwd adj (negate . sin) kb tb + 4 -> bwdUni' nb fwd adj (\x -> 1 / (cos x * cos x)) kb tb + 5 -> bwdUni' nb fwd adj cosh kb tb + 6 -> bwdUni' nb fwd adj sinh kb tb + 7 -> bwdUni' nb fwd adj (\x -> 1 - tanh x * tanh x) kb tb + 8 -> bwdUni' nb fwd adj (\x -> 1 / sqrt (1 - x * x)) kb tb + 9 -> bwdUni' nb fwd adj (\x -> -1 / sqrt (1 - x * x)) kb tb + 10 -> bwdUni' nb fwd adj (\x -> 1 / (1 + x * x)) kb tb + 11 -> bwdUni' nb fwd adj (\x -> 1 / sqrt (1 + x * x)) kb tb + 12 -> bwdUni' nb fwd adj (\x -> 1 / (sqrt (x - 1) * sqrt (x + 1))) kb tb + 13 -> bwdUni' nb fwd adj (\x -> 1 / (1 - x * x)) kb tb + 14 -> bwdUni' nb fwd adj (\x -> 1 / (2 * sqrt x)) kb tb + 15 -> bwdUni' nb fwd adj (\x -> x / (2 * abs x ** (3 / 2))) kb tb + 16 -> bwdUni' nb fwd adj (\x -> 1 / (3 * (x * x) ** (1 / 3))) kb tb + 17 -> bwdUni' nb fwd adj (\x -> 2 * x) kb tb + 18 -> bwdUni' nb fwd adj recip kb tb + 19 -> bwdUni' nb fwd adj recip kb tb + 20 -> bwdUni' nb fwd adj exp kb tb + 21 -> bwdUni' nb fwd adj (\x -> -1 / (x * x)) kb tb + 22 -> bwdUni' nb fwd adj (\x -> 3 * x * x) kb tb + _ -> pure () +{-# INLINE bwdUni #-} + +bwdUni' nb fwd adj f kb tb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + x <- VUM.unsafeRead fwd (tb + i) + c <- VUM.unsafeRead adj (tb + i) + VUM.unsafeWrite adj (tb + i) (c + v * f x) +{-# INLINE bwdUni' #-} +-- Unboxed ST loop over the m data rows; always inlined so the per-node +-- bodies above are fused into a single tail-recursive kernel per node. +forRows :: Int -> (Int -> ST s ()) -> ST s () +forRows !n f = go 0 + where + go !i | i >= n = pure () + | otherwise = f i >> go (i + 1) +{-# INLINE forRows #-} +evalOpCode :: Int -> Double -> Double -> Double +evalOpCode 0 = (+) +evalOpCode 1 = (-) +evalOpCode 2 = (*) +evalOpCode 3 = (/) +evalOpCode 4 = (**) +evalOpCode 5 = \l r -> abs l ** r +evalOpCode 6 = \l r -> l / sqrt (1 + r * r) +evalOpCode _ = error "evalOpCode: bad op code" +{-# INLINE evalOpCode #-} + +evalFunCode :: Int -> Double -> Double +evalFunCode 0 = id +evalFunCode 1 = abs +evalFunCode 2 = sin +evalFunCode 3 = cos +evalFunCode 4 = tan +evalFunCode 5 = sinh +evalFunCode 6 = cosh +evalFunCode 7 = tanh +evalFunCode 8 = asin +evalFunCode 9 = acos +evalFunCode 10 = atan +evalFunCode 11 = asinh +evalFunCode 12 = acosh +evalFunCode 13 = atanh +evalFunCode 14 = sqrt +evalFunCode 15 = \x -> sqrt (abs x) +evalFunCode 16 = \x -> signum x * abs x ** (1 / 3) +evalFunCode 17 = \x -> x * x +evalFunCode 18 = log +evalFunCode 19 = \x -> log (abs x) +evalFunCode 20 = exp +evalFunCode 21 = recip +evalFunCode 22 = \x -> x * x * x +evalFunCode _ = error "evalFunCode: bad function code" +{-# INLINE evalFunCode #-} + +derivFunCode :: Int -> Double -> Double +derivFunCode 0 = const 1 +derivFunCode 1 = \x -> x / abs x +derivFunCode 2 = cos +derivFunCode 3 = negate . sin +derivFunCode 4 = \x -> 1 / (cos x * cos x) +derivFunCode 5 = cosh +derivFunCode 6 = sinh +derivFunCode 7 = \x -> 1 - tanh x * tanh x +derivFunCode 8 = \x -> 1 / sqrt (1 - x * x) +derivFunCode 9 = \x -> -1 / sqrt (1 - x * x) +derivFunCode 10 = \x -> 1 / (1 + x * x) +derivFunCode 11 = \x -> 1 / sqrt (1 + x * x) +derivFunCode 12 = \x -> 1 / (sqrt (x - 1) * sqrt (x + 1)) +derivFunCode 13 = \x -> 1 / (1 - x * x) +derivFunCode 14 = \x -> 1 / (2 * sqrt x) +derivFunCode 15 = \x -> x / (2 * abs x ** (3 / 2)) +derivFunCode 16 = \x -> 1 / (3 * (x * x) ** (1 / 3)) +derivFunCode 17 = (* 2) +derivFunCode 18 = recip +derivFunCode 19 = recip +derivFunCode 20 = exp +derivFunCode 21 = \x -> -1 / (x * x) +derivFunCode 22 = \x -> 3 * x * x +derivFunCode _ = error "derivFunCode: bad function code" +{-# INLINE derivFunCode #-} + +-- Pure local-derivative rules keyed on fromEnum Op, scalar version (same +-- math as the original vectorized `diffPure`, applied per-row above). +diffScalarCode :: Int -> Double -> Double -> Double -> Double -> (Double, Double) +diffScalarCode 0 dx _ _ _ = (dx, dx) +diffScalarCode 1 dx _ _ _ = (dx, negate dx) +diffScalarCode 2 dx fx gy _ = (dx * gy, dx * fx) +diffScalarCode 3 dx _ gy fg = (dx / gy, dx * (negate fg / gy)) +diffScalarCode 4 dx fx gy fg = + ( fixNaN (dx * gy * fg / fx) + , fixNaN (dx * fg * log fx) ) +diffScalarCode 5 dx fx gy fg = + let v2 = abs fx + in ( fixNaN (dx * (fx * gy) * fg / (v2 * v2)) + , fixNaN (dx * fg * log (abs fx)) ) +diffScalarCode 6 dx fx gy _ = + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy) +diffScalarCode _ _ _ _ _ = error "diffScalarCode: bad op code" +{-# INLINE diffScalarCode #-} + +fixNaN :: Double -> Double +fixNaN x = if isNaN x then 0 else x +{-# INLINE fixNaN #-} + +-- --------------------------------------------------------------------- +-- Drop-in-compatible wrapper -- same signature as your original function. +-- Use this ONLY to verify correctness against your existing implementation +-- (e.g. QuickCheck / golden tests comparing outputs). It gets you ZERO +-- speedup on its own, since it calls compileTree fresh every time, same +-- as before. The actual win requires changing the NLopt-facing call site. +-- --------------------------------------------------------------------- + +--reverseModeGraphO :: [V.Vector Double] -> V.Vector Double -> Maybe (V.Vector Double) -> V.Vector Double -> Fix SRTree -> (V.Vector Double, V.Vector Double) +--reverseModeGraphO xss ys mYErr theta tree = evalGrad (compileTree xss ys mYErr tree) theta + +-- | Safely chunk an unboxed vector into 'n' roughly equal parts. +chunkVector :: Int -> VU.Vector Double -> [VU.Vector Double] +chunkVector numChunks v + | VU.null v = [] + | otherwise = + let n = VU.length v + chunkSize = max 1 (n `div` numChunks) + go vec | VU.null vec = [] + | VU.length vec <= chunkSize = [vec] + | otherwise = let (h, t) = VU.splitAt chunkSize vec + in h : go t + in go v + +-- | Compiles the tree for multiple data chunks independently. +compileTreeMulti :: [VU.Vector Double] + -> VU.Vector Double + -> Maybe (VU.Vector Double) + -> Fix SRTree + -> [CompiledTree] +compileTreeMulti xss ys mYErr tree = + let nRows = VU.length ys + minChunkSize = 2000 + numChunks = max 1 (min cap (nRows `div` minChunkSize)) + cap = if mtSingleChunk then 1 else unsafePerformIO getNumCapabilities + ysChunks = chunkVector numChunks ys + -- transpose groups the chunks by slice rather than by feature + xssChunks = Data.List.transpose (map (chunkVector numChunks) xss) + errChunks = case mYErr of + Just e -> map Just (chunkVector numChunks e) + Nothing -> replicate (length ysChunks) Nothing + in [ compileTree xs y err tree | (xs, y, err) <- zip3 xssChunks ysChunks errChunks ] + +-- | When True, the MultiThread backend compiles/evaluates each tree on a +-- single chunk so a higher-level population-parallel driver (eggp's fitness +-- batch) owns the cores instead of oversubscribing the per-tree chunk split. +{-# NOINLINE mtSingleChunk #-} +mtSingleChunk :: Bool +mtSingleChunk = unsafePerformIO (readIORef mtParGate) + +mtParGate :: IORef Bool +mtParGate = unsafePerformIO (newIORef False) +{-# NOINLINE mtParGate #-} + +-- | Enable/disable single-chunk (non-oversubscribing) mode for the MultiThread +-- backend; called around a population-parallel fitness batch. +setMTPopParallel :: Bool -> IO () +setMTPopParallel b = writeIORef mtParGate b + +-- | Evaluates the gradient across all compiled chunks in parallel. +-- Each chunk is evaluated by the fast node-outer `evalGradVec` kernel on +-- its own slice of the data. The kernel is now chunked internally (O(stride +-- * 1024) per-call buffers, L3-resident) and is compute-bound rather than +-- memory-bandwidth-bound, so splitting the data into one chunk per core and +-- running the kernels concurrently scales almost linearly. The objective +-- and gradient accumulate across chunks (same math per row; only the FP +-- summation order across chunk boundaries differs). +evalGradMulti :: [CompiledTree] -> V.Vector Double -> (Double, V.Vector Double) +evalGradMulti [ct] theta = evalGradVec ct theta +evalGradMulti cts theta = unsafePerformIO $ do + results <- forConcurrently cts $ \ct -> evaluate (evalGradVec ct theta) + let totalObj = sum $ map fst results + totalGrad = foldl1' (V.zipWith (+)) (map snd results) + pure (totalObj, totalGrad) diff --git a/src/Algorithm/SRTree/Compile.hs b/src/Algorithm/SRTree/Compile.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Compile.hs @@ -0,0 +1,106 @@ +{-# LANGUAGE GADTs #-} + +module Algorithm.SRTree.Compile where + +import Data.SRTree +import Data.SRTree.Eval (compileLoss, Target, Columns, Theta) +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Generic as G +import Algorithm.SRTree.AD +import Algorithm.SRTree.Utils +import Algorithm.SRTree.Likelihoods (Distribution(..), Loss(..), buildLoss, hessianNLL) +import Algorithm.SRTree.NonlinearOpt (minimizeNLL, minimizeNLLWithFixedParam) +import Data.SRTree.Recursion (cata) + +data EvalTree = EvalTree { + ctDist :: Distribution, + ctLoss :: Theta -> Double, + ctAD :: VS.Vector Double -> (Double, VS.Vector Double), + ctOptimizer :: Target -> Target, + ctOptimizerFixed :: Int -> Target -> Target, + ctNLL :: Target -> Double, + ctGradNLL :: Target -> (Double, Target), + ctHessianNLL :: Target -> Columns, + ctTree :: Fix SRTree, + ctRows :: Int, + ctVar :: Double +} + +-- | Compile a tree and store it in a CompiledTree data structure +compileTree :: Distribution -> Columns -> Target -> Maybe Target -> Fix SRTree -> EvalTree +compileTree dist xss ys mYerr tree = EvalTree { + ctDist = dist, + ctLoss = compileLoss xss tree ys mYerr, + ctAD = compileFunAndGrad MultiThread xss ys mYerr tree, + ctOptimizer = fst3 . minimizeNLL MultiThread (NLL dist) mYerr 100 xss ys tree, + ctOptimizerFixed = minimizeNLLWithFixedParam MultiThread (NLL dist) mYerr 100 xss ys tree, + ctNLL = compileLoss xss lossTree ys mYerr, + ctGradNLL = \theta -> let fg = compileFunAndGrad MultiThread xss ys mYerr lossTree + (obj, gradStorable) = fg (G.convert theta) + in (obj, G.convert gradStorable), + ctHessianNLL = hessianNLL dist mYerr xss ys tree, + ctTree = tree, + ctRows = n, + ctVar = let ym = U.sum ys / fromIntegral n + in U.foldr (\yi acc -> acc + (yi - ym)^2) 0 ys +} + where + n = U.length ys + lossTree = buildLoss (NLL dist) (fromIntegral n) tree + fst3 (a, _, _) = a + +data EvaluatedTree = EvaluatedTree { + valLoss :: Double, + valTheta :: Theta, + valRows :: Double, + valParams :: Double, + valTree :: Fix SRTree, + valLogParams :: Double, + valLogParamsLattice :: Double, + valVar :: Double +} + +evaluateTree :: EvalTree -> Target -> [[Double]] -> Theta -> EvaluatedTree +evaluateTree et fisher hessian theta = EvaluatedTree { + valLoss = ctLoss et theta, + valTheta = theta, + valRows = fromIntegral (ctRows et), + valParams = fromIntegral (U.length theta), + valTree = ctTree et, + valLogParams = logParameters fisher theta, + valLogParamsLattice = logParametersLatt hessian fisher theta, + valVar = ctVar et +} + + +-- log of the parameters complexity +logParameters :: U.Vector Double -> Target -> Double +logParameters fisher theta = -(p / 2) * log 3 + 0.5 * logFisher + logTheta + where + (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (U.toList theta) (U.toList fisher) + +-- same as above but for the Lattice +logParametersLatt :: [[Double]] -> U.Vector Double -> Target -> Double +logParametersLatt hessian fisher theta = 0.5 * p * (1 - log 3) + 0.5 * log detFisher + where + detFisher = det $ map U.fromList hessian + + (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (U.toList theta) (U.toList fisher) + +addIfSignificant (v, f) (acc_v, acc_f, acc_p) + | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1) + | otherwise = (acc_v, acc_f, acc_p) +{-# INLINE addIfSignificant #-} + +isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +{-# INLINE isSignificant #-} + +fixParam :: Int -> Double -> Fix SRTree -> Fix SRTree +fixParam ix val = cata alg + where + alg (Param i) | i == ix = Fix $ Const val + | i > ix = Fix $ Param (i-1) + | otherwise = Fix $ Param i + alg other = Fix other +{-# INLINE fixParam #-} diff --git a/src/Algorithm/SRTree/ConfidenceIntervals.hs b/src/Algorithm/SRTree/ConfidenceIntervals.hs --- a/src/Algorithm/SRTree/ConfidenceIntervals.hs +++ b/src/Algorithm/SRTree/ConfidenceIntervals.hs @@ -1,7 +1,7 @@ {-# language ViewPatterns, ScopedTypeVariables, MultiWayIf, FlexibleContexts #-} ------------------------------------------------------------------------------ +------------------------------------------------------------------------------- -- | --- Module : Algorithm.SRTree.ConfidenceIntervals +-- Module : Algorithm.SRTree.ConfidenceIntervals -- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com @@ -9,283 +9,271 @@ -- Portability : ConstraintKinds -- -- Functions to optimize the parameters of an expression. --- ------------------------------------------------------------------------------ +------------------------------------------------------------------------------- module Algorithm.SRTree.ConfidenceIntervals where -import qualified Data.Massiv.Array as A -import Data.Massiv.Array (Ix2(..), (*.), (!+!), (!*!)) -import Data.Massiv.Array.Numeric ( identityMatrix ) import Statistics.Distribution ( ContDistr(quantile) ) import Statistics.Distribution.StudentT ( studentT ) import Statistics.Distribution.FDistribution ( fDistribution ) +import qualified Data.Vector.Unboxed as U import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Generic as G import Data.SRTree import Data.SRTree.Eval import Data.SRTree.Recursion ( cata ) import Algorithm.SRTree.Likelihoods -import Algorithm.SRTree.Opt - ( minimizeNLL, minimizeNLLWithFixedParam ) +import Algorithm.SRTree.Compile import Data.List ( sortOn, nubBy ) -import Data.Maybe ( fromMaybe ) -import Algorithm.SRTree.NonlinearOpt -import Algorithm.Massiv.Utils +import Data.Maybe ( listToMaybe ) +import Algorithm.SRTree.Utils +import Numeric.Optimization.NLOPT import System.IO.Unsafe ( unsafePerformIO ) -import Control.Monad.Catch ( catch ) +import Control.Monad.Catch ( catch, SomeException ) -import Debug.Trace ( trace, traceShow ) +import Debug.Trace ( trace ) -- | profile likelihood algorithms: Bates (classical), ODE (faster), Constrained (fastest) -- The Constrained approach returns only the endpoints. -data PType = Bates | ODE | Constrained deriving (Show, Read) +data PType = Bates | ODE | Constrained deriving (Show, Read, Eq) -- | Confidence Interval using Laplace approximation or profile likelihood. data CIType = Laplace BasicStats | Profile BasicStats [ProfileT] --- | Basic stats of the data: covariance of parameters, correlation, standard errors -data BasicStats = MkStats { _cov :: SRMatrix - , _corr :: SRMatrix - , _stdErr :: PVector - } deriving (Eq, Show) +-- | Basic stats of the data: covariance of parameters, correlation, standard errors +data BasicStats = MkStats + { _cov :: Columns + , _corr :: Columns + , _stdErr :: Target + } deriving (Eq, Show) -- | a confience interval is composed of the point estimate (`est_`), lower bound (`_lower_`) -- and upper bound (`upper_`) -data CI = CI { est_ :: Double - , lower_ :: Double - , upper_ :: Double - } deriving (Eq, Show, Read) +data CI = CI + { est_ :: Double + , lower_ :: Double + , upper_ :: Double + } deriving (Eq, Show, Read) --- | A profile likelihood is composed of a vector of tau values that traces the likelihood, --- the matrix of thetas for each profile, the local optima, and two splines that converts --- taus to theta and vice-versa. -data ProfileT = ProfileT { _taus :: PVector - , _thetas :: SRMatrix - , _opt :: Double - , _tau2theta :: Double -> Double - , _theta2tau :: Double -> Double - } +-- | A profile likelihood is composed of a vector of tau values that traces the likelihood, +-- the matrix of thetas for each profile, the local optima, and two splines that converts +-- taus to theta and vice-versa. +data ProfileT = ProfileT + { _taus :: Target + , _thetas :: Columns + , _opt :: Double + , _tau2theta :: Double -> Double + , _theta2tau :: Double -> Double + } --- shows the CI with n places +-- shows the CI with n places showCI :: Int -> CI -> String showCI n (CI x l h) = show (rnd l) <> " <= " <> show (rnd x) <> " <= " <> show (rnd h) - where - rnd = (/10^n) . (fromIntegral . round) . (*10^n) + where rnd = (/10^n) . (fromIntegral . round) . (*10^n) + printCI :: Int -> CI -> IO () printCI n = putStrLn . showCI n --- | Calculates the confidence interval of the parameters using +-- | Calculates the confidence interval of the parameters using -- Laplace approximation or Profile likelihood -paramCI :: CIType -> Int -> PVector -> Double -> [CI] -paramCI (Laplace stats) nSamples theta alpha = zipWith3 CI (A.toList theta) lows highs +paramCI :: CIType -> Int -> Target -> Double -> [CI] +paramCI (Laplace stats) nSamples theta alpha = zipWith3 CI (U.toList theta) lows highs where - -- the Laplace approximation is theta +/- t(1-alpha/2) * standard error - (A.Sz k) = A.size theta - t = quantile (studentT . fromIntegral $ nSamples - k) (1 - alpha / 2.0) - stdErr = _stdErr stats - lows = A.toList $ A.zipWith (-) theta $ A.map (*t) stdErr - highs = A.toList $ A.zipWith (+) theta $ A.map (*t) stdErr + -- the Laplace approximation is theta +/- t(1-alpha/2) * standard error + k = U.length theta + t = quantile (studentT . fromIntegral $ nSamples - k) (1 - alpha / 2.0) + stdErr = _stdErr stats + lows = U.toList $ U.zipWith (-) theta $ U.map (*t) stdErr + highs = U.toList $ U.zipWith (+) theta $ U.map (*t) stdErr paramCI (Profile stats profiles) nSamples _ alpha = zipWith3 CI theta lows highs where -- for the profile likelihood we use the square root of the F-distribution with (1-alpha) - k = length theta - t = sqrt $ quantile (fDistribution k (fromIntegral $ nSamples - k)) (1 - alpha) - stdErr = _stdErr stats - lows = map (`_tau2theta` (-t)) profiles - highs = map (`_tau2theta` t) profiles - theta = map _opt profiles + k = length theta + t = sqrt $ quantile (fDistribution k (fromIntegral $ nSamples - k)) (1 - alpha) + stdErr = _stdErr stats + lows = map (`_tau2theta` (-t)) profiles + highs = map (`_tau2theta` t) profiles + theta = map _opt profiles --- | calculates the prediction confidence interval using Laplace approximation or profile likelihood. --- -predictionCI :: CIType -> Distribution -> (SRMatrix -> PVector) -> (SRMatrix -> [PVector]) -> (CI -> PVector -> Fix SRTree -> (Double -> Double, Double)) -> SRMatrix -> Fix SRTree -> PVector -> Double -> [CI] -> [CI] +-- | calculates the prediction confidence interval using Laplace approximation or profile likelihood. +-- predictionCI +predictionCI :: CIType -> Distribution -> (Columns -> Target) -> (Columns -> [Target]) -> (CI -> Target -> Fix SRTree -> (Double -> Double, Double)) -> Columns -> Fix SRTree -> Target -> Double -> [CI] -> [CI] predictionCI (Laplace stats) _ predFun jacFun _ xss tree theta alpha _ = zipWith3 CI yhat lows highs where - yhat = A.toList $ predFun xss - jac' :: A.Matrix A.S Double - jac' = A.fromLists' compMode $ map A.toList $ jacFun xss - jac :: [PVector] - jac = A.toList $ A.outerSlices $ A.computeAs A.S $ A.transpose jac' - n = length yhat - (A.Sz k) = A.size theta - t = quantile (studentT . fromIntegral $ n - k) (1 - alpha / 2.0) - covs = A.toList $ A.outerSlices $ _cov stats - lows = zipWith (-) yhat $ map (*t) resStdErr - highs = zipWith (+) yhat $ map (*t) resStdErr + yhat = U.toList $ predFun xss + jac' = jacFun xss + k = U.length theta + n = length yhat + t = quantile (studentT . fromIntegral $ n - k) (1 - alpha / 2.0) - getResStdError row = sqrt $ (A.!.!) row $ A.fromList compMode $ map (row A.!.!) covs - resStdErr = map getResStdError jac + covMat = toRowMajor (_cov stats) + nCov = k - 1 -predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat xss' -- $ take 10 xss' - where - yhat = A.toList $ predFun xss - theta' = A.toStorableVector theta + lows = zipWith (-) yhat $ map (*t) resStdErr + highs = zipWith (+) yhat $ map (*t) resStdErr - t = sqrt $ quantile (fDistribution k (fromIntegral $ n - k)) (1 - alpha) - (A.Sz k) = A.size theta - n = length yhat + getResStdError row = + sqrt $ U.sum $ U.generate nCov $ \i -> + (row U.! i) * U.sum (U.zipWith (*) row (U.slice (i * k) nCov covMat)) + resStdErr = map (getResStdError . U.slice 0 nCov) (getRows jac') - theta0 = calcTheta0 dist tree - xss' = A.toList $ A.outerSlices xss +predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat xss' + where + yhat = U.toList $ predFun xss + k = U.length theta + n = length yhat + t = sqrt $ quantile (fDistribution k (fromIntegral $ n - k)) (1 - alpha) - f estPI yh xs = - let t' = replaceParam0 tree $ evalVar xs theta0 - (spline, yh') = profFun estPI (A.fromStorableVector compMode (theta' VS.// [(0, yh)])) t' - in CI yh' (spline (-t)) (spline t) + theta0 = calcTheta0 dist tree + xss' = getRows xss --- inverse function of the distributions + f estPI yh xs = let + t' = replaceParam0 tree $ evalVar xs theta0 + (spline, yh') = profFun estPI (theta U.// [(0, yh)]) t' + in CI yh' (spline (-t)) (spline t) + +-- inverse function of the distributions inverseDist :: Floating p => Distribution -> p -> p -inverseDist MSE y = y -inverseDist Gaussian y = y +inverseDist Gaussian y = y inverseDist Bernoulli y = log (y/(1-y)) -inverseDist Poisson y = log y +inverseDist Poisson y = log y +inverseDist _ y = y --- rewrite the tree by fixing theta 0 to optimal value +-- rewrite the tree by fixing theta 0 to optimal value replaceParam0 :: Fix SRTree -> Fix SRTree -> Fix SRTree replaceParam0 tree t0 = cata alg tree where - alg (Var ix) = Fix $ Var ix - alg (Param 0) = t0 - alg (Param ix) = Fix $ Param ix - alg (Const c) = Fix $ Const c - alg (Uni g t) = Fix $ Uni g t + alg (Var ix) = Fix $ Var ix + alg (Param 0) = t0 + alg (Param ix) = Fix $ Param ix + alg (Const c) = Fix $ Const c + alg (Y ix) = Fix $ Y ix + alg (Uni g t) = Fix $ Uni g t alg (Bin op l r) = Fix $ Bin op l r -evalVar :: PVector -> Fix SRTree -> Fix SRTree +evalVar :: Target -> Fix SRTree -> Fix SRTree evalVar xs = cata alg where - alg (Var ix) = Fix $ Const (xs A.! ix) - alg (Param ix) = Fix $ Param ix - alg (Const c) = Fix $ Const c - alg (Uni g t) = Fix $ Uni g t + alg (Var ix) = Fix $ Const (xs U.! ix) + alg (Param ix) = Fix $ Param ix + alg (Const c) = Fix $ Const c + alg (Y ix) = Fix $ Y ix + alg (Uni g t) = Fix $ Uni g t alg (Bin op l r) = Fix $ Bin op l r calcTheta0 :: Distribution -> Fix SRTree -> Fix SRTree calcTheta0 dist tree = case cata alg tree of - Left g -> g $ inverseDist dist (Fix $ Param 0) - Right _ -> error "No theta0?" + Left g -> g $ inverseDist dist (Fix $ Param 0) + Right _ -> error "No theta0?" where - alg (Var ix) = Right $ Fix $ Var ix - alg (Param 0) = Left id - alg (Param ix) = Right $ Fix $ Param ix - alg (Const c) = Right $ Fix $ Const c - alg (Uni g t) = case t of - Left f -> Left $ f . evalInverse g - Right v -> Right $ evalFun g v + alg (Var ix) = Right $ Fix $ Var ix + alg (Param 0) = Left id + alg (Param ix) = Right $ Fix $ Param ix + alg (Const c) = Right $ Fix $ Const c + alg (Y ix) = Right $ Fix $ Y ix + alg (Uni g t) = case t of + Left f -> Left $ f . evalInverse g + Right v -> Right $ evalFun g v alg (Bin op l r) = case l of - Left f -> case r of - Left _ -> error "This shouldn't happen!" - Right v -> Left $ f . invright op v - Right vl -> case r of - Left g -> Left $ g . invleft op vl - Right vr -> Right $ evalOp op vl vr + Left f -> case r of + Left _ -> error "This shouldn't happen!" + Right v -> Left $ f . invright op v + Right vl -> case r of + Left g -> Left $ g . invleft op vl + Right vr -> Right $ evalOp op vl vr --- calculate the profile likelihood of every parameter -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 []) +-- calculate the profile likelihood of every parameter +getAllProfiles :: PType -> EvalTree -> Target -> Target -> [CI] -> Double -> [ProfileT] +getAllProfiles ptype et theta stdErr estCIs alpha = getAll 0 [] where - (A.Sz k) = A.size theta - (A.Sz n) = A.size ys - tau_max = sqrt $ quantile (fDistribution k (n - k)) (1 - 0.01) - tau_max' = sqrt $ quantile (fDistribution k (n - k)) (1 - alpha) + k = U.length theta + n = ctRows et + tau_max = sqrt $ quantile (fDistribution k (n - k)) (1 - 0.01) + tau_max' = sqrt $ quantile (fDistribution k (n - k)) (1 - alpha) profFun ix = case ptype of - 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 + Bates -> getProfile et theta (stdErr U.! ix) tau_max ix + ODE -> getProfileODE et theta (stdErr U.! ix) (estCIs !! ix) tau_max ix + Constrained -> getProfileCnstr et theta (stdErr U.! ix) tau_max' ix getAll ix acc | ix == k = acc + | ix == k-1 && ptype == Constrained && ctDist et == Gaussian = case getProfileODE et theta (stdErr U.! ix) (estCIs !! ix) tau_max ix of + Left t -> getAllProfiles ptype et t stdErr estCIs alpha + Right p -> getAll (ix + 1) (acc <> [p]) | otherwise = case profFun ix of - Left t -> getAllProfiles ptype dist mYerr xss ys tree t stdErr estCIs alpha - Right p -> getAll (ix + 1) (p : acc) + Left t -> getAllProfiles ptype et t stdErr estCIs alpha + Right p -> getAll (ix + 1) (acc <> [p]) --- calculates the profile likelihood of a single parameter -getProfile :: Distribution - -> Maybe PVector - -> SRMatrix - -> PVector - -> Fix SRTree - -> PVector - -> Double - -> Double - -> Int - -> Either PVector ProfileT -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) +-- calculates the profile likelihood of a single parameter +getProfile :: EvalTree -> Target -> Double -> Double -> Int -> Either Target ProfileT +getProfile et theta stdErr_i tau_max ix + | stdErr_i == 0.0 = pure $ ProfileT (U.fromList [-tau_max, tau_max]) [theta, theta] (theta U.! ix) (const (theta U.! ix)) (const tau_max) | otherwise = do negDelta <- go kmax (-stdErr_i / 8) 0 1 mempty posDelta <- go kmax (stdErr_i / 8) 0 1 p0 - let (A.fromList compMode -> taus, A.fromLists' compMode. map A.toList -> thetas) = negDelta <> posDelta - (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix + let (taus', thetas') = negDelta <> posDelta + taus = U.fromList taus' + thetas = thetas' + (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix pure $ ProfileT taus thetas optTh tau2theta theta2tau where - theta' = A.toList theta p0 = ([0], [theta_opt]) kmax = 300 - 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 mYerr 100 xss ys tree ix + nll_opt = ctNLL et theta_opt + theta_opt = ctOptimizer et theta + optTh = theta_opt U.! ix + minimizer = ctOptimizerFixed et ix - -- after k iterations, interpolates to the endpoint go 0 delta _ _ acc = Right acc go k delta t inv_slope acc@(taus, thetas) - | isNaN inv_slope = Right acc -- stop since we cannot move forward on discontinuity - | nll_cond < nll_opt = Left theta_t -- found a better optima - | abs tau > tau_max = Right acc' -- we reached the endpoint + | isNaN inv_slope = Right acc + | nll_cond < nll_opt = Left theta_t + | abs tau > tau_max = Right acc' + | otherwise = go (k-1) delta (t + inv_slope) inv_slope' acc' where - t_delta = (theta_opt A.! ix) + delta * (t + inv_slope) + t_delta = (theta_opt U.! 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 mYerr xss ys tree theta_t) A.! ix - zvs = snd $ gradNLL dist mYerr xss ys tree theta_t + (nll_cond, grad) = ctGradNLL et theta_t + zv = grad U.! ix inv_slope' = min 4.0 . max 0.0625 . abs $ (tau / (stdErr_i * zv)) - nll_cond = nll dist mYerr xss ys tree theta_t - acc' = if nll_cond == nll_opt || ( (not.null) taus && tau == head taus ) || isNaN tau + tau = signum delta * sqrt (2*nll_cond - 2*nll_opt) + acc' = if nll_cond == nll_opt || maybe False (tau ==) (listToMaybe taus) || isNaN tau then acc else (tau:taus, theta_t:thetas) - tau = signum delta * sqrt (2*nll_cond - 2*nll_opt) -- 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 PVector - -> SRMatrix - -> PVector - -> Fix SRTree - -> PVector - -> Double -> Double - -> Int - -> Either PVector ProfileT -getProfileCnstr dist mYerr xss ys tree theta stdErr_i tau_max ix +getProfileCnstr :: EvalTree -> Target -> Double -> Double -> Int -> Either Target ProfileT +getProfileCnstr et 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 - taus = A.fromList compMode [-tau_max, tau_max] - theta' = A.toList theta - thetas = A.fromLists' compMode [theta', theta'] - theta_i = theta A.! ix - getPoint = getEndPoint dist mYerr xss ys tree theta tau_max ix + taus = U.fromList [-tau_max, tau_max] + thetas = [theta, theta] + theta_i = theta U.! ix + getPoint = getEndPoint et theta tau_max ix leftPt = getPoint True rightPt = getPoint False tau2theta tau = if tau < 0 then leftPt else rightPt -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 +getEndPoint :: EvalTree -> Target -> Double -> Int -> Bool -> Double +getEndPoint et theta tau_max ix isLeft = + case minimizeAugLag problem (G.convert theta_opt) of Right sol -> solutionParams sol VS.! ix - Left e -> traceShow e $ theta_opt A.! ix + Left _ -> theta_opt U.! ix where - (A.Sz1 n) = A.size theta + n = U.length theta - (theta_opt, _, _) = minimizeNLL dist mYerr 100 xss ys tree theta - nll_opt = nll dist mYerr xss ys tree theta_opt + theta_opt = ctOptimizer et theta + nll_opt = ctNLL et theta_opt loss_crit = nll_opt + tau_max - loss = subtract loss_crit . nll dist mYerr xss ys tree . A.fromStorableVector compMode + loss = subtract loss_crit . ctNLL et . G.convert obj = (if isLeft then id else negate) . (VS.! ix) - stop = ObjectiveRelativeTolerance 1e-4 :| [] + stop = ObjectiveRelativeTolerance 1e-4 :| [MaximumEvaluations 1000] localAlg = NELDERMEAD obj [] Nothing local = LocalProblem (fromIntegral n) stop localAlg constraint = InequalityConstraint (Scalar loss) 1e-6 @@ -296,153 +284,139 @@ -- Based on -- 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 PVector - -> SRMatrix - -> PVector - -> Fix SRTree - -> PVector - -> Double - -> CI - -> Double - -> Int - -> Either PVector ProfileT -getProfileODE dist mYerr xss ys tree theta stdErr_i estCI tau_max ix +getProfileODE :: EvalTree -> Target -> Double -> CI -> Double -> Int -> Either Target ProfileT +getProfileODE et 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 + | otherwise = let (taus', thetas') = solLeft <> ([0], [theta_opt]) <> solRight + taus = U.fromList taus' + thetas = thetas' (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 = (\(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 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 mYerr xss ys tree + dflt = ProfileT (U.fromList [-tau_max, tau_max]) [theta, theta] (theta U.! ix) (const (theta U.! ix)) (const tau_max) + theta_opt = ctOptimizer et theta + grader = snd . ctGradNLL et + nll_opt = ctNLL et theta_opt + optTh = theta_opt U.! ix + p = U.length theta + p' = p + 1 odeFun gamma _ u = let grad = grader u - w = hessianNLL dist mYerr xss ys tree u - m = A.makeArray compMode (A.Sz (p' :. p')) - (\ (i :. j) -> if | i

w A.! (i :. j) - | i==ix -> 1 - | j==ix -> 1 - | otherwise -> 0 - ) - - v = A.computeAs A.S $ A.snoc (A.map (*(-gamma)) grad) 1 + w = ctHessianNLL et u + m = [ U.generate p' (\i -> + if i < p && j < p then (w !! j) U.! i + else if i == ix || j == ix then 1 + else 0 + ) + | j <- [0 .. p'-1] ] + v = U.snoc (U.map (*(-gamma)) grad) 1 dotTheta = unsafePerformIO $ luSolve m v - in A.fromStorableVector compMode $ VS.init $ A.toStorableVector dotTheta - tsHi = linSpace 50 (optTh, upper_ estCI) - tsLo = linSpace 50 (optTh, lower_ estCI) + in U.init dotTheta + + minRange = max (abs (upper_ estCI - optTh)) (abs (lower_ estCI - optTh)) + scanRange = max minRange (tau_max * abs stdErr_i) + nPts = max 50 (min 100 (ceiling (scanRange / minRange * 49) + 1)) + tsHi = linSpace nPts (optTh, optTh + scanRange) + tsLo = linSpace nPts (optTh, optTh - scanRange) scanOn sig = foldMap (calcTau sig) . f . scanl (rk (odeFun sig)) (optTh, theta_opt) 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 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]) + calcTau s t = let nll_i = ctNLL et (snd t) + z = signum ((snd t U.! ix) - optTh) * sqrt (2 * nll_i - 2 * nll_opt) + in if z == 0 || isNaN z then ([], []) else ([z], [snd t]) -rk :: (Double -> PVector -> PVector) -> (Double, PVector) -> Double -> (Double, PVector) -rk f (t, y) t' = (t', y !+! ((1.0/6.0) *. h' !*! (k1 !+! (2.0 *. k2) !+! (2.0 *. k3) !+! k4))) +rk :: (Double -> Target -> Target) -> (Double, Target) -> Double -> (Double, Target) +rk f (t, y) t' = (t', U.zipWith5 (\y0 k1 k2 k3 k4 -> y0 + h/6 * (k1 + 2*k2 + 2*k3 + k4)) y k1 k2 k3 k4) where h = t' - t - h', k1, k2, k3, k4 :: PVector - h' = A.replicate compMode (A.size y) h k1 = f t y - k2 = f (t + 0.5*h) (A.computeAs A.S $ A.zipWith3 (g 0.5) y h' k1) -- (y !+! 0.5*.h' A.!*! k1) - k3 = f (t + 0.5*h) (A.computeAs A.S $ A.zipWith3 (g 0.5) y h' k2) -- (y !+! 0.5*.h' A.!*! k2) - k4 = f (t + 1.0*h) (A.computeAs A.S $ A.zipWith3 (g 1.0) y h' k3) -- (y !+! 1.0*.h'!*!k3) - g a yi hi ki = yi + a * hi * ki + k2 = f (t + 0.5*h) (U.zipWith (\y0 k -> y0 + 0.5*h*k) y k1) + k3 = f (t + 0.5*h) (U.zipWith (\y0 k -> y0 + 0.5*h*k) y k2) + k4 = f (t + 1.0*h) (U.zipWith (\y0 k -> y0 + 1.0*h*k) y k3) {-# INLINE rk #-} --- tau0, tau1 theta0, thetaX = tau1 theta0 / tau0 -getStatsFromModel :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> BasicStats +-- tau0, tau1 theta0, thetaX = tau1 theta0 / tau0 +getStatsFromModel :: Distribution -> Maybe Target -> Columns -> Target -> Fix SRTree -> Target -> 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 + k = U.length theta + n = U.length ys nParams = fromIntegral k - ssr = sse xss ys tree theta - ident = A.computeAs A.S $ identityMatrix nParams + ident = fromRowMajor k k (U.generate (k * k) (\ix -> let (i, j) = ix `divMod` k in if i == j then 1.0 else 0.0)) - -- only for gaussian - sErr = sqrt $ ssr / fromIntegral (n - k) + hess = hessianNLL dist mYerr 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 - cov = unsafePerformIO $ catch (invChol hess) fexcept + fexcept :: SomeException -> IO Columns + fexcept e = trace ("cov NegDef" <> show (toRowMajor hess)) $ pure ident - stdErr = A.makeArray compMode (A.Sz1 k) (\ix -> sqrt $ cov A.! (ix :. ix)) + cov = unsafePerformIO $ catch (invChol hess) fexcept + + covMat = toRowMajor cov + stdErr = U.generate k (\ix -> sqrt $ covMat U.! (ix * k + ix)) + stdErrSq = case outer stdErr stdErr of - Left _ -> error "stdErr size mismatch?" - Right v -> v + Right v -> v + Left _ -> [] - corr = A.computeAs A.S $ A.zipWith (/) cov stdErrSq + stdErrSqMat = toRowMajor stdErrSq + corr = fromRowMajor k k $ U.generate (k * k) (\ix -> covMat U.! ix / stdErrSqMat U.! ix) -- Create splines for profile-t -createSplines :: PVector -> SRMatrix -> Double -> Double -> Int -> (Double -> Double, Double -> Double) +createSplines :: Target -> Columns -> Double -> Double -> Int -> (Double -> Double, Double -> Double) createSplines taus thetas se tau_max ix - | n < 2 = (genSplineFun [(-tau_max, -se), (tau_max, se)], genSplineFun [(-se, 0), (se, 1)]) + | n < 2 = (genSplineFun [(-tau_max, -se), (tau_max, se)], genSplineFun [(-se, 0), (se, 1)]) | otherwise = (tau2theta, theta2tau) where - (A.Sz n) = A.size taus - cols = getCol ix thetas + n = U.length taus + cols = getCol ix thetas nubOnFirst = nubBy (\x y -> fst x == fst y) - tau2theta = genSplineFun $ nubOnFirst $ sortOnFirst taus cols - theta2tau = genSplineFun $ nubOnFirst $ sortOnFirst cols taus + tau2theta = genSplineFun $ nubOnFirst $ sortOnFirst taus cols + theta2tau = genSplineFun $ nubOnFirst $ sortOnFirst cols taus -getCol :: Int -> SRMatrix -> PVector -getCol ix mtx = getCols mtx A.! ix +getCol :: Int -> Columns -> Target +getCol ix mtx = U.generate (length mtx) (\j -> (mtx !! j) U.! ix) {-# inline getCol #-} -sortOnFirst :: PVector -> PVector -> [(Double, Double)] -sortOnFirst xs ys = sortOn fst $ zip (A.toList xs) (A.toList ys) +sortOnFirst :: Target -> Target -> [(Double, Double)] +sortOnFirst xs ys = sortOn fst $ zip (U.toList xs) (U.toList ys) {-# inline sortOnFirst #-} -splinesSketches :: Double -> PVector -> PVector -> (Double -> Double) -> (Double -> Double) -splinesSketches tauScale (A.toList -> tau) (A.toList -> theta) theta2tau +splinesSketches :: Double -> Target -> Target -> (Double -> Double) -> (Double -> Double) +splinesSketches tauScale (U.toList -> tau) (U.toList -> theta) theta2tau | length tau < 2 = id - | otherwise = genSplineFun gpq + | otherwise = genSplineFun gpq where - gpq = sortOn fst [(x, acos y') | (x, y) <- zip tau theta - , let y' = theta2tau y / tauScale - , abs y' < 1 ] + gpq = sortOn fst [ (x, acos y') | (x, y) <- zip tau theta, let y' = theta2tau y / tauScale, abs y' < 1 ] approximateContour :: Int -> Int -> [ProfileT] -> Int -> Int -> Double -> [(Double, Double)] approximateContour nParams nPoints profs ix1 ix2 alpha = go 0 where - -- get the info for ix1 and ix2 - (prof1, prof2) = (profs !! ix1, profs !! ix2) + (prof1, prof2) = (profs !! ix1, profs !! ix2) (tau2theta1, theta2tau1) = (_tau2theta prof1, _theta2tau prof1) (tau2theta2, theta2tau2) = (_tau2theta prof2, _theta2tau prof2) - -- calculate the spline for A-D - tauScale = sqrt (fromIntegral nParams * quantile (fDistribution nParams (nPoints - nParams)) (1 - alpha)) - splineG1 = splinesSketches tauScale (_taus prof1) (getCol ix2 (_thetas prof1)) theta2tau2 - splineG2 = splinesSketches tauScale (_taus prof2) (getCol ix1 (_thetas prof2)) theta2tau1 - angles = [ (0, splineG1 1), (splineG2 1, 0), (pi, splineG1 (-1)), (splineG2 (-1), pi) ] - splineAD = genSplineFun points + tauScale = sqrt (fromIntegral nParams * quantile (fDistribution nParams (fromIntegral nPoints - fromIntegral nParams)) (1 - alpha)) + splineG1 = splinesSketches tauScale (_taus prof2) (getCol ix1 (_thetas prof2)) theta2tau1 + splineG2 = splinesSketches tauScale (_taus prof1) (getCol ix2 (_thetas prof1)) theta2tau2 + angles = [ (0, splineG2 1), (splineG1 1, 0), (pi, splineG2 (-1)), (splineG1 (-1), pi) ] applyIfNeg (x, y) = if y < 0 then (-x, -y) else (x ,y) - points = sortOn fst - $ [applyIfNeg ((x+y)/2, x - y) | (x, y) <- angles] - <> (\(x,y) -> [(x + 2*pi, y)]) (head points) + points' = [applyIfNeg ((x+y)/2, x - y) | (x, y) <- angles] + points = sortOn fst $ points' <> maybe [] (\(x,y) -> [(x + 2*pi, y)]) (listToMaybe points') + splineAD = genSplineFun points - -- generate the points of the curve + fmod a b = a - b * fromIntegral (truncate (a / b)) + + tot = 100 go 100 = [] - go ix = (p, q) : go (ix+1) + go ix = (p, q) : go (ix+1) where - ai = ix * 2 * pi / 99 - pi + ai = fromIntegral ix * 2 * pi / 99 - pi di = splineAD ai - taup = cos (ai + di / 2) * tauScale - tauq = cos (ai - di / 2) * tauScale - p = tau2theta1 taup - q = tau2theta2 tauq + t1i = tauScale * cos (ai + di) + t2i = tauScale * cos (ai - di) + p = tau2theta1 t1i + q = tau2theta2 t2i + diff --git a/src/Algorithm/SRTree/Likelihoods.hs b/src/Algorithm/SRTree/Likelihoods.hs --- a/src/Algorithm/SRTree/Likelihoods.hs +++ b/src/Algorithm/SRTree/Likelihoods.hs @@ -1,9 +1,11 @@ {-# LANGUAGE ViewPatterns #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE UnboxedTuples #-} ----------------------------------------------------------------------------- -- | --- Module : Algorithm.SRTree.Likelihoods +-- Module : AlgorithV.SRTree.Likelihoods -- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com @@ -15,94 +17,87 @@ ----------------------------------------------------------------------------- module Algorithm.SRTree.Likelihoods ( Distribution (..) - , PVector - , SRMatrix - , sse - , mse - , rmse - , r2 - , nll - , predict - , buildNLL - , buildNLLEGraph - , gradNLL - , gradNLLArr - , gradNLLGraph - , gradNLLEGraph + , Loss (..) + , readLoss + , Target + , Columns + , buildDistLoss + , buildLoss + , buildPredictor , fisherNLL , getSErr , hessianNLL - , tree2arr ) where -import Algorithm.SRTree.AD ( reverseModeArr, reverseModeGraph, reverseModeEGraph ) -import Data.Massiv.Array hiding (all, map, read, replicate, tail, take, zip) -import qualified Data.Massiv.Array as M -import qualified Data.Massiv.Array.Mutable as Mut -import Data.Maybe (fromMaybe) import Data.SRTree import Data.SRTree.Recursion ( cata, accu ) -import Data.SRTree.Derivative (deriveByParam, deriveByVar, derivative) +import Data.SRTree.Derivative (deriveByParam, deriveByVar, derivative, derivOp) import Data.SRTree.Eval import qualified Data.IntMap.Strict as IntMap import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM + import GHC.IO (unsafePerformIO) import Data.Maybe +import Text.Read (readMaybe) +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Unboxed.Mutable as VM +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (forConcurrently) + import Debug.Trace import Data.SRTree.Print -import Algorithm.EqSat.Egraph -import Algorithm.EqSat.Simplify -import Algorithm.EqSat.Build import Control.Monad.State.Strict import Control.Monad.Identity import Data.SRTree.Print +import qualified Data.Vector.Generic as G --- | Supported distributions for negative log-likelihood --- MSE refers to mean squared error --- HGaussian is Gaussian with heteroscedasticity, where the error should be provided -data Distribution = MSE | Gaussian | HGaussian | Bernoulli | Poisson | ROXY | LOG10 +-- | Supported distributions for negative log-likelihood. +-- | HGaussian is Gaussian with heteroscedasticity, where the error should be provided. +data Distribution = Gaussian | HGaussian | Bernoulli | Poisson | ROXY | LeastSquares deriving (Show, Read, Enum, Bounded, Eq) --- | Sum-of-square errors or Sum-of-square residues -sse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double -sse xss ys 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) +-- | Loss functions used to build the per-row optimization objective (see +-- 'buildLoss'), to be used by e.g. "Algorithm.SRTree.Opt". 'NLL' wraps a +-- 'Distribution' to use its negative log-likelihood as the loss -- +-- including the plain \'MSE\' and \'LOG10\' losses, reached via @NLL MSE@ +-- and @NLL LOG10@ respectively (kept on 'Distribution', rather than +-- duplicated here, since Haskell does not allow two data constructors +-- with the same name -- 'MSE' and 'LOG10' -- to coexist in the same +-- module). +data Loss = MSE | LOG10 | MAE | MAPE | Pinball Double | NLL Distribution + deriving (Show, Read, Eq) -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)) +instance Enum Loss where + fromEnum MSE = 0 + fromEnum LOG10 = 1 + fromEnum MAE = 2 + fromEnum MAPE = 3 + fromEnum (Pinball _) = 4 + fromEnum (NLL dist) = 5 + fromEnum dist --- | Total Sum-of-squares -sseTot :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double -sseTot xss ys tree theta = err - where - (Sz m) = M.size ys - cmp = getComp xss - ym = M.sum ys / fromIntegral m - err = M.sum $ (M.map (subtract ym) ys) ^ (2 :: Int) - --- | Mean squared errors -mse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double -mse xss ys tree theta = let (Sz m) = M.size ys in sse xss ys tree theta / fromIntegral m + toEnum 0 = MSE + toEnum 1 = LOG10 + toEnum 2 = MAE + toEnum 3 = MAPE + toEnum 4 = Pinball 0.95 + toEnum x | x >= 5 = NLL (toEnum (x-5)) --- | Root of the mean squared errors -rmse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double -rmse xss ys tree = sqrt . mse xss ys tree +instance Bounded Loss where + minBound = MSE + maxBound = NLL ROXY --- | Coefficient of determination -r2 :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double -r2 xss ys tree theta = 1 - sse xss ys tree theta / sseTot xss ys tree theta +-- | Parse a loss from its CLI string. Accepts both the direct 'Loss' +-- names ('MSE', 'LOG10', 'MAE', 'MAPE', @Pinball tau@) and the bare +-- 'Distribution' names ('Gaussian', 'HGaussian', 'Bernoulli', 'Poisson', +-- 'ROXY', 'LeastSquares'), which are wrapped in 'NLL'. +readLoss :: String -> Maybe Loss +readLoss s = case readMaybe s of + Just l -> Just l + Nothing -> NLL <$> (readMaybe s :: Maybe Distribution) -- | logistic function logistic :: Floating a => a -> a @@ -118,518 +113,217 @@ {-# inline getSErr #-} -- negation of the sum of values in a vector -negSum :: PVector -> Double -negSum = negate . M.sum +negSum :: Target -> Double +negSum = negate . V.sum {-# inline negSum #-} --- | Negative log-likelihood -nll :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> Double - --- | Mean Squared error (not a distribution) -nll MSE _ xss ys t theta = mse xss ys t theta - -nll LOG10 _ xss ys t theta = M.sum $ (M.map (logBase 10) $ (f (delay ys) / f yhat)) ^ (2 :: Int) - where - yhat = evalTree xss theta t - (Sz m) = M.size ys - f :: Array D Ix1 Double -> Array D Ix1 Double - f z = (z + M.map (\zi -> sqrt (zi^2 + 1e-10)) z) - -- log ys - log y = log (ys/y) - --- | Gaussian distribution, theta must contain an additional parameter corresponding --- to variance. -nll Gaussian mYerr xss ys t theta - | nParams == (p'-1) = error "For Gaussian distribution theta must contain the variance as its last value." - | otherwise = 0.5*(sse xss ys t theta / s + m*log (2*pi*s)) - where - s = sqrt $ mse xss ys t theta -- theta M.! (p' - 1) - (Sz m') = M.size ys - (Sz p') = M.size theta - nParams = countParamsUniq 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 = M.sum $ (M.map (1-) (delay ys)) * yhat + log (M.map (1+) $ exp (M.map negate yhat)) - where - (Sz m) = M.size ys - yhat = evalTree xss theta tree - notValid = M.any (\x -> x /= 0 && x /= 1) - -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 - where - ys' = delay ys - yhat = evalTree xss theta tree - notValid = M.any (<0) - -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 = countParamsUniq tree - - x0 = xss 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 +checkAssumptions :: Distribution -> Maybe Target -> Target -> Bool +checkAssumptions Gaussian _ _ = True +checkAssumptions HGaussian (Just yErr) _ = True +checkAssumptions HGaussian Nothing _ = False +checkAssumptions Bernoulli _ ys = V.all (\x -> x /= 0 && x /= 1) ys +checkAssumptions Poisson _ ys = V.all (>0) ys +checkAssumptions LeastSquares _ _ = True +checkAssumptions ROXY mYerr ys = isJust mYerr -- WARNING: pass tree with parameters -- TODO: handle error similar to ROXY -buildNLL MSE m tree = ((tree - var (-1)) ** 2) / constv m -buildNLL LOG10 m tree = (((log (y / tree')) / log 10) ** 2) / constv m - where - tree' = (tree + sqrt(tree^2 + 1e-10)) - y = (var (-1) + sqrt(var (-1) ^ 2 + 1e-10)) -buildNLL Gaussian m tree = (square(tree - var (-1)) / square (param p)) + log ((square (param p))) +-- | Builds the per-row negative log-likelihood expression for a given +-- 'Distribution', to be summed across rows (e.g. by +-- 'Algorithm.SRTree.AD.evalGradMulti') and differentiated by automatic +-- differentiation. The special variable index @-1@ refers to the target +-- ('ys') and @-2@ to the target's measurement error ('yErr'), following +-- the convention used by "Algorithm.SRTree.AD". +-- +-- 'buildLoss' delegates to this function for the @'NLL' dist@ loss. +buildDistLoss :: Distribution -> Double -> Fix SRTree -> Fix SRTree +buildDistLoss Gaussian m tree = (square(tree - var (-1)) * (e (negate (param p)))) + (((param p))) where square = Fix . Uni Square - p = countParamsUniq tree -buildNLL HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2)) -buildNLL Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree -buildNLL Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree -buildNLL ROXY m tree = neglogP + e = Fix. Uni Exp + p = countParamsUniq tree +buildDistLoss HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2)) +buildDistLoss Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree +buildDistLoss Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree +buildDistLoss LeastSquares m tree = ((tree - var (-1)) ** 2) / constv m +buildDistLoss ROXY m tree = neglogP where - p = countParamsUniq 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 + p = countParamsUniq 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_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) + 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 -buildNLLEGraph MSE m egraph root = runIdentity $ addToEg `runStateT` egraph - where - addToEg :: EGraphST Identity EClassId - addToEg = do v <- add myCost (Var (-1)) - c1 <- add myCost (Const 2) - c2 <- add myCost (Const m) - x <- add myCost (Bin Sub root v) - y <- add myCost (Bin Power x c1) - add myCost (Bin Div y c2) -buildNLLEGraph LOG10 m egraph root = runIdentity $ addToEg `runStateT` egraph - where - addToEg :: EGraphST Identity EClassId - addToEg = do v <- add myCost (Var (-1)) - c1 <- add myCost (Const 2) - c2 <- add myCost (Const m) - c3 <- add myCost (Const 10) - c4 <- add myCost (Const 1e-10) - -- log (x + sqrt (x^2 + 1)) / log 10 - log10 <- add myCost (Uni Log c3) - t2 <- add myCost (Uni Square root) - t2p1 <- add myCost (Bin Add t2 c4) - sqt <- add myCost (Uni Sqrt t2p1) - tpt <- add myCost (Bin Add root sqt) - - -- same with y - y2 <- add myCost (Uni Square v) - y2p1 <- add myCost (Bin Add y2 c4) - sqy <- add myCost (Uni Sqrt y2p1) - ypy <- add myCost (Bin Add v sqy) - - tptypy <- add myCost (Bin Div ypy tpt) - - logy <- add myCost (Uni Log tptypy) - log10y <- add myCost (Bin Div logy log10) - - --x <- add myCost (Bin Sub log10t v) - y <- add myCost (Bin Power tptypy c1) - add myCost (Bin Div y c2) - -buildNLLEGraph Gaussian m egraph root = runIdentity (addToEg `runStateT` egraph) - where - p = countParamsUniqEg egraph root - addToEg :: EGraphST Identity EClassId - addToEg = do v <- add myCost (Var (-1)) - p <- add myCost (Param p) - sp <- add myCost (Uni Square p) - lsp <- add myCost (Uni Log sp) - d <- add myCost (Bin Sub root v) - sd <- add myCost (Uni Square d) - x <- add myCost (Bin Div sd sp) - add myCost (Bin Add x lsp) - -buildNLLEGraph HGaussian m egraph root = runIdentity $ addToEg `runStateT` egraph - where - addToEg :: EGraphST Identity EClassId - addToEg = do v1 <- add myCost (Var (-1)) - v2 <- add myCost (Var (-2)) - c1 <- add myCost (Const (2*pi)) - c2 <- add myCost (Const m) - x <- add myCost (Bin Sub root v1) - y <- add myCost (Uni Square x) - z <- add myCost (Bin Div y v2) - w <- add myCost (Bin Mul c1 v2) - lw <- add myCost (Uni Log w) - p <- add myCost (Bin Mul c2 lw) - add myCost (Bin Add z p) - - -buildNLLEGraph Poisson m egraph root = runIdentity $ addToEg `runStateT` egraph - where - addToEg :: EGraphST Identity EClassId - addToEg = do v1 <- add myCost (Var (-1)) - lv <- add myCost (Uni Log v1) - x <- add myCost (Bin Mul v1 lv) - y <- add myCost (Uni Exp root) - z <- add myCost (Bin Add x y) - vt <- add myCost (Bin Mul v1 root) - add myCost (Bin Sub z vt) - -buildNLLEGraph Bernoulli m egraph root = runIdentity $ addToEg `runStateT` egraph - where - addToEg :: EGraphST Identity EClassId - addToEg = do v <- add myCost (Var (-1)) - c1 <- add myCost (Const 1) - c2 <- add myCost (Const (-1)) - mr <- add myCost (Bin Mul c2 root) - er <- add myCost (Uni Exp mr) - er1 <- add myCost (Bin Add c1 er) - ler1 <- add myCost (Uni Log er1) - v1 <- add myCost (Bin Sub c1 v) - v1r <- add myCost (Bin Mul v1 root) - add myCost (Bin Add ler1 v1r) - -buildNLLEGraph ROXY m egraph root = error "ROXY not supported with cache" - --- | Prediction for different distributions -predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector -predict MSE tree theta xss = evalTree xss theta tree -predict LOG10 tree theta xss = evalTree xss theta tree -predict Gaussian tree theta xss = evalTree xss theta tree -predict Bernoulli tree theta xss = logistic $ evalTree xss theta tree -predict Poisson tree theta xss = exp $ evalTree xss theta tree -predict ROXY tree theta xss = evalTree xss theta tree - --- | Gradient of the negative log-likelihood -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) +-- | Builds the per-row loss expression for a given 'Loss', to be summed +-- across rows (e.g. by 'Algorithm.SRTree.AD.evalGradMulti') and +-- differentiated by automatic differentiation. Same special variable +-- convention as 'buildDistLoss'. +buildLoss :: Loss -> Double -> Fix SRTree -> Fix SRTree +buildLoss MSE m tree = ((tree - var (-1)) ** 2) / constv m +buildLoss LOG10 m tree = (((log (y / tree')) / log 10) ** 2) / constv m where - grad :: PVector - grad = M.fromList M.Seq [finitediff ix | ix <- [0..p-1]] - (Sz p) = M.size theta - - 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) - ys' - in (f' - f)/eps - (Sz2 m _) = M.size xss - tree' = buildNLL dist (fromIntegral m) tree - treeArr = IntMap.toAscList $ tree2arr tree' - j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] - flog :: Array D Ix1 Double -> Array D Ix1 Double - flog z = M.map (logBase 10) (z + M.map sqrt (z^2 + 1e-10)) - ys' = (if dist==LOG10 then id else id) (delay ys) + tree' = (tree + sqrt(tree^2 + 1e-10)) + y = (var (-1) + sqrt(var (-1) ^ 2 + 1e-10)) +buildLoss MAE m tree = abs (tree - var (-1)) / constv m -nanTo0 x = x -- if isNaN x || isInfinite x then 0 else x -{-# INLINE nanTo0 #-} +-- | Mean absolute percentage error. A small epsilon is added to the +-- denominator's magnitude to avoid division by zero when the target is +-- (close to) zero. +buildLoss MAPE m tree = (abs (tree - var (-1)) / (abs (var (-1)) + constv 1e-8)) / constv m --- | Gradient of the negative log-likelihood -gradNLLArr MSE 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 LOG10 xss ys mYerr tree j2ix theta = - (M.sum yhat, delay grad') - where - (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix - grad' = M.map nanTo0 grad -gradNLLArr Gaussian xss ys mYerr tree j2ix theta = - (M.sum yhat, delay grad') - where - (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 = (M.sum yhat, delay grad') - where - (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix - grad' = M.map nanTo0 grad -gradNLLArr Poisson xss ys mYerr tree j2ix theta - | M.any (<0) ys = error "For Poisson distribution the output must be non-negative." - | otherwise = (M.sum yhat, delay grad') - where - (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 +-- | Pinball (quantile) loss for a residual @r = y - yhat@: +-- @tau * r@ if @r >= 0@, @(tau - 1) * r@ otherwise. Both cases are +-- captured in closed form by @0.5 * ((2*tau - 1) * r + abs r)@, which +-- avoids branching in the symbolic tree. +buildLoss (Pinball tau) m tree = ((constv (2*tau - 1) * r + abs r) / 2) / constv m + where r = var (-1) - tree --- | Gradient of the negative log-likelihood -gradNLLGraph MSE xss ys mYerr tree theta = - (M.sum yhat, grad') - where - (yhat, grad) = reverseModeGraph xss ys mYerr theta tree - grad' = VS.map nanTo0 grad -gradNLLGraph LOG10 xss ys mYerr tree theta = - (M.sum yhat, grad') - where - (yhat, grad) = reverseModeGraph xss ys mYerr theta tree - grad' = VS.map nanTo0 grad -gradNLLGraph Gaussian xss ys mYerr tree theta = - (M.sum yhat, grad') - where - (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 = (M.sum yhat, grad') - where - (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." - | otherwise = (M.sum yhat, grad') - where - (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 +buildLoss (NLL dist) m tree = buildDistLoss dist m tree --- | e-graph support -gradNLLEGraph MSE xss ys mYerr egraph cache root theta = - (M.sum yhat, grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad -gradNLLEGraph LOG10 xss ys mYerr egraph cache root theta = - (M.sum yhat, grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad - ys' :: PVector - ys' = M.computeAs M.S $ M.map (logBase 10) (delay ys + M.map sqrt (delay ys^2 + 1e-10)) -gradNLLEGraph Gaussian xss ys mYerr egraph cache root theta = - (M.sum yhat, grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad -gradNLLEGraph Bernoulli xss ys mYerr egraph cache root theta - | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1." - | otherwise = (M.sum yhat, grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad -gradNLLEGraph Poisson xss ys mYerr egraph cache root theta - | M.any (<0) ys = error "For Poisson distribution the output must be non-negative." - | otherwise = (M.sum yhat, grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad -gradNLLEGraph ROXY xss ys mYerr egraph cache root theta = - ((*0.5) $ M.sum yhat, VS.map (*(0.5)) $ grad') - where - (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta - grad' = VS.map nanTo0 grad +-- | Builds the predictor expression from a fitted model tree by applying +-- the inverse link function implied by the 'Distribution': @exp@ for +-- 'Poisson', the logistic function for 'Bernoulli', and the identity +-- otherwise. +buildPredictor :: Distribution -> Fix SRTree -> Fix SRTree +buildPredictor Poisson tree = exp tree +buildPredictor Bernoulli tree = 1 / (1 + exp (negate tree)) +buildPredictor _ tree = tree -- | Fisher information of negative log-likelihood -fisherNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector -fisherNLL ROXY mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff +fisherNLL :: Distribution -> Maybe Target -> Columns -> Target -> Fix SRTree -> Target -> Target +fisherNLL ROXY mYerr xss ys tree theta = V.generate 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 + m = V.length ys + p = V.length theta + loss = compileLoss xss (buildDistLoss ROXY (fromIntegral m) tree) ys mYerr + f = loss 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 + theta' <- V.thaw theta + v <- VM.read theta' ix + VM.write theta' ix (v + eps) + thetaPlus <- V.freeze theta' + VM.write theta' ix (v - eps) + thetaMinus <- V.freeze theta' + let fPlus = loss thetaPlus + fMinus = loss thetaMinus pure $ (fPlus + fMinus - 2*f)/(eps*eps) -fisherNLL Gaussian mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff +fisherNLL Gaussian mYerr xss ys tree theta = V.generate 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 + m = V.length ys + p = V.length theta + loss = compileLoss xss (buildDistLoss Gaussian (fromIntegral m) tree) ys mYerr + f = loss 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 + theta' <- V.thaw theta + v <- VM.read theta' ix + VM.write theta' ix (v + eps) + thetaPlus <- V.freeze theta' + VM.write theta' ix (v - eps) + thetaMinus <- V.freeze theta' + let fPlus = loss thetaPlus + fMinus = loss thetaMinus pure $ (fPlus + fMinus - 2*f)/(eps*eps) -fisherNLL dist mYerr xss ys tree theta = makeArray cmp (Sz p) build +fisherNLL dist mYerr xss ys tree theta = V.generate p build where build ix = let dtdix = deriveByParam ix t' d2tdix2 = deriveByParam ix dtdix f' = eval dtdix f'' = eval d2tdix2 - in M.sum $ phi' * f'^2 - res * f'' + in V.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 + -- Gaussian -> V.sum . (/(theta V.! (p-1))) $ phi' * f'^2 - res * f'' + -- _ -> V.sum $ phi' * f'^2 - res * f'' + m = V.length ys + p = V.length theta t' = fst $ floatConstsToParam tree - eval = evalTree xss theta + eval = \t -> compile xss t theta yhat = eval t' - res = delay ys - phi + res = ys - phi yErr = case mYerr of - Nothing -> M.replicate (getComp xss) (Sz m) est + Nothing -> V.replicate 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) + Gaussian -> (yhat, V.replicate m 1) + LeastSquares -> (yhat, V.replicate m 1) + Bernoulli -> (logistic yhat, phi*(V.replicate m 1 - phi)) + Poisson -> (exp yhat, phi) -- | Hessian of negative log-likelihood -- -- 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 PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRMatrix +hessianNLL :: Distribution -> Maybe Target -> Columns -> Target -> Fix SRTree -> Target -> Columns hessianNLL ROXY mYerr xss ys tree theta = undefined -hessianNLL dist mYerr xss ys tree theta = makeArray cmp (Sz (p :. p)) build +hessianNLL Gaussian mYerr xss ys tree theta = [V.generate p (build iy) | iy <- [0..p-1]] where - build (ix :. iy) = let dtdix = deriveByParam ix t' - dtdiy = deriveByParam iy t' - d2tdixy = deriveByParam iy dtdix - fx = eval dtdix - fy = eval dtdiy - fxy = eval d2tdixy - 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 - 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 - - (phi, phi') = case dist of - MSE -> (yhat, M.replicate cmp (Sz m) 1) - LOG10 -> (yhat, M.replicate cmp (Sz m) 1) - Gaussian -> (yhat, M.replicate cmp (Sz m) 1) - Bernoulli -> (logistic yhat, phi*(M.replicate cmp (Sz m) 1 - phi)) - Poisson -> (exp yhat, phi) + build iy ix = let dtdix = deriveByParam ix tree + dtdiy = deriveByParam iy tree + d2tdixy = deriveByParam iy dtdix + fx = eval dtdix + fy = eval dtdiy + fxy = eval d2tdixy + in if ix < p-1 && iy < p-1 + then V.sum . (/yErr) $ fx * fy - res * fxy + else if ix == p-1 && iy == p-1 + then (*0.5) . V.sum . (/ yErr ) $ res*res + else if ix == p-1 + then V.sum . (/yErr) $ res * fy + else V.sum . (/yErr) $ res * fx + m = V.length ys + p = V.length theta + yErr :: Target + yErr = V.replicate m $ exp (theta V.! (p-1)) / est + yhat = eval tree + res = ys - yhat + eval = \t -> compile xss t theta + est = fromIntegral (m - p + 1) -tree2arr :: Fix SRTree -> IntMap.IntMap (Int, Int, Int, Double) -tree2arr tree = IntMap.fromList listTree +hessianNLL dist mYerr xss ys tree theta = [V.generate p (build iy) | iy <- [0..p-1]] 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 + build iy ix = let dtdix = deriveByParam ix t' + dtdiy = deriveByParam iy t' + d2tdixy = deriveByParam iy dtdix + fx = eval dtdix + fy = eval dtdiy + fxy = eval d2tdixy + in case dist of + Gaussian -> V.sum . (/yErr) $ phi' * fx * fy - res * fxy + _ -> V.sum $ phi' * fx * fy - res * fxy - 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) + m = V.length ys + p = V.length theta + t' = tree -- relabelParams tree -- $ floatConstsToParam tree + eval = \t -> compile xss t theta + yErr = case mYerr of + Nothing -> V.replicate m est + Just e -> e + est = fromIntegral (m - p) + yhat = eval t' + res = ys - phi - 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 #-} + (phi, phi') = case dist of + Gaussian -> (yhat, V.replicate m 1) + LeastSquares -> (yhat, V.replicate m 1) + Bernoulli -> (logistic yhat, phi*(V.replicate m 1 - phi)) + Poisson -> (exp yhat, phi) + diff --git a/src/Algorithm/SRTree/ModelSelection.hs b/src/Algorithm/SRTree/ModelSelection.hs --- a/src/Algorithm/SRTree/ModelSelection.hs +++ b/src/Algorithm/SRTree/ModelSelection.hs @@ -1,9 +1,9 @@ {-# LANGUAGE ViewPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE LambdaCase #-} ------------------------------------------------------------------------------ +------------------------------------------------------------------------------- -- | --- Module : Algorithm.SRTree.ModelSelection +-- Module : Algorithm.SRTree.ModelSelection -- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com @@ -11,160 +11,169 @@ -- Portability : ConstraintKinds -- -- Helper functions for model selection criteria --- ------------------------------------------------------------------------------ +------------------------------------------------------------------------------- -module Algorithm.SRTree.ModelSelection where +module Algorithm.SRTree.ModelSelection + ( bic + , aic + , evidence + , fractionalBayesFactor + , mdl + , mdlLatt + , mdlFreq + , logFunctional + , logFunctionalFreq + , ModelEval (..) + , module Algorithm.SRTree.Compile + ) where -import Algorithm.Massiv.Utils ( det ) +import Algorithm.SRTree.Utils ( det ) import Algorithm.SRTree.Likelihoods - ( PVector, SRMatrix, fisherNLL, hessianNLL, nll, Distribution(..) ) -import Data.Massiv.Array (Ix2 (..), Sz (..), (!-!)) -import qualified Data.Massiv.Array as A + ( fisherNLL, hessianNLL + , Distribution(..), Loss(..), buildDistLoss + ) import Data.SRTree -import Data.SRTree.Eval (evalTree) +import Data.SRTree.Eval (Target, Columns, compileLoss) import Data.SRTree.Recursion (cata) -import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as U +import Algorithm.SRTree.Compile import Debug.Trace -- | Bayesian information criterion -bic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -bic dist mYerr xss ys theta tree = p * log n + 2 * nll dist mYerr xss ys tree theta - where - (A.Sz (fromIntegral -> p)) = A.size theta - (A.Sz (fromIntegral -> n)) = A.size ys +bic :: EvaluatedTree -> Double +bic et = valParams et * log (valRows et) + 2 * valLoss et {-# INLINE bic #-} -- | Akaike information criterion -aic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -aic dist mYerr xss ys theta tree = 2 * p + 2 * nll dist mYerr xss ys tree theta - where - (A.Sz (fromIntegral -> p)) = A.size theta - (A.Sz (fromIntegral -> n)) = A.size ys +aic :: EvaluatedTree -> Double +aic et = 2 * valParams et + 2 * valLoss et {-# INLINE aic #-} --- | Evidence -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 +-- | Evidence +evidence :: EvaluatedTree -> Double +evidence et = (1 - b) * valLoss et - valParams et / 2 * log b where - (A.Sz (fromIntegral -> p)) = A.size theta - (A.Sz (fromIntegral -> n)) = A.size ys - b = 1 / sqrt n + b = 1 / sqrt (valRows et) {-# INLINE evidence #-} -fractionalBayesFactor :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -fractionalBayesFactor dist mYerr xss ys theta tree = (1 - b) * nll' - p / 2 * log b + f_compl + p / 2 * log(2*pi*nup) +fractionalBayesFactor :: EvaluatedTree -> Double +fractionalBayesFactor et = (1 - b) * valLoss et - valParams et / 2 * log b + f_compl + valParams et / 2 * log(2*pi*nup) where - nll_val = nll dist mYerr xss ys tree theta - nll_gaus = nll Gaussian mYerr xss ys tree theta - nll' = if dist == MSE then nll_gaus else nll_val - (A.Sz (fromIntegral -> p)) = A.size theta - (A.Sz (fromIntegral -> n)) = A.size ys - b = 1 / sqrt n + b = 1 / sqrt (valRows et) nup = exp(1 - log 3) - f_compl = countNodes tree * log (countUniqueTokens tree) + f_compl = countNodes (valTree et) * log (countUniqueTokens (valTree et)) {-# INLINE fractionalBayesFactor #-} --- | MDL as described in +-- | MDL as described in -- Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023). -mdl :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double -mdl dist mYerr xss ys theta tree = nll' dist mYerr xss ys theta tree - + logFunctional tree - + logParameters dist mYerr xss ys theta tree - where - fisher = fisherNLL dist mYerr xss ys tree theta - theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher - isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +mdl :: EvaluatedTree -> Double +mdl et = valLoss et + logFunctional (valTree et) + valLogParams et {-# 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 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 mYerr xss ys theta tree - where - fisher = fisherNLL dist mYerr xss ys tree theta - theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher - isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +mdlLatt :: EvaluatedTree -> Double +mdlLatt et = valLoss et + logFunctional (valTree et) + valLogParamsLattice et {-# 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 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 mYerr xss ys theta tree +mdlFreq :: EvaluatedTree -> Double +mdlFreq et = valLoss et + logFunctionalFreq (valTree et) + valLogParams et {-# INLINE mdlFreq #-} +-- | The possible metrics used to evaluate\/select a fitted model, +-- ranging from plain loss functions ('EvalLoss', wrapping any 'Loss' -- +-- including a distribution's negative log-likelihood via @EvalLoss (NLL +-- dist)@) to the error metrics and model-selection criteria already +-- provided by this module ('RMSE', 'R2', 'AIC', 'BIC', 'Evidence', 'FBF', +-- 'MDL', 'MDLLatt', 'MDLFreq'). +data ModelEval + = RMSE + | R2 + | AIC + | BIC + | Evidence + | FBF + | MDL + | MDLLatt + | MDLFreq + | EvalLoss Loss + deriving (Show, Read, Eq) + +instance Enum ModelEval where + fromEnum RMSE = 0 + fromEnum R2 = 1 + fromEnum AIC = 2 + fromEnum BIC = 3 + fromEnum Evidence = 4 + fromEnum FBF = 5 + fromEnum MDL = 6 + fromEnum MDLLatt = 7 + fromEnum MDLFreq = 8 + fromEnum (EvalLoss l) = 9 + fromEnum l + + toEnum 0 = RMSE + toEnum 1 = R2 + toEnum 2 = AIC + toEnum 3 = BIC + toEnum 4 = Evidence + toEnum 5 = FBF + toEnum 6 = MDL + toEnum 7 = MDLLatt + toEnum 8 = MDLFreq + toEnum x | x >= 9 = EvalLoss (toEnum (x-9)) + +instance Bounded ModelEval where + minBound = RMSE + maxBound = EvalLoss maxBound + +-- | Evaluates the requested 'ModelEval' metric. +-- +-- for 'RMSE', and 'R2' the tree must have been compiled +-- with MSE loss. + +evalModelSelection :: ModelEval -> EvaluatedTree -> Double +evalModelSelection (EvalLoss MAE) et = valLoss et +evalModelSelection (EvalLoss MAPE) et = valLoss et +evalModelSelection (EvalLoss (Pinball tau)) et = valLoss et +evalModelSelection (EvalLoss (NLL dist)) et = valLoss et +evalModelSelection RMSE et = sqrt (valLoss et) -- assumes MSE +evalModelSelection R2 et = 1 - (valRows et * valLoss et) / valVar et -- assumes MSE +evalModelSelection AIC et = aic et +evalModelSelection BIC et = bic et +evalModelSelection Evidence et = evidence et +evalModelSelection FBF et = fractionalBayesFactor et +evalModelSelection MDL et = mdl et +evalModelSelection MDLLatt et = mdlLatt et +evalModelSelection MDLFreq et = mdlFreq et +{-# INLINE evalModelSelection #-} + -- log of the functional complexity logFunctional :: Fix SRTree -> Double -logFunctional tree = countNodes tree * log (countUniqueTokens tree') - + foldr (\c acc -> log (abs c) + acc) 0 consts - + log(2) * numberOfConsts +logFunctional tree = countNodes tree * log (countUniqueTokens tree') + foldr (\c acc -> log (abs c) + acc) 0 consts + log(2) * numberOfConsts where - tree' = fst $ floatConstsToParam tree - consts = getIntConsts tree + tree' = fst $ floatConstsToParam tree + consts = getIntConsts tree numberOfConsts = fromIntegral $ length consts - signs = sum [1 | a <- getIntConsts tree, a < 0] -- TODO: will we use that? {-# INLINE logFunctional #-} --- same as above but weighted by frequency -logFunctionalFreq :: Fix SRTree -> Double -logFunctionalFreq tree = treeToNat tree' - + foldr (\c acc -> log (abs c) + acc) 0 consts - + countVarNodes tree * log (numberOfVars tree) +-- same as above but weighted by frequency +logFunctionalFreq :: Fix SRTree -> Double +logFunctionalFreq tree = treeToNat tree' + foldr (\c acc -> log (abs c) + acc) 0 consts + countVarNodes tree * log (numberOfVars tree) where - tree' = fst $ floatConstsToParam tree + tree' = fst $ floatConstsToParam tree consts = getIntConsts tree {-# INLINE logFunctionalFreq #-} --- log of the parameters complexity -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 mYerr xss ys tree theta - (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) - $ zip (A.toList theta) (A.toList fisher) - - addIfSignificant (v, f) (acc_v, acc_f, acc_p) - | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1) - | otherwise = (acc_v, acc_f, acc_p) - - isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 - --- same as above but for the Lattice -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 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) - - addIfSignificant (v, f) (acc_v, acc_f, acc_p) - | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1) - | otherwise = (acc_v, acc_f, acc_p) - - isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 - --- flipped version of nll -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 -treeToNat = cata $ - \case - Uni f t -> funToNat f + t - Bin op l r -> opToNat op + l + r - _ -> 0.6610799229372109 +treeToNat = cata $ \case + Uni f t -> funToNat f + t + Bin op l r -> opToNat op + l + r + _ -> 0.6610799229372109 where - opToNat :: Op -> Double opToNat Add = 2.500842464597881 opToNat Sub = 2.500842464597881 @@ -176,15 +185,14 @@ funToNat :: Function -> Double funToNat Sqrt = 4.780867285331753 - funToNat Log = 4.765599813200964 - funToNat Exp = 4.788589331425663 - funToNat Abs = 6.352564869783006 - funToNat Sin = 5.9848400896576885 - funToNat Cos = 5.474014465891698 + funToNat Log = 4.765599813200964 + funToNat Exp = 4.788589331425663 + funToNat Abs = 6.352564869783006 + funToNat Sin = 5.9848400896576885 + funToNat Cos = 5.474014465891698 funToNat Sinh = 8.038963823353235 funToNat Cosh = 8.262107374667444 funToNat Tanh = 7.85664226655928 - funToNat Tan = 8.262107374667444 - funToNat _ = 8.262107374667444 - --funToNat Factorial = 7.702491586732021 + funToNat Tan = 8.262107374667444 + funToNat _ = 8.262107374667444 {-# INLINE treeToNat #-} diff --git a/src/Algorithm/SRTree/NonlinearOpt.hs b/src/Algorithm/SRTree/NonlinearOpt.hs --- a/src/Algorithm/SRTree/NonlinearOpt.hs +++ b/src/Algorithm/SRTree/NonlinearOpt.hs @@ -1,976 +1,98 @@ -{-# OPTIONS_GHC -Wall #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE TypeApplications #-} - -{- | -Module : Numeric.NLOPT -Copyright : (c) Matthew Peddie 2017 -License : BSD3 -Maintainer : Matthew Peddie -Stability : provisional -Portability : GHC - -This module provides a high-level, @hmatrix@-compatible interface to -the library by -Steven G. Johnson. - -NOTE: This is an adaptation from https://hackage.haskell.org/package/hmatrix-nlopt-0.2.0.0 -that removes the dependency to hmatrix and support any Vector Storage. - -= Documentation - -Most non-numerical details are documented, but for specific -information on what the optimization methods do, how constraints are -handled, etc., you should consult: - - * The - - * The - - * The - -= Example program - -The following interactive session example uses the Nelder-Mead simplex -algorithm, a derivative-free local optimizer, to minimize a trivial -function with a minimum of 22.0 at @(0, 0)@. - ->>> import Numeric.LinearAlgebra ( dot, fromList ) ->>> let objf x = x `dot` x + 22 -- define objective ->>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion ->>> let algorithm = NELDERMEAD objf [] Nothing -- specify algorithm ->>> let problem = LocalProblem 2 stop algorithm -- specify problem ->>> let x0 = fromList [5, 10] -- specify initial guess ->>> minimizeLocal problem x0 -Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) - --} - -module Algorithm.SRTree.NonlinearOpt ( - -- * Specifying the objective function - Objective - , ObjectiveD - , Preconditioner - -- * Specifying the constraints - -- ** Bound constraints - , Bounds(..) - -- ** Nonlinear constraints - -- - -- $nonlinearconstraints - - -- *** Constraint functions - , ScalarConstraint - , ScalarConstraintD - , VectorConstraint - , VectorConstraintD - -- *** Constraint types - , Constraint(..) - , EqualityConstraint(..) - , InequalityConstraint(..) - -- *** Collections of constraints - , EqualityConstraints - , EqualityConstraintsD - , InequalityConstraints - , InequalityConstraintsD - -- * Stopping conditions - -- - -- $nonempty - , StoppingCondition(..) - , NonEmpty(..) - -- * Additional configuration - , RandomSeed(..) - , Population(..) - , VectorStorage(..) - , InitialStep(..) - -- * Minimization problems - -- ** Local minimization - , LocalAlgorithm(..) - , LocalProblem(..) - , minimizeLocal - -- ** Global minimization - , GlobalAlgorithm(..) - , GlobalProblem(..) - , minimizeGlobal - -- ** Minimization by augmented Lagrangian - , AugLagAlgorithm(..) - , AugLagProblem(..) - , minimizeAugLag - -- ** Results - , Solution(..) - , N.Result(..) - ) where - -import qualified Numeric.Optimization.NLOPT.Bindings as N - -import Data.List.NonEmpty (NonEmpty(..)) - -import qualified Data.Vector.Storable as V -import Data.Vector.Storable ( Vector ) - -import Control.Exception ( Exception ) -import qualified Control.Exception as Ex -import Data.Typeable ( Typeable ) -import Data.Foldable ( traverse_ ) - -import System.IO.Unsafe ( unsafePerformIO ) - --- each element i contains a row vec -type Matrix a = [Vector a] - -flatten :: V.Storable a => Matrix a -> Vector a -flatten = V.concat -{-# INLINE flatten #-} - -{- Function wrapping for the immutable HMatrix interface -} -wrapScalarFunction :: (Vector Double -> Double) -> N.ScalarFunction () -wrapScalarFunction f params _ _ = return $ f params - -wrapScalarFunctionD :: (Vector Double -> (Double, Vector Double)) - -> N.ScalarFunction () -wrapScalarFunctionD f params grad _ = do - case grad of - Nothing -> return () - Just g -> V.copy g usergrad - return result - where - (result, usergrad) = f params - -wrapVectorFunction :: (Vector Double -> Word -> Vector Double) - -> Word -> N.VectorFunction () -wrapVectorFunction f n params vout _ _ = V.copy vout $ f params n - -wrapVectorFunctionD :: (Vector Double -> Word -> (Vector Double, Matrix Double)) - -> Word -> N.VectorFunction () -wrapVectorFunctionD f n params vout jac _ = do - V.copy vout result - case jac of - Nothing -> return () - Just j -> V.copy j (flatten userjac) - where - (result, userjac) = f params n - -wrapPreconditionerFunction :: (Vector Double -> Vector Double -> Vector Double) - -> N.PreconditionerFunction () -wrapPreconditionerFunction f params v vpre _ = V.copy vpre (f params v) - -{- Objective functions -} --- | An objective function that calculates the objective value at the --- given parameter vector. -type Objective - = Vector Double -- ^ Parameter vector - -> Double -- ^ Objective function value - --- | An objective function that calculates both the objective value --- and the gradient of the objective with respect to the input --- parameter vector, at the given parameter vector. -type ObjectiveD - = Vector Double -- ^ Parameter vector - -> (Double, Vector Double) -- ^ (Objective function value, gradient) - --- | A preconditioner function, which computes @vpre = H(x) v@, where --- @H@ is the Hessian matrix: the positive semi-definite second --- derivative at the given parameter vector @x@, or an approximation --- thereof. -type Preconditioner - = Vector Double -- ^ Parameter vector @x@ - -> Vector Double -- ^ Vector @v@ to precondition at @x@ - -> Vector Double -- ^ Preconditioned vector @vpre@ - -data ObjectiveFunction f - = MinimumObjective f - | PreconditionedMinimumObjective Preconditioner f - -applyObjective :: N.Opt -> ObjectiveFunction Objective -> IO N.Result -applyObjective opt (MinimumObjective f) = - N.set_min_objective opt (wrapScalarFunction f) () -applyObjective opt (PreconditionedMinimumObjective p f) = - N.set_precond_min_objective opt (wrapScalarFunction f) - (wrapPreconditionerFunction p) () - -applyObjectiveD :: N.Opt -> ObjectiveFunction ObjectiveD -> IO N.Result -applyObjectiveD opt (MinimumObjective f) = - N.set_min_objective opt (wrapScalarFunctionD f) () -applyObjectiveD opt (PreconditionedMinimumObjective p f) = - N.set_precond_min_objective opt (wrapScalarFunctionD f) - (wrapPreconditionerFunction p) () - -{- Constraint functions -} --- | A constraint function which returns @c(x)@ given the parameter --- vector @x@. The constraint will enforce that @c(x) == 0@ (equality --- constraint) or @c(x) <= 0@ (inequality constraint). -type ScalarConstraint - = Vector Double -- ^ Parameter vector @x@ - -> Double -- ^ Constraint violation (deviation from 0) - --- | A constraint function which returns @c(x)@ given the parameter --- vector @x@ along with the gradient of @c(x)@ with respect to @x@ at --- that point. The constraint will enforce that @c(x) == 0@ (equality --- constraint) or @c(x) <= 0@ (inequality constraint). -type ScalarConstraintD - = Vector Double -- ^ Parameter vector - -> (Double, Vector Double) -- ^ (Constraint violation, constraint gradient) - --- | A constraint function which returns a vector @c(x)@ given the --- parameter vector @x@. The constraint will enforce that @c(x) == 0@ --- (equality constraint) or @c(x) <= 0@ (inequality constraint). -type VectorConstraint - = Vector Double -- ^ Parameter vector - -> Word -- ^ Constraint Vectorize - -> Vector Double -- ^ Constraint violation vector - --- | A constraint function which returns @c(x)@ given the parameter --- vector @x@ along with the Jacobian (first derivative) matrix of --- @c(x)@ with respect to @x@ at that point. The constraint will --- enforce that @c(x) == 0@ (equality constraint) or @c(x) <= 0@ --- (inequality constraint). -type VectorConstraintD - = Vector Double -- ^ Parameter vector - -> Word -- ^ Constraint Vectorize - -> (Vector Double, Matrix Double) -- ^ (Constraint violation vector, - -- constraint Jacobian) - --- $nonlinearconstraints --- --- Note that most NLOPT algorithms do not support nonlinear --- constraints natively; if you need to enforce nonlinear constraints, --- you may want to use the 'AugLagAlgorithm' family of solvers, which --- can add nonlinear constraints to some algorithm that does not --- support them by a principled modification of the objective --- function. --- --- == Example program --- --- The following interactive session example enforces a scalar --- constraint on the problem given in the beginning of the module: the --- parameters must always sum to 1. The minimizer finds a constrained --- minimum of 22.5 at @(0.5, 0.5)@. --- --- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) --- >>> let objf x = x `dot` x + 22 --- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] --- >>> -- define constraint function: --- >>> let constraintf x = sum (toList x) - 1.0 --- >>> -- define constraint object to pass to the algorithm: --- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 --- >>> let algorithm = COBYLA objf [] [] [constraint] Nothing --- >>> let problem = LocalProblem 2 stop algorithm --- >>> let x0 = fromList [5, 10] --- >>> minimizeLocal problem x0 --- Right (Solution {solutionCost = 22.500000000013028, solutionParams = [0.5000025521533521,0.49999744784664796], solutionResult = FTOL_REACHED}) - - -data Constraint s v - -- | A scalar constraint. - = Scalar s - -- | A vector constraint. - | Vector Word v - -- | A scalar constraint with an attached preconditioning function. - | Preconditioned Preconditioner s - --- | An equality constraint, comprised of both the constraint function --- (or functions, if a preconditioner is used) along with the desired --- tolerance. -data EqualityConstraint s v = EqualityConstraint - { eqConstraintFunctions :: Constraint s v - , eqConstraintTolerance :: Double - } - --- | An inequality constraint, comprised of both the constraint --- function (or functions, if a preconditioner is used) along with the --- desired tolerance. -data InequalityConstraint s v = InequalityConstraint - { ineqConstraintFunctions :: Constraint s v - , ineqConstraintTolerance :: Double - } - --- | A collection of equality constraints that do not supply --- constraint derivatives. -type EqualityConstraints = - [EqualityConstraint ScalarConstraint VectorConstraint] - --- | A collection of inequality constraints that do not supply --- constraint derivatives. -type InequalityConstraints = - [InequalityConstraint ScalarConstraint VectorConstraint] - --- | A collection of equality constraints that supply constraint --- derivatives. -type EqualityConstraintsD = [EqualityConstraint ScalarConstraintD VectorConstraintD] - --- | A collection of inequality constraints that supply constraint --- derivatives. -type InequalityConstraintsD = [InequalityConstraint ScalarConstraintD VectorConstraintD] - -class ApplyConstraint constraint where - applyConstraint :: N.Opt -> constraint -> IO N.Result - -instance ApplyConstraint (EqualityConstraint ScalarConstraint VectorConstraint) where - applyConstraint opt (EqualityConstraint ty tol) = case ty of - Scalar s -> - N.add_equality_constraint opt (wrapScalarFunction s) () tol - Vector n v -> - N.add_equality_mconstraint opt n (wrapVectorFunction v n) () tol - Preconditioned p s -> - N.add_precond_equality_constraint opt (wrapScalarFunction s) - (wrapPreconditionerFunction p) () tol - -instance ApplyConstraint (InequalityConstraint ScalarConstraint VectorConstraint) where - applyConstraint opt (InequalityConstraint ty tol) = case ty of - Scalar s -> - N.add_inequality_constraint opt (wrapScalarFunction s) () tol - Vector n v -> - N.add_inequality_mconstraint opt n (wrapVectorFunction v n) () tol - Preconditioned p s -> - N.add_precond_inequality_constraint opt (wrapScalarFunction s) - (wrapPreconditionerFunction p) () tol - -instance ApplyConstraint (EqualityConstraint ScalarConstraintD VectorConstraintD) where - applyConstraint opt (EqualityConstraint ty tol) = case ty of - Scalar s -> - N.add_equality_constraint opt (wrapScalarFunctionD s) () tol - Vector n v -> - N.add_equality_mconstraint opt n (wrapVectorFunctionD v n) () tol - Preconditioned p s -> - N.add_precond_equality_constraint opt (wrapScalarFunctionD s) - (wrapPreconditionerFunction p) () tol - -instance ApplyConstraint (InequalityConstraint ScalarConstraintD VectorConstraintD) where - applyConstraint opt (InequalityConstraint ty tol) = case ty of - Scalar s -> - N.add_inequality_constraint opt (wrapScalarFunctionD s) () tol - Vector n v -> - N.add_inequality_mconstraint opt n (wrapVectorFunctionD v n) () tol - Preconditioned p s -> - N.add_precond_inequality_constraint opt (wrapScalarFunctionD s) - (wrapPreconditionerFunction p) () tol - -{- Bounds -} - --- | Bound constraints are specified by vectors of the same dimension --- as the parameter space. --- --- == Example program --- --- The following interactive session example enforces lower bounds on --- the example from the beginning of the module. This prevents the --- optimizer from locating the true minimum at @(0, 0)@; a slightly --- higher constrained minimum at @(1, 1)@ is found. Note that the --- optimizer returns 'N.XTOL_REACHED' rather than 'N.FTOL_REACHED', --- because the bound constraint is active at the final minimum. --- --- >>> import Numeric.LinearAlgebra ( dot, fromList ) --- >>> let objf x = x `dot` x + 22 -- define objective --- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion --- >>> let lowerbound = LowerBounds $ fromList [1, 1] -- specify bounds --- >>> let algorithm = NELDERMEAD objf [lowerbound] Nothing -- specify algorithm --- >>> let problem = LocalProblem 2 stop algorithm -- specify problem --- >>> let x0 = fromList [5, 10] -- specify initial guess --- >>> minimizeLocal problem x0 --- Right (Solution {solutionCost = 24.0, solutionParams = [1.0,1.0], solutionResult = XTOL_REACHED}) -data Bounds - -- | Lower bound vector @v@ means we want @x >= v@. - = LowerBounds (Vector Double) - -- | Upper bound vector @u@ means we want @x <= u@. - | UpperBounds (Vector Double) - deriving (Eq, Show, Read) - -applyBounds :: N.Opt -> Bounds -> IO N.Result -applyBounds opt (LowerBounds lbvec) = N.set_lower_bounds opt lbvec -applyBounds opt (UpperBounds ubvec) = N.set_upper_bounds opt ubvec - -{- Stopping conditions -} - --- | A 'StoppingCondition' tells NLOPT when to stop working on a --- minimization problem. When multiple 'StoppingCondition's are --- provided, the problem will stop when any one condition is met. -data StoppingCondition - -- | Stop minimizing when an objective value @J@ less than or equal - -- to the provided value is found. - = MinimumValue Double - -- | Stop minimizing when an optimization step changes the objective - -- value @J@ by less than the provided tolerance multiplied by @|J|@. - | ObjectiveRelativeTolerance Double - -- | Stop minimizing when an optimization step changes the objective - -- value by less than the provided tolerance. - | ObjectiveAbsoluteTolerance Double - -- | Stop when an optimization step changes /every element/ of the - -- parameter vector @x@ by less than @x@ scaled by the provided - -- tolerance. - | ParameterRelativeTolerance Double - -- | Stop when an optimization step changes /every element/ of the - -- parameter vector @x@ by less than the corresponding element in - -- the provided vector of tolerances values. - | ParameterAbsoluteTolerance (Vector Double) - -- | Stop when the number of evaluations of the objective function - -- exceeds the provided count. - | MaximumEvaluations Word - -- | Stop when the optimization time exceeds the provided time (in - -- seconds). This is not a precise limit. - | MaximumTime Double - deriving (Eq, Show, Read) - --- $nonempty --- --- The 'NonEmpty' data type from 'Data.List.NonEmpty' is re-exported --- here, because it is used to ensure that you always specify at least --- one stopping condition. - -applyStoppingCondition :: N.Opt -> StoppingCondition -> IO N.Result -applyStoppingCondition opt (MinimumValue x) = N.set_stopval opt x -applyStoppingCondition opt (ObjectiveRelativeTolerance x) = N.set_ftol_rel opt x -applyStoppingCondition opt (ObjectiveAbsoluteTolerance x) = N.set_ftol_abs opt x -applyStoppingCondition opt (ParameterRelativeTolerance x) = N.set_xtol_rel opt x -applyStoppingCondition opt (ParameterAbsoluteTolerance v) = N.set_xtol_abs opt v -applyStoppingCondition opt (MaximumEvaluations n) = N.set_maxeval opt n -applyStoppingCondition opt (MaximumTime deltat) = N.set_maxtime opt deltat - -{- Random seed control -} - --- | This specifies how to initialize the random number generator for --- stochastic algorithms. -data RandomSeed - -- | Seed the RNG with the provided value. - = SeedValue Word - -- | Seed the RNG using the system clock. - | SeedFromTime - -- | Don't perform any explicit initialization of the RNG. - | Don'tSeed - deriving (Eq, Show, Read) - -applyRandomSeed :: RandomSeed -> IO () -applyRandomSeed Don'tSeed = return () -applyRandomSeed (SeedValue n) = N.srand n -applyRandomSeed SeedFromTime = N.srand_time - -{- Random stuff -} - --- | This specifies the population size for algorithms that use a pool --- of solutions. -newtype Population = Population Word deriving (Eq, Show, Read) - -applyPopulation :: N.Opt -> Population -> IO N.Result -applyPopulation opt (Population n) = N.set_population opt n - --- | This specifies the memory size to be used by algorithms like --- 'LBFGS' which store approximate Hessian or Jacobian matrices. -newtype VectorStorage = VectorStorage Word deriving (Eq, Show, Read) - -applyVectorStorage :: N.Opt -> VectorStorage -> IO N.Result -applyVectorStorage opt (VectorStorage n) = N.set_vector_storage opt n - --- | This vector with the same dimension as the parameter vector @x@ --- specifies the initial step for the optimizer to take. (This --- applies to local gradient-free algorithms, which cannot use --- gradients to estimate how big a step to take.) -newtype InitialStep = InitialStep (Vector Double) deriving (Eq, Show, Read) - -applyInitialStep :: N.Opt -> InitialStep -> IO N.Result -applyInitialStep opt (InitialStep v) = N.set_initial_step opt v - -{- Algorithms -} - -data GlobalProblem = GlobalProblem - { lowerBounds :: Vector Double -- ^ Lower bounds for @x@ - , upperBounds :: Vector Double -- ^ Upper bounds for @x@ - , gstop :: NonEmpty StoppingCondition -- ^ At least one stopping - -- condition - , galgorithm :: GlobalAlgorithm -- ^ Algorithm specification - } - --- | These are the global minimization algorithms provided by NLOPT. Please see --- --- for more details on how the methods work and how they relate to one another. --- --- Optional parameters are wrapped in a 'Maybe'; for example, if you --- see 'Maybe' 'Population', you can simply specify 'Nothing' to use --- the default behavior. -data GlobalAlgorithm - -- | DIviding RECTangles - = DIRECT Objective - -- | DIviding RECTangles, locally-biased variant - | DIRECT_L Objective - -- | DIviding RECTangles, "slightly randomized" - | DIRECT_L_RAND Objective RandomSeed - -- | DIviding RECTangles, unscaled version - | DIRECT_NOSCAL Objective - -- | DIviding RECTangles, locally-biased and unscaled - | DIRECT_L_NOSCAL Objective - -- | DIviding RECTangles, locally-biased, unscaled and "slightly - -- randomized" - | DIRECT_L_RAND_NOSCAL Objective RandomSeed - -- | DIviding RECTangles, original FORTRAN implementation - | ORIG_DIRECT Objective InequalityConstraints - -- | DIviding RECTangles, locally-biased, original FORTRAN - -- implementation - | ORIG_DIRECT_L Objective InequalityConstraints - -- | Stochastic Global Optimization. - -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ - | STOGO ObjectiveD - -- | Stochastic Global Optimization, randomized variant. - -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ - | STOGO_RAND ObjectiveD RandomSeed - -- | Controlled Random Search with Local Mutation - | CRS2_LM Objective RandomSeed (Maybe Population) - -- | Improved Stochastic Ranking Evolution Strategy - | ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population) - -- | Evolutionary Algorithm - | ESCH Objective - -- | Original Multi-Level Single-Linkage - | MLSL Objective LocalProblem (Maybe Population) - -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy - -- Sequence for starting points - | MLSL_LDS Objective LocalProblem (Maybe Population) - -algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm -algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT -algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L -algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND -algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL -algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL -algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL -algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT -algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L -algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO -algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND -algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM -algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES -algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH -algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL -algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS - -applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO () -applyGlobalObjective opt alg = go alg - where - obj = tryTo . applyObjective opt . MinimumObjective - objD = tryTo . applyObjectiveD opt . MinimumObjective - - go (DIRECT o) = obj o - go (DIRECT_L o) = obj o - go (DIRECT_NOSCAL o) = obj o - go (DIRECT_L_NOSCAL o) = obj o - go (ESCH o) = obj o - go (STOGO o) = objD o - go (DIRECT_L_RAND o _) = obj o - go (DIRECT_L_RAND_NOSCAL o _) = obj o - go (ORIG_DIRECT o _) = obj o - go (ORIG_DIRECT_L o _) = obj o - go (STOGO_RAND o _) = objD o - go (CRS2_LM o _ _) = obj o - go (ISRES o _ _ _ _) = obj o - go (MLSL o _ _) = obj o - go (MLSL_LDS o _ _) = obj o - -applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO () -applyGlobalAlgorithm opt alg = do - applyGlobalObjective opt alg - go alg - where - seed = applyRandomSeed - pop = maybe (return ()) (tryTo . applyPopulation opt) - ic = traverse_ (tryTo . applyConstraint opt) - ec = traverse_ (tryTo . applyConstraint opt) - - local lp = setupLocalProblem lp >>= N.set_local_optimizer opt - - go (DIRECT_L_RAND _ s) = seed s - go (DIRECT_L_RAND_NOSCAL _ s) = seed s - go (ORIG_DIRECT _ ineq) = ic ineq - go (ORIG_DIRECT_L _ ineq) = ic ineq - go (STOGO_RAND _ s) = seed s - go (CRS2_LM _ s p) = seed s *> pop p - go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p - go (MLSL _ lp p) = local lp *> pop p - go (MLSL_LDS _ lp p) = local lp *> pop p - go _ = return () - -tryTo :: IO N.Result -> IO () -tryTo act = do - result <- act - if (N.isSuccess result) - then return () - else Ex.throw $ NloptException result - -data NloptException = NloptException N.Result deriving (Show, Typeable) -instance Exception NloptException - --- | Solve the specified global optimization problem. --- --- = Example program --- --- The following interactive session example uses the 'ISRES' --- algorithm, a stochastic, derivative-free global optimizer, to --- minimize a trivial function with a minimum of 22.0 at @(0, 0)@. --- The search is conducted within a box from -10 to 10 in each --- dimension. --- --- >>> import Numeric.LinearAlgebra ( dot, fromList ) --- >>> let objf x = x `dot` x + 22 -- define objective --- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion --- >>> let algorithm = ISRES objf [] [] (SeedValue 22) Nothing -- specify algorithm --- >>> let lowerbounds = fromList [-10, -10] -- specify bounds --- >>> let upperbounds = fromList [10, 10] -- specify bounds --- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm --- >>> let x0 = fromList [5, 8] -- specify initial guess --- >>> minimizeGlobal problem x0 --- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED}) -minimizeGlobal :: GlobalProblem -- ^ Problem specification - -> Vector Double -- ^ Initial parameter guess - -> Either N.Result Solution -- ^ Optimization results -minimizeGlobal prob x0 = - unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler - where - handler :: NloptException -> IO (Either N.Result a) - handler (NloptException retcode) = return $ Left retcode - -applyGlobalProblem :: N.Opt -> GlobalProblem -> IO () -applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do - tryTo $ applyBounds opt (LowerBounds lb) - tryTo $ applyBounds opt (UpperBounds ub) - traverse_ (tryTo . applyStoppingCondition opt) stop - applyGlobalAlgorithm opt alg - -newOpt :: N.Algorithm -> Word -> IO N.Opt -newOpt alg sz = do - opt' <- N.create alg sz - case opt' of - Nothing -> Ex.throw $ NloptException N.FAILURE - Just opt -> return opt - -setupGlobalProblem :: GlobalProblem -> IO N.Opt -setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do - opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp) - applyGlobalProblem opt gp - return opt - -solveProblem :: N.Opt -> Vector Double -> IO Solution -solveProblem opt x0 = do - (N.Output outret outcost outx nevals) <- N.optimize opt x0 - if (N.isSuccess outret) - then return $ Solution outcost outx outret nevals - else Ex.throw $ NloptException outret - -minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution -minimizeGlobal' gp x0 = do - opt <- setupGlobalProblem gp - solveProblem opt x0 - -data LocalProblem = LocalProblem - { lsize :: Word -- ^ The dimension of the - -- parameter vector. - , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping - -- condition - , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification - } - --- | These are the local minimization algorithms provided by NLOPT. Please see --- --- for more details on how the methods work and how they relate to one --- another. Note that some local methods require you provide --- derivatives (gradients or Jacobians) for your objective function --- and constraint functions. --- --- Optional parameters are wrapped in a 'Maybe'; for example, if you --- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to --- use the default behavior. -data LocalAlgorithm - -- | Limited-memory BFGS - = LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage) - -- | Limited-memory BFGS - | LBFGS ObjectiveD (Maybe VectorStorage) - -- | Shifted limited-memory variable-metric, rank-2 - | VAR2 ObjectiveD (Maybe VectorStorage) - -- | Shifted limited-memory variable-metric, rank-1 - | VAR1 ObjectiveD (Maybe VectorStorage) - -- | Truncated Newton's method - | TNEWTON ObjectiveD (Maybe VectorStorage) - -- | Truncated Newton's method with automatic restarting - | TNEWTON_RESTART ObjectiveD (Maybe VectorStorage) - -- | Preconditioned truncated Newton's method - | TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage) - -- | Preconditioned truncated Newton's method with automatic - -- restarting - | TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage) - -- | Method of moving averages - | MMA ObjectiveD InequalityConstraintsD - -- | Sequential Least-Squares Quadratic Programming - | SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD - -- | Conservative Convex Separable Approximation - | CCSAQ ObjectiveD Preconditioner - -- | PRincipal AXIS gradient-free local optimization - | PRAXIS Objective [Bounds] (Maybe InitialStep) - -- | Constrained Optimization BY Linear Approximations - | COBYLA Objective [Bounds] InequalityConstraints EqualityConstraints - (Maybe InitialStep) - -- | Powell's NEWUOA algorithm - | NEWUOA Objective (Maybe InitialStep) - -- | Powell's NEWUOA algorithm with bounds by SGJ - | NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep) - -- | Nelder-Mead Simplex gradient-free method - | NELDERMEAD Objective [Bounds] (Maybe InitialStep) - -- | NLOPT implementation of Rowan's Subplex algorithm - | SBPLX Objective [Bounds] (Maybe InitialStep) - -- | Bounded Optimization BY Quadratic Approximations - | BOBYQA Objective [Bounds] (Maybe InitialStep) - -algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm -algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL -algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS -algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2 -algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1 -algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON -algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART -algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND -algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART -algorithmEnumOfLocal (MMA _ _) = N.LD_MMA -algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP -algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ -algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS -algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA -algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA -algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA -algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD -algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX -algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA - -applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO () -applyLocalObjective opt alg = go alg - where - obj = tryTo . applyObjective opt . MinimumObjective - objD = tryTo . applyObjectiveD opt . MinimumObjective - precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p - - go (LBFGS_NOCEDAL o _) = objD o - go (LBFGS o _) = objD o - go (VAR2 o _) = objD o - go (VAR1 o _) = objD o - go (TNEWTON o _) = objD o - go (TNEWTON_RESTART o _) = objD o - go (TNEWTON_PRECOND o _) = objD o - go (TNEWTON_PRECOND_RESTART o _) = objD o - go (MMA o _) = objD o - go (SLSQP o _ _ _) = objD o - go (CCSAQ o prec) = precond prec o - go (PRAXIS o _ _) = obj o - go (COBYLA o _ _ _ _) = obj o - go (NEWUOA o _) = obj o - go (NEWUOA_BOUND o _ _) = obj o - go (NELDERMEAD o _ _) = obj o - go (SBPLX o _ _) = obj o - go (BOBYQA o _ _) = obj o - -applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO () -applyLocalAlgorithm opt alg = do - applyLocalObjective opt alg - go alg - where - ic = traverse_ (tryTo . applyConstraint opt) - icd = traverse_ (tryTo . applyConstraint opt) - ec = traverse_ (tryTo . applyConstraint opt) - ecd = traverse_ (tryTo . applyConstraint opt) - store = maybe (return ()) (tryTo . applyVectorStorage opt) - bound = traverse_ (tryTo . applyBounds opt) - step0 = maybe (return ()) (tryTo . applyInitialStep opt) - - go (LBFGS_NOCEDAL _ vs) = store vs - go (LBFGS _ vs) = store vs - go (VAR2 _ vs) = store vs - go (VAR1 _ vs) = store vs - go (TNEWTON _ vs) = store vs - go (TNEWTON_RESTART _ vs) = store vs - go (TNEWTON_PRECOND _ vs) = store vs - go (TNEWTON_PRECOND_RESTART _ vs) = store vs - go (MMA _ ineqd) = icd ineqd - go (SLSQP _ b ineqd eqd) = - bound b *> icd ineqd *> ecd eqd - go (CCSAQ _ _ ) = return () - go (PRAXIS _ b s) = bound b *> step0 s - go (COBYLA _ b ineq eq s) = - bound b *> ic ineq *> ec eq *> step0 s - go (NEWUOA _ s) = step0 s - go (NEWUOA_BOUND _ b s) = bound b *> step0 s - go (NELDERMEAD _ b s) = bound b *> step0 s - go (SBPLX _ b s) = bound b *> step0 s - go (BOBYQA _ b s) = bound b *> step0 s - -applyLocalProblem :: N.Opt -> LocalProblem -> IO () -applyLocalProblem opt (LocalProblem _ stop alg) = do - traverse_ (tryTo . applyStoppingCondition opt) stop - applyLocalAlgorithm opt alg - -setupLocalProblem :: LocalProblem -> IO N.Opt -setupLocalProblem lp@(LocalProblem sz _ alg) = do - opt <- newOpt (algorithmEnumOfLocal alg) sz - applyLocalProblem opt lp - return opt - -minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution -minimizeLocal' lp x0 = do - opt <- setupLocalProblem lp - solveProblem opt x0 - --- | --- == Example program --- --- The following interactive session example enforces the same scalar --- constraint as the nonlinear constraint example, but this time it --- uses the SLSQP solver to find the minimum. --- --- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale ) --- >>> let objf x = (x `dot` x + 22, 2 `scale` x) --- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] --- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1]) --- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 --- >>> let algorithm = SLSQP objf [] [] [constraint] --- >>> let problem = LocalProblem 2 stop algorithm --- >>> let x0 = fromList [5, 10] --- >>> minimizeLocal problem x0 --- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED}) -minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution -minimizeLocal prob x0 = - unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler - where - handler :: NloptException -> IO (Either N.Result a) - handler (NloptException retcode) = return $ Left retcode - -class ProblemSize c where - problemSize :: c -> Word - -instance ProblemSize LocalProblem where - problemSize = lsize - -instance ProblemSize GlobalProblem where - problemSize = fromIntegral . V.length . lowerBounds - -instance ProblemSize AugLagProblem where - problemSize (AugLagProblem _ _ alg) = case alg of - AUGLAG_LOCAL lp _ _ -> problemSize lp - AUGLAG_EQ_LOCAL lp -> problemSize lp - AUGLAG_GLOBAL gp _ _ -> problemSize gp - AUGLAG_EQ_GLOBAL gp -> problemSize gp - - --- | __IMPORTANT NOTE__ --- --- For augmented lagrangian problems, you, the user, are responsible --- for providing the appropriate type of constraint. If the --- subsidiary problem requires an `ObjectiveD`, then you should --- provide constraint functions with derivatives. If the subsidiary --- problem requires an `Objective`, you should provide constraint --- functions without derivatives. If you don't do this, you may get a --- runtime error. -data AugLagProblem = AugLagProblem - { alEquality :: EqualityConstraints -- ^ Possibly empty set of - -- equality constraints - , alEqualityD :: EqualityConstraintsD -- ^ Possibly empty set of - -- equality constraints with - -- derivatives - , alalgorithm :: AugLagAlgorithm -- ^ Algorithm specification. - } - --- | The Augmented Lagrangian solvers allow you to enforce nonlinear --- constraints while using local or global algorithms that don't --- natively support them. The subsidiary problem is used to do the --- minimization, but the @AUGLAG@ methods modify the objective to --- enforce the constraints. Please see --- --- for more details on how the methods work and how they relate to one another. --- --- See the documentation for 'AugLagProblem' for an important note --- about the constraint functions. -data AugLagAlgorithm - -- | AUGmented LAGrangian with a local subsidiary method - = AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD - -- | AUGmented LAGrangian with a local subsidiary method and with - -- penalty functions only for equality constraints - | AUGLAG_EQ_LOCAL LocalProblem - -- | AUGmented LAGrangian with a global subsidiary method - | AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD - -- | AUGmented LAGrangian with a global subsidiary method and with - -- penalty functions only for equality constraints. - | AUGLAG_EQ_GLOBAL GlobalProblem - -algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm -algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG -algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ -algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG -algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ - --- | This structure is returned in the event of a successful --- optimization. -data Solution = Solution - { solutionCost :: Double -- ^ The objective function value - -- at the minimum - , 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 () -applyAugLagAlgorithm opt alg = go alg - where - ic = traverse_ (tryTo . applyConstraint opt) - icd = traverse_ (tryTo . applyConstraint opt) - -- AUGLAG won't work at all if you don't pass it the same - -- objective as the subproblem -- here we pull out the subproblem - -- objectives from the algorithm spec and set the same objective - -- function so the user can't mess it up. - local lp = tryTo $ do - localopt <- setupLocalProblem lp - applyLocalObjective opt (lalgorithm lp) - N.set_local_optimizer opt localopt - global gp = do - tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt - applyGlobalObjective opt (galgorithm gp) - - go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd - go (AUGLAG_EQ_LOCAL lp) = local lp - go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd - go (AUGLAG_EQ_GLOBAL gp) = global gp - -applyAugLagProblem :: N.Opt -> AugLagProblem -> IO () -applyAugLagProblem opt (AugLagProblem eq eqd alg) = do - traverse_ (tryTo . applyConstraint opt) eq - traverse_ (tryTo . applyConstraint opt) eqd - applyAugLagAlgorithm opt alg - -minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution -minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do - opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap) - applyAugLagProblem opt ap - solveProblem opt x0 - --- | --- == Example program --- --- The following interactive session example enforces the same scalar --- constraint as the nonlinear constraint example, but this time it --- uses the augmented Lagrangian method to enforce the constraint and --- the 'SBPLX' algorithm, which does not support nonlinear constraints --- itself, to perform the minimization. As before, the parameters --- must always sum to 1, and the minimizer finds the same constrained --- minimum of 22.5 at @(0.5, 0.5)@. --- --- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) --- >>> let objf x = x `dot` x + 22 --- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] --- >>> let algorithm = SBPLX objf [] Nothing --- >>> let subproblem = LocalProblem 2 stop algorithm --- >>> let x0 = fromList [5, 10] --- >>> minimizeLocal subproblem x0 --- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) --- >>> -- define constraint function: --- >>> let constraintf x = sum (toList x) - 1.0 --- >>> -- define constraint object to pass to the algorithm: --- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 --- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem) --- >>> minimizeAugLag problem x0 --- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED}) - -minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution -minimizeAugLag prob x0 = - unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler - where - handler :: NloptException -> IO (Either N.Result a) - handler (NloptException retcode) = return $ Left retcode +{-# LANGUAGE BangPatterns #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.SRTree.Opt +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to optimize the parameters of an expression. +-- +----------------------------------------------------------------------------- +module Algorithm.SRTree.NonlinearOpt + where + +import Algorithm.SRTree.Likelihoods +import Numeric.Optimization.NLOPT +import Data.Bifunctor (bimap, second) +import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams, countNodes, convertProtectedOps) +import Data.SRTree.Eval +import Algorithm.SRTree.AD + +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed.Mutable as VM +import qualified Data.Vector.Generic as G + +import qualified Data.IntMap.Strict as IntMap +import Data.SRTree.Recursion +import Control.Monad.State.Strict +import Control.Monad.Identity + +import Debug.Trace + +minimizeNLLWith :: (VS.Vector Double -> (Double, VS.Vector Double)) -> (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Int -> Target -> (Target, Double, Int) +minimizeNLLWith funAndGrad alg niter t0 + | niter == 0 = (t0, f, 0) + | n == 0 = (t0, f, 0) + | otherwise = (t_opt', fst (funAndGrad t_opt), nEvs) + where + t0' = G.convert t0 + n = V.length t0 + + (f, _) = funAndGrad t0' -- if there's no parameter or no iterations + + algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) + stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + (t_opt, nEvs) = case minimizeLocal problem t0' of + Right sol -> (solutionParams sol, nEvals sol) + Left e -> (t0', 0) + t_opt' = G.convert t_opt +{-# INLINE minimizeNLLWith #-} + +-- | minimizes the negative log-likelihood of the expression +minimizeNLL' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Target -> (Target, Double, Int) +minimizeNLL' alg backend dist mYerr niter xss ys tree t0 = minimizeNLLWith funAndGrad alg niter t0 + where + m = V.length ys + tree' = buildLoss dist (fromIntegral m) tree + funAndGrad = compileFunAndGrad backend xss ys mYerr tree' + + +minimizeNLL :: ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Target -> (Target, Double, Int) +minimizeNLL = minimizeNLL' TNEWTON + +minimizeNLLWithFixedParam' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Int -> Target -> Target +minimizeNLLWithFixedParam' alg backend dist mYerr' niter xss' ys' tree ix t0 = result + where + m = V.length ys' + tree' = buildLoss dist (fromIntegral m) tree + fixedVal = t0 V.! ix + p = V.length t0 + + evalFull = compileFunAndGrad backend xss' ys' mYerr' tree' + + wrapRed thRed = let (lo, hi) = VS.splitAt ix thRed + in (lo `VS.snoc` fixedVal) VS.++ hi + unwrapRed th = let (lo, hi) = VS.splitAt ix th + in lo VS.++ VS.tail hi + + wrap thRed = let (lo, hi) = V.splitAt ix thRed in (lo `V.snoc` fixedVal) V.++ hi + unwrap th = let (lo, hi) = V.splitAt ix th in lo V.++ V.tail hi + + fgRed :: VS.Vector Double -> (Double, VS.Vector Double) + fgRed thRed = + let thFull = wrapRed thRed + (nll, gradFull) = evalFull thFull + gradRed = unwrapRed gradFull + in (nll, gradRed) + + t0Red = unwrap t0 + (tRawRed,_,_) = minimizeNLLWith fgRed alg niter t0Red + result = wrap tRawRed + +minimizeNLLWithFixedParam = minimizeNLLWithFixedParam' TNEWTON + diff --git a/src/Algorithm/SRTree/Opt.hs b/src/Algorithm/SRTree/Opt.hs deleted file mode 100644 --- a/src/Algorithm/SRTree/Opt.hs +++ /dev/null @@ -1,135 +0,0 @@ -{-# LANGUAGE BangPatterns #-} ------------------------------------------------------------------------------ --- | --- Module : Algorithm.SRTree.Opt --- Copyright : (c) Fabricio Olivetti 2021 - 2024 --- License : BSD3 --- Maintainer : fabricio.olivetti@gmail.com --- Stability : experimental --- Portability : ConstraintKinds --- --- Functions to optimize the parameters of an expression. --- ------------------------------------------------------------------------------ -module Algorithm.SRTree.Opt - where - -import Algorithm.SRTree.Likelihoods -import Algorithm.SRTree.NonlinearOpt -import Data.Bifunctor (bimap, second) -import Data.Massiv.Array -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 -import Data.SRTree.Recursion -import Algorithm.EqSat.Egraph hiding ( size ) -import Algorithm.EqSat.Build -import Control.Monad.State.Strict -import Control.Monad.Identity -import Algorithm.SRTree.AD (evalCache) - -import Debug.Trace - --- | minimizes the negative log-likelihood of the expression -minimizeNLLEGraph :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> EGraph -> EClassId -> ECache -> PVector -> (PVector, Double, Int, ECache) -minimizeNLLEGraph alg dist mYerr niter xss ys egraph root cache t0 - | niter == 0 = (t0, f, 0, cache') - | n == 0 = (t0, f, 0, cache') - | otherwise = (t_opt', fst aa, nEvs, cache') -- (t_opt', nll dist mYerr xss ys tree t_opt', nEvs, cache') - where - (rt, eg) = buildNLLEGraph dist (fromIntegral m) egraph root -- convertProtectedOps - t0' = toStorableVector t0 - (Sz n) = size t0 - (Sz m) = size ys - tree = runIdentity $ getBestExpr root `evalStateT` egraph - aa = gradNLLEGraph dist xss ys mYerr eg cache' rt t_opt - - funAndGrad = gradNLLEGraph dist xss ys mYerr eg cache' rt - (f, _) = gradNLLEGraph dist xss ys mYerr eg cache' rt t0' -- if there's no parameter or no iterations - cache' = evalCache xss egraph cache root t0' - - - algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) - stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)] - problem = LocalProblem (fromIntegral n) stop algorithm - (t_opt, nEvs) = case minimizeLocal problem t0' of - Right sol -> (solutionParams sol, nEvals sol) - Left e -> (t0', 0) - t_opt' = fromStorableVector compMode t_opt - - --- | minimizes the negative log-likelihood of the expression -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' = buildNLL dist (fromIntegral m) 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 = gradNLLGraph dist xss ys mYerr tree' -- second (toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix - - (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 = 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, 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' :: (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 - | otherwise = t_opt' - where - tree' = buildNLL dist (fromIntegral m) tree -- relabelParams - t0' = toStorableVector t0 - treeArr = IntMap.toAscList $ tree2arr tree' - j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..] - (Sz n) = size t0 - (Sz m) = size ys - setTo0 = (VS.// [(ix, 0.0)]) - funAndGrad = second (setTo0 . toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix - - (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, 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, Int) -minimizeGaussian = minimizeNLL Gaussian Nothing - --- | minimizes using Binomial likelihood -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, Int) -minimizePoisson = minimizeNLL Poisson Nothing diff --git a/src/Algorithm/SRTree/Utils.hs b/src/Algorithm/SRTree/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Utils.hs @@ -0,0 +1,320 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +module Algorithm.SRTree.Utils where + +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import Control.Monad +import Control.Monad.Catch +import Control.Monad.Primitive +import Control.Monad.IO.Class +import System.IO.Unsafe + +-- taken from https://hackage.haskell.org/package/cubicspline-0.1.2 +import Control.Arrow +import Data.List (unfoldr) + +import Data.SRTree.Eval +import Debug.Trace (traceShow) + +-- | Internal helper to get dimensions (rows, columns) +matSize :: Columns -> (Int, Int) +matSize [] = (0, 0) +matSize cs@(c:_) = (U.length c, length cs) + +getRows :: Columns -> [Target] +getRows mtx + | n == 0 = [] + | otherwise = [ U.fromListN n [ c U.! i | c <- mtx ] | i <- [0 .. m - 1] ] + where (m, n) = matSize mtx +{-# INLINE getRows #-} + +getCols :: Columns -> [Target] +getCols = id +{-# INLINE getCols #-} + +appendRow :: MonadThrow m => Columns -> Target -> m Columns +appendRow xs v = pure $ zipWith U.snoc xs (U.toList v) +{-# INLINE appendRow #-} + +appendCol :: MonadThrow m => Columns -> Target -> m Columns +appendCol xs v = pure $ xs ++ [v] +{-# INLINE appendCol #-} + +updateS :: Target -> [(Int, Double)] -> Target +updateS vec new = vec U.// new + +linSpace :: Int -> (Double, Double) -> [Double] +linSpace num (lo, hi) = Prelude.take num $ iterate (\x -> x + step) lo + where step = (hi - lo) / (fromIntegral num - 1) +{-# INLINE linSpace #-} + +outer :: (MonadThrow m) => Target -> Target -> m Columns +outer arr1 arr2 + | U.null arr1 || U.null arr2 = pure [] + | otherwise = pure [ U.map (* (arr2 U.! j)) arr1 | j <- [0 .. U.length arr2 - 1] ] +{-# INLINE outer #-} + +-- | Flatten list of column vectors to a row-major U.Vector Double +toRowMajor :: Columns -> U.Vector Double +toRowMajor cols = U.generate (m * n) (\ix -> let (i, j) = ix `divMod` n in (cols !! j) U.! i) + where (m, n) = matSize cols + +-- | Restore a row-major continuous U.Vector Double back to Columns +fromRowMajor :: Int -> Int -> U.Vector Double -> Columns +fromRowMajor m n vec = [ U.generate m (\i -> vec U.! (i * n + j)) | j <- [0 .. n - 1] ] + +unsafeRead :: PrimMonad m => Int -> UM.MVector (PrimState m) Double -> (Int, Int) -> m Double +unsafeRead stride arr (i, j) = UM.unsafeRead arr (i * stride + j) +{-# INLINE unsafeRead #-} + +unsafeWrite :: PrimMonad m => Int -> UM.MVector (PrimState m) Double -> (Int, Int) -> Double -> m () +unsafeWrite stride arr (i, j) val = UM.unsafeWrite arr (i * stride + j) val +{-# INLINE unsafeWrite #-} + +det :: Columns -> Double +det mtx + | m == 0 || n == 0 = 1 + | otherwise = (^2) $ product [ (toRowMajor l) U.! (i * n + i) | i <- [0 .. m - 1] ] + where + (m, n) = matSize mtx + (l, _) = unsafePerformIO (lu mtx) + +detChol :: Columns -> Double +detChol mtx + | m == 0 || n == 0 = 1 + | otherwise = (^2) $ product [ (toRowMajor cho) U.! (i * m + i) | i <- [0 .. m - 1] ] + where + (m, n) = matSize mtx + cho = unsafePerformIO (cholesky mtx) +{-# INLINE det #-} + +rangedLinearDotProd :: PrimMonad m => Int -> Int -> Int -> UM.MVector (PrimState m) Double -> m Double +rangedLinearDotProd r1 r2 len arr = go 0 0 + where + go !acc k + | k < len = do + x <- UM.unsafeRead arr (r1 + k) + y <- UM.unsafeRead arr (r2 + k) + go (acc + x * y) (k + 1) + | otherwise = pure acc +{-# INLINE rangedLinearDotProd #-} + +data NegDef = NegDef deriving Show +instance Exception NegDef + +cholesky :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m Columns +cholesky arr + | m /= n = error $ "cholesky dimension mismatch " <> show m <> " X " <> show n + | m == 0 = pure [] + | otherwise = do + l <- UM.new (m * m) + let orig = toRowMajor arr + forM_ [0 .. m - 1] $ \i -> + forM_ [0 .. m - 1] $ \j -> + if i < j then unsafeWrite m l (i, j) 0 + else do + let cur = orig U.! (i * m + j) + rowI = i * m + rowJ = j * m + xjj <- UM.unsafeRead l (rowJ + j) + tot <- rangedLinearDotProd rowI rowJ j l + let delta = cur - tot + if i == j + then if delta <= 0 + then throwM NegDef + else UM.unsafeWrite l (rowI + j) (sqrt delta) + else UM.unsafeWrite l (rowI + j) (delta / xjj) + frozen <- U.unsafeFreeze l + pure $ fromRowMajor m m frozen + where (m, n) = matSize arr +{-# INLINE cholesky #-} + +invChol :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m Columns +invChol arr = do + lMtx <- cholesky arr + let (m, _) = matSize arr + mtx <- U.thaw (toRowMajor lMtx) + forM_ [0 .. m - 1] $ \i -> do + lII <- unsafeRead m mtx (i, i) + unsafeWrite m mtx (i, i) (1 / lII) + forM_ [0 .. i - 1] $ \j -> do + tot <- rangedLinearDotProd (i * m + j) (j * m + j) (i - j) mtx + unsafeWrite m mtx (j, i) ((-tot) / lII) + unsafeWrite m mtx (i, j) 0 + + mm <- UM.replicate (m * m) 0 + forM_ [0 .. m - 1] $ \i -> do + dii <- rangedLinearDotProd (i * m + i) (i * m + i) (m - i) mtx + unsafeWrite m mm (i, i) dii + forM_ [i + 1 .. m - 1] $ \j -> do + dij <- rangedLinearDotProd (i * m + j) (j * m + j) (m - j) mtx + unsafeWrite m mm (i, j) dij + unsafeWrite m mm (j, i) dij + frozen <- U.unsafeFreeze mm + pure $ fromRowMajor m m frozen +{-# INLINE invChol #-} + +lu :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m (Columns, Columns) +lu mtx = do + let (m, n) = matSize mtx + orig = toRowMajor mtx + u <- UM.replicate (m * n) 0 + forM_ [0 .. min m n - 1] $ \i -> unsafeWrite n u (i, i) 1 + l <- UM.replicate (m * n) 0 + + let buildLVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do + lik <- unsafeRead n l (i, k) + ukj <- unsafeRead n u (k, j) + go (k+1) (s + lik * ukj) + s' <- go 0 0 + unsafeWrite n l (i, j) ((orig U.! (i * n + j)) - s') + + buildL !i !j = when (i /= m) $ do + buildLVal i j + buildL (i+1) j + + buildUVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do + ljk <- unsafeRead n l (j, k) + uki <- unsafeRead n u (k, i) + go (k+1) (s + ljk * uki) + s' <- go 0 0 + ljj <- unsafeRead n l (j, j) + unsafeWrite n u (j, i) (((orig U.! (j * n + i)) - s') / ljj) + + buildU !i !j = when (i /= n) $ do + buildUVal i j + buildU (i+1) j + + buildLU !j = when (j /= n && j /= m) $ do + buildL j j + buildU j j + buildLU (j+1) + + buildLU 0 + finalL <- U.unsafeFreeze l + finalU <- U.unsafeFreeze u + pure (fromRowMajor m n finalL, fromRowMajor m n finalU) + +forwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +forwardSub a b = do + let m = U.length b + n = length a + aMat = toRowMajor a + x <- UM.replicate m 0 + let coeff !i !j !s + | j == i = pure s + | otherwise = do + let aij = aMat U.! (i * n + j) + xj <- UM.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i /= m) $ do + let bi = b U.! i + aii = aMat U.! (i * n + i) + c <- coeff i 0 0 + UM.unsafeWrite x i ((bi - c) / aii) + go (i+1) + go 0 + U.unsafeFreeze x + +backwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +backwardSub a b = do + let m = U.length b + n = length a + aMat = toRowMajor a + x <- UM.replicate m 0 + let coeff !i !j !s + | j == m = pure s + | otherwise = do + let aij = aMat U.! (i * n + j) + xj <- UM.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i >= 0) $ do + let bi = b U.! i + aii = aMat U.! (i * n + i) + c <- coeff i (i+1) 0 + UM.unsafeWrite x i ((bi - c) / aii) + go (i-1) + go (m-1) + U.unsafeFreeze x + +luSolve :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +luSolve a b = do + (l, u) <- lu a + forwardSub l b >>= backwardSub u + +type PolyCos = (Double, Double, Double) + +cubicSplineCoefficients :: [(Double, Double)] -> [PolyCos] +cubicSplineCoefficients xs = Prelude.zip3 x y z' + where + x = map fst xs + y = map snd xs + xdiff = zipWith (-) (tail x) x + xdiff' = U.fromList xdiff + + dydx :: U.Vector Double + dydx = U.fromList $ Prelude.zipWith3 (\y0 y1 xd -> (y0 - y1) / xd) (tail y) y xdiff + + n = length x + + w :: [Double] + w = 0 : nextW 1 w + where + nextW ix (wi : t) + | ix == n - 1 = [] + | otherwise = + let m = (xdiff' U.! (ix - 1)) * (2 - wi) + 2 * (xdiff' U.! ix) + wn = (xdiff' U.! ix) / m + in wn : nextW (ix + 1) t + + z :: [Double] + z = 0 : nextZ 1 z + where + nextZ ix (zi : t) + | ix == n - 1 = [0] + | otherwise = + let m = (xdiff' U.! (ix - 1)) * (2 - (w !! (ix - 1))) + 2 * (xdiff' U.! ix) + zn = (6 * ((dydx U.! ix) - (dydx U.! (ix - 1))) - (xdiff' U.! (ix - 1)) * zi) / m + in zn : nextZ (ix + 1) t + + z' :: [Double] + z' = Prelude.reverse $ 0 : [z !! i - w !! i * z !! (i + 1) | i <- [n - 2, n - 3 .. 0]] + +chunkBy :: Int -> [t] -> [[t]] +chunkBy n = unfoldr go + where + go [] = Nothing + go x = Just $ splitAt n x + +genSplineFun :: [(Double, Double)] -> Double -> Double +genSplineFun pts x + | length xs < 2 = x + | x < head xs = y1 + (x - x1) * (y2 - y1) / (x2 - x1) + | x > last xs = y_1 + (x - x_1) * (y_n - y_1) / (x_n - x_1) + | otherwise = go xs $ zip coefs (tail coefs) + where + xs = map fst pts + ys = map snd pts + coefs = cubicSplineCoefficients pts + x1 = head xs; y1 = head ys + x2 = xs !! 1; y2 = ys !! 1 + x_1 = xs !! (len - 2); y_1 = ys !! (len - 2) + x_n = last xs; y_n = last ys + len = length xs + + evalAt (a1, b1, c1) (a2, b2, c2) y = + let hi1 = a2 - a1 + in c1 / (6 * hi1) * (a2 - y)^3 + c2 / (6 * hi1) * (y - a1)^3 + + (b2 / hi1 - c2 * hi1 / 6) * (y - a1) + (b1 / hi1 - c1 * hi1 / 6) * (a2 - y) + + go [x1, x2] [(c1, c2)] = evalAt c1 c2 x + go (x1 : x2 : xs') ((c1, c2) : cs) + | x >= x1 && x <= x2 = evalAt c1 c2 x + | otherwise = go (x2 : xs') cs diff --git a/src/Data/SRTree/Datasets.hs b/src/Data/SRTree/Datasets.hs --- a/src/Data/SRTree/Datasets.hs +++ b/src/Data/SRTree/Datasets.hs @@ -1,6 +1,11 @@ {-# language ImportQualifiedPost #-} {-# language ViewPatterns #-} {-# language OverloadedStrings #-} +{-# language BlockArguments #-} +{-# language ExplicitForAll #-} +{-# language BangPatterns #-} +{-# language LambdaCase #-} +{-# language RankNTypes, ScopedTypeVariables #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Datasets @@ -14,41 +19,119 @@ -- this module exports only the `loadDataset` function. -- ----------------------------------------------------------------------------- -module Data.SRTree.Datasets ( loadDataset, loadTrainingOnly, getX, splitData, DataSet(..) ) +module Data.SRTree.Datasets ( loadDataset, loadTrainingOnly, getX, splitData, DataSet(..), splitFileNameParams, getRows, getColumns ) where import Codec.Compression.GZip (decompress) import Data.ByteString.Char8 qualified as B import Data.ByteString.Lazy qualified as BS import Data.List (delete, find, intercalate) -import Data.Massiv.Array - ( Array, - Comp (Seq, Par), - Ix2 ((:.)), - S (..), - Sz (Sz1), - ( Array S Ix2 Double -loadMtx = M.fromLists' compMode . map (map (read . B.unpack)) +loadMtx :: [[B.ByteString]] -> [Vector Double] +loadMtx [] = [] +loadMtx rows = map V.fromList + $ foldr (zipWith (:) . map parseDouble) (replicate ncols []) rows + where ncols = length (head rows) {-# INLINE loadMtx #-} +-- | Powers of ten as exact 'Integer's, precomputed once and shared by every +-- 'parseDouble' call. The per-value @10 ^ k@ exponentiation previously ran a +-- growing-Integer multiply loop on every parsed number, which showed up as a +-- measurable chunk of the corpus-load allocation. The table is the exact same +-- integer, so conversions stay bit-identical. +maxPow10 :: Int +maxPow10 = 400 + +pow10 :: VB.Vector Integer +pow10 = VB.generate (maxPow10 + 1) (\k -> 10 ^ k) +{-# NOINLINE pow10 #-} + +-- | @10^k@ as an exact 'Integer'; falls back to direct exponentiation for +-- exponents beyond the precomputed range (only reachable with absurd inputs). +pow10E :: Int -> Integer +pow10E k | k >= 0 && k <= maxPow10 = VB.unsafeIndex pow10 k + | otherwise = 10 ^ k +{-# INLINE pow10E #-} + +-- | Fast decimal double parser over a 'B.ByteString'. Handles an optional +-- sign, a fractional part and an optional 'e'/'E' exponent. The mantissa is +-- accumulated exactly as an 'Integer' and converted to 'Double' through a +-- single 'fromRational', which matches the correctly-rounded result of 'read'. +-- Falls back to 'read' (the slow Show-derived parser) for anything it can't +-- parse (NaN, Infinity, hex floats, etc.), so behavior is unchanged for odd +-- input. +parseDouble :: B.ByteString -> Double +parseDouble bs = case go 0 1 0 False 0 of + Just (m, s, nd, e) + -- when e >= nd the rational m * 10^e / 10^nd is an exact integer, so a + -- single fromInteger is bit-identical to fromRational (which would only + -- gcd-reduce it) but skips the rational machinery entirely. + | e >= nd -> fromInteger (s * (m * pow10E (e - nd))) + -- otherwise the value is m / 10^(nd-e); keep fromRational so the single + -- rounding matches `read` exactly (a Double division by a rounded power + -- of ten would be off by up to an ulp). + | otherwise -> fromRational (s * m % (pow10E (nd - e))) + Nothing -> read (B.unpack bs) + where + n = B.length bs + -- i: index, sgn: +/-1, acc: accumulated mantissa digits (exact Integer), + -- dot: whether a '.' has been seen, nd: number of digits following the + -- decimal point, expo: signed integer exponent from the 'e' tail + go :: Int -> Integer -> Integer -> Bool -> Int -> Maybe (Integer, Integer, Int, Int) + go !i !sgn !acc !dot !nd + | i >= n = Just (acc, sgn, nd, 0) + | otherwise = + let c = fromEnum (B.index bs i) + in case c of + 45 -> if i == 0 then go (i+1) (-sgn) acc dot nd else Nothing -- '-' + 43 -> if i == 0 then go (i+1) sgn acc dot nd else Nothing -- '+' + 46 -> if dot then Nothing else go (i+1) sgn acc True nd -- '.' + _ | c >= 48 && c <= 57 -> + let d = fromIntegral (c - 48) :: Integer + nd' = if dot then nd + 1 else nd + in go (i+1) sgn (acc * 10 + d) dot nd' + | (c == 101 || c == 69) && i > 0 -> -- 'e' / 'E' + parseExp (i+1) sgn acc dot nd + | otherwise -> Nothing + -- parse the (optional) exponent tail: an optional sign then digits + parseExp :: Int -> Integer -> Integer -> Bool -> Int -> Maybe (Integer, Integer, Int, Int) + parseExp !i !sgn !acc !dot !nd + | i >= n = Just (acc, sgn, nd, 0) + | otherwise = + let c = fromEnum (B.index bs i) + in case c of + 45 -> expDig (i+1) sgn acc dot nd (-1) 0 -- '-' + 43 -> expDig (i+1) sgn acc dot nd 1 0 -- '+' + _ -> expDig i sgn acc dot nd 1 0 + where + -- es: exponent sign (+/-1); e: accumulated exponent magnitude + expDig :: Int -> Integer -> Integer -> Bool -> Int -> Int -> Int -> Maybe (Integer, Integer, Int, Int) + expDig !i !sgn !acc !dot !nd !es !e + | i >= n = Just (acc, sgn, nd, es * e) + | otherwise = + let c = fromEnum (B.index bs i) + in if c >= 48 && c <= 57 + then expDig (i+1) sgn acc dot nd es (e * 10 + fromIntegral (c - 48)) + else Nothing + -- | Returns true if the extension is .gz isGZip :: FilePath -> Bool isGZip = (== ".gz") . takeExtension @@ -83,7 +166,7 @@ -- The first row can be a header. readFileToLines :: FilePath -> IO [[B.ByteString]] readFileToLines filename = do - content <- removeBEmpty . toLines . toChar8 . unzip <$> BS.readFile filename + content <- removeBEmpty . toLines . toStrict . unzip <$> BS.readFile filename let sep = getSep content pure . removeEmpty . map (B.split sep) $ content where @@ -92,7 +175,10 @@ removeEmpty = filter (not . null) toLines = B.split '\n' unzip = if isGZip filename then decompress else id - toChar8 = B.pack . map (toEnum . fromEnum) . BS.unpack + -- lazy -> strict without going through a [Word8]/[Char] list (the old + -- B.pack . map toEnum . BS.unpack round trip allocated ~1GB on a 14MB + -- CSV); BS.toStrict is a single O(n) copy. + toStrict = BS.toStrict {-# INLINE readFileToLines #-} -- | Splits the parameters from the filename @@ -105,7 +191,9 @@ -- 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 6 params) +splitFileNameParams (B.pack -> filename) + | B.null filename = ("", replicate 6 B.empty) + | otherwise = (B.unpack fname, take 6 params) where (fname : params') = B.split ':' filename -- fill up the empty parameters with an empty string @@ -155,8 +243,8 @@ getRows :: B.ByteString -> B.ByteString -> Int -> (Int, Int) getRows (B.unpack -> start) (B.unpack -> end) nRows | st_ix >= end_ix = error $ "Invalid range: " <> show start <> ":" <> show end <> "." - | st_ix == 0 && end_ix == nRows-1 = (0, nRows - 1) - | otherwise = (st_ix, end_ix) + | st_ix == 0 && end_ix == nRows-1 = (0, nRows) + | otherwise = (st_ix, end_ix + 1) where st_ix = if null start then 0 @@ -188,7 +276,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), (Maybe PVector, Maybe PVector), String, String) +loadDataset :: FilePath -> Bool -> IO (([Vector Double], Vector Double, [Vector Double], Vector Double), (Maybe (Vector Double), Maybe (Vector Double)), String, String) loadDataset filename hasHeader = do csv <- readFileToLines fname pure $ processData csv params hasHeader @@ -196,7 +284,7 @@ (fname, params) = splitFileNameParams filename -- support function that does everything for loadDataset -processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> ((SRMatrix, PVector, SRMatrix, PVector), (Maybe PVector, Maybe PVector), String, String) +processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> (([Vector Double], Vector Double, [Vector Double], Vector Double), (Maybe (Vector Double), Maybe (Vector Double)), 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 @@ -209,24 +297,24 @@ ] 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 + (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 p = length ixs - x = M.computeAs S $ M.throwEither $ M.stackInnerSlicesM $ map (datum [e] -> [[e]] @@ -238,7 +326,7 @@ build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a] build g = g (:) [] -splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet) +splitData :: DataSet -> Int -> State StdGen (DataSet, DataSet) splitData (x, y, mYErr) k = do if k == 1 then pure ((x, y, mYErr), (x, y, mYErr)) @@ -246,36 +334,83 @@ 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 + let tr_ix = [ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i] + val_ix = [ix | ixs_i <- ixs, let ix = Prelude.head ixs_i] + (x_tr, x_te) = getX tr_ix val_ix x + (y_tr, y_te) = getY tr_ix val_ix y + + mY = fmap (getY tr_ix val_ix) 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 + sz = V.length 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] + getX :: [Int] -> [Int] -> [Vector Double] -> ([Vector Double], [Vector Double]) + getX tr_ix val_ix xs = ( [ V.fromList [x V.! ix | ix <- tr_ix] | x <- xs ] + , [ V.fromList [x V.! ix | ix <- val_ix] | x <- xs ] + ) + getY :: [Int] -> [Int] -> Vector Double -> (Vector Double, Vector Double) + getY tr_ix val_ix ys = ( V.fromList [ys V.! ix | ix <- tr_ix] + , V.fromList [ys V.! ix | ix <- val_ix] ) getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2) getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c) -getX :: DataSet -> SRMatrix +getX :: DataSet -> [Vector Double] getX (a, _, _) = a -getTarget :: DataSet -> PVector +getTarget :: DataSet -> Vector Double getTarget (_, b, _) = b -getError :: DataSet -> Maybe PVector +getError :: DataSet -> Maybe (Vector Double) getError (_, _, c) = c loadTrainingOnly fname b = getTrain <$> loadDataset fname b + +-- | Shuffles a list, taken from list-shuffle +shuffle :: (RandomGen g) => [a] -> g -> ([a], g) +shuffle list gen0 = + runST do + array <- listToMutableArray list + gen1 <- shuffleN (Array.sizeofMutableArray array - 1) array gen0 + array1 <- Array.unsafeFreezeArray array + pure (Foldable.toList array1, gen1) + +listToMutableArray :: forall a s. [a] -> ST s (Array.MutableArray s a) +listToMutableArray list = do + array <- Array.newArray (length list) undefined + let writeElems :: Int -> [a] -> ST s () + writeElems !i = \case + [] -> pure () + x : xs -> do + Array.writeArray array i x + writeElems (i + 1) xs + writeElems 0 list + pure array +{-# INLINE listToMutableArray #-} + +shuffleN :: forall a g s. (RandomGen g) => Int -> Array.MutableArray s a -> g -> ST s g +shuffleN n0 array = + go 0 + where + go :: Int -> g -> ST s g + go !i gen0 + | i >= n = pure gen0 + | otherwise = do + let (j, gen1) = uniformR (i, m) gen0 + swapArrayElems i j array + go (i + 1) gen1 + + n = min n0 m + m = Array.sizeofMutableArray array - 1 +{-# SPECIALIZE shuffleN :: Int -> Array.MutableArray s a -> StdGen -> ST s StdGen #-} + +-- Swap two elements in a mutable array. +swapArrayElems :: Int -> Int -> Array.MutableArray s a -> ST s () +swapArrayElems i j array = do + x <- Array.readArray array i + y <- Array.readArray array j + Array.writeArray array i y + Array.writeArray array j x +{-# INLINE swapArrayElems #-} diff --git a/src/Data/SRTree/Derivative.hs b/src/Data/SRTree/Derivative.hs --- a/src/Data/SRTree/Derivative.hs +++ b/src/Data/SRTree/Derivative.hs @@ -16,6 +16,7 @@ , doubleDerivative , deriveByVar , deriveByParam + , derivOp ) where @@ -115,6 +116,18 @@ doubleDerivative Recip = (*2) . recip . (^3) doubleDerivative Cube = (6*) {-# INLINE doubleDerivative #-} + +-- | Returns (d(Output)/d(Left), d(Output)/d(Right)) +-- used for AD +derivOp :: Op -> Double -> Double -> (Double, Double) +derivOp Add _ _ = (1.0, 1.0) +derivOp Sub _ _ = (1.0, -1.0) +derivOp Mul v1 v2 = (v2, v1) +derivOp Div v1 v2 = (1.0 / v2, -(v1) / (v2 * v2)) +-- e.g., Power: d(x^y)/dx = y*x^(y-1), d(x^y)/dy = x^y * ln(x) +derivOp Power v1 v2 = (v2 * (v1 ** (v2 - 1)), (v1 ** v2) * log v1) +derivOp _ _ _ = (0.0, 0.0) -- Add remaining ops +{-# INLINE derivOp #-} -- | Symbolic derivative by a variable deriveByVar :: Int -> Fix SRTree -> Fix SRTree diff --git a/src/Data/SRTree/Eval.hs b/src/Data/SRTree/Eval.hs --- a/src/Data/SRTree/Eval.hs +++ b/src/Data/SRTree/Eval.hs @@ -1,4 +1,5 @@ -{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE LambdaCase, BangPatterns #-} + ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Eval @@ -13,8 +14,7 @@ ----------------------------------------------------------------------------- {-# LANGUAGE FlexibleInstances #-} module Data.SRTree.Eval - ( evalTree - , evalOp + ( evalOp , evalFun , cbrt , inverseFunc @@ -23,73 +23,249 @@ , invright , invleft , replicateAs - , SRVector, PVector, SRMatrix - , compMode + , Target, Theta, Columns + , compile + , compileLoss ) where -import Data.Massiv.Array -import qualified Data.Massiv.Array as M import Data.SRTree.Internal import Data.SRTree.Recursion (Fix (..), cata) +import Data.Vector.Unboxed (Vector) +import qualified Data.Vector.Unboxed as V +import Control.Monad.ST (runST) +import qualified Data.Vector as VB -- Boxed vector for instructions +import qualified Data.Vector.Unboxed.Mutable as VM +import Control.Concurrent.Async (forConcurrently_) +import System.IO.Unsafe (unsafePerformIO) +import Control.Concurrent (getNumCapabilities) +import Data.Maybe (fromJust) -- | Vector of target values -type SRVector = M.Array D Ix1 Double +type Target = Vector Double -- | Vector of parameter values. Needs to be strict to be readily accesible. -type PVector = M.Array S Ix1 Double +type Theta = Vector Double -- | Matrix of features values -type SRMatrix = M.Array S Ix2 Double +type Columns = [Vector Double] -compMode :: M.Comp -compMode = M.Seq +-- A multi-threaded replacement for V.sum +sumParallel :: Int -> (Int -> Double) -> Double +sumParallel n f = unsafePerformIO $ do + numThreads <- getNumCapabilities + let chunkSize = n `quot` numThreads + -- 1. Allocate a single block of unboxed memory EXACTLY ONCE + out <- VM.unsafeNew numThreads + + -- 2. Spawn threads. Each thread gets a unique ID and a slice of memory. + forConcurrently_ [0 .. numThreads - 1] $ \tId -> do + let !start = tId * chunkSize + -- The last thread cleans up the remainder + !end = if tId == numThreads - 1 then n else start + chunkSize + + -- 3. The inner thread loop. Strict, unboxed, and bounds-check free. + let loop !i !acc + | i >= end = return acc + | otherwise = loop (i + 1) (acc + f i) + + total <- loop start 0.0 + VM.unsafeWrite out tId total + + -- 4. Instantly cast the mutable memory to an immutable Vector (O(1) cost) + totals <- V.unsafeFreeze out + return (V.sum totals) +{-# NOINLINE sumParallel #-} + -- Improve quality of life with Num and Floating instances for our matrices -instance Index ix => Num (M.Array D ix Double) where - (+) = (!+!) - (-) = (!-!) - (*) = (!*!) - abs = absA - signum = signumA - fromInteger = fromInteger - negate = negateA +instance Num Target where + (+) = V.zipWith (+) + (-) = V.zipWith (-) + (*) = V.zipWith (*) + abs = V.map abs + signum = V.map signum + fromInteger = V.singleton . fromInteger + negate = V.map negate -instance Index ix => Floating (M.Array D ix Double) where - pi = pi - exp = expA - log = logA - sqrt = sqrtA - sin = sinA - cos = cosA - tan = tanA - asin = asinA - acos = acosA - atan = atanA - sinh = sinhA - cosh = coshA - tanh = tanhA - asinh = asinhA - acosh = acoshA - atanh = atanhA - (**) = (.**) -instance Index ix => Fractional (M.Array D ix Double) where - fromRational = fromRational - (/) = (!/!) - recip = recipA +instance Floating Target where + pi = V.singleton pi + exp = V.map exp + log = V.map log + sqrt = V.map sqrt + sin = V.map sin + cos = V.map cos + tan = V.map tan + asin = V.map asin + acos = V.map acos + atan = V.map atan + sinh = V.map sinh + cosh = V.map cosh + tanh = V.map tanh + asinh = V.map asinh + acosh = V.map acosh + atanh = V.map atanh + (**) = V.zipWith (**) +instance Fractional Target where + fromRational = V.singleton . fromRational + (/) = V.zipWith (/) + recip = V.map recip +-- We change the Dynamic type to evaluate a single scalar at a specific row index (Int) +data Staged = + Scl Double + | Static (Vector Double) + | Dynamic (Vector Double -> Int -> Double) -- (Theta -> RowIndex -> Result) + +-- A multi-threaded replacement for V.generate +generateParallel :: Int -> (Int -> Double) -> V.Vector Double +generateParallel n f = unsafePerformIO $ do + numThreads <- getNumCapabilities + let chunkSize = n `quot` numThreads + + -- 1. Allocate a single block of unboxed memory EXACTLY ONCE + out <- VM.unsafeNew n + + -- 2. Spawn threads. Each thread gets a unique ID and a slice of memory. + forConcurrently_ [0 .. numThreads - 1] $ \tId -> do + let !start = tId * chunkSize + -- The last thread cleans up the remainder + !end = if tId == numThreads - 1 then n else start + chunkSize + + -- 3. The inner thread loop. Strict, unboxed, and bounds-check free. + let loop !i + | i >= end = return () + | otherwise = do + -- Write directly to the shared memory pointer + VM.unsafeWrite out i (f i) + loop (i + 1) + + loop start + + -- 4. Instantly cast the mutable memory to an immutable Vector (O(1) cost) + V.unsafeFreeze out +{-# NOINLINE generateParallel #-} + +compileLoss :: [Vector Double] -> Fix SRTree -> Target -> Maybe Target -> (Vector Double -> Double) +compileLoss dataset tree y mYerr = + case cata alg tree of + Scl c -> \_ -> V.sum $ V.replicate n c + Static v -> \_ -> V.sum v + -- We only allocate memory EXACTLY ONCE here at the top level + --Dynamic f -> \th -> V.generate n (f th) + Dynamic f -> \th -> V.sum (V.generate n (f th)) + where + n = V.length (head dataset) + yErr = fromJust mYerr + + alg :: SRTree Staged -> Staged + + -- 1. Base Cases + alg (Const c) = Scl c + alg (Var (-1)) = Static y + alg (Var (-2)) = Static yErr + alg (Var i) = Static (dataset !! i) + alg (Param i) = Dynamic (\th !idx -> th `V.unsafeIndex` i) + + -- 2. Univariate Functions + alg (Uni f (Scl c)) = Scl (evalFun f c) + alg (Uni f (Static v)) = Static (V.map (evalFun f) v) + + -- We map the function over the scalar result of the inner closure + alg (Uni f (Dynamic g)) = let !rawFun = evalFun f in Dynamic (\th !i -> rawFun (g th i)) + + -- 3. Binary Functions + alg (Bin op (Scl c1) (Scl c2)) = Scl (evalOp op c1 c2) + alg (Bin op (Scl c) (Static v)) = Static (V.map (evalOp op c) v) + alg (Bin op (Static v) (Scl c)) = Static (V.map (\c2 -> evalOp op c2 c) v) + alg (Bin op (Static v1) (Static v2)) = Static (V.zipWith (evalOp op) v1 v2) + + -- 4. Dynamic Combinations (The Core Optimization) + + alg (Bin op (Scl c) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp c (g th i)) + + alg (Bin op (Dynamic g) (Scl c)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) c) + + -- When combining a Static array with a Dynamic closure, + -- we use unsafeIndex to fetch the static value at row 'i' directly. + alg (Bin op (Static v) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (v `V.unsafeIndex` i) (g th i)) + + alg (Bin op (Dynamic g) (Static v)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) (v `V.unsafeIndex` i)) + + alg (Bin op (Dynamic g1) (Dynamic g2)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g1 th i) (g2 th i)) + + +compile :: [Vector Double] -> Fix SRTree -> (Vector Double -> Vector Double) +compile dataset tree = + case cata alg tree of + Scl c -> \_ -> V.replicate n c + Static v -> \_ -> v + -- We only allocate memory EXACTLY ONCE here at the top level + --Dynamic f -> \th -> V.generate n (f th) + Dynamic f -> \th -> V.generate n (f th) + where + n = V.length (head dataset) + + alg :: SRTree Staged -> Staged + + -- 1. Base Cases + alg (Const c) = Scl c + alg (Var i) = Static (dataset !! i) + -- Look at this! No more V.replicate. It just fetches the scalar directly. + alg (Param i) = Dynamic (\th !idx -> th `V.unsafeIndex` i) + alg (Y i) = undefined -- this shouldn't be called + + -- 2. Univariate Functions + alg (Uni f (Scl c)) = Scl (evalFun f c) + alg (Uni f (Static v)) = Static (V.map (evalFun f) v) + + -- We map the function over the scalar result of the inner closure + alg (Uni f (Dynamic g)) = let !rawFun = evalFun f in Dynamic (\th !i -> rawFun (g th i)) + + -- 3. Binary Functions + alg (Bin op (Scl c1) (Scl c2)) = Scl (evalOp op c1 c2) + alg (Bin op (Scl c) (Static v)) = Static (V.map (evalOp op c) v) + alg (Bin op (Static v) (Scl c)) = Static (V.map (\c2 -> evalOp op c2 c) v) + alg (Bin op (Static v1) (Static v2)) = Static (V.zipWith (evalOp op) v1 v2) + + -- 4. Dynamic Combinations (The Core Optimization) + + alg (Bin op (Scl c) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp c (g th i)) + + alg (Bin op (Dynamic g) (Scl c)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) c) + + -- When combining a Static array with a Dynamic closure, + -- we use unsafeIndex to fetch the static value at row 'i' directly. + alg (Bin op (Static v) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (v `V.unsafeIndex` i) (g th i)) + + alg (Bin op (Dynamic g) (Static v)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) (v `V.unsafeIndex` i)) + + alg (Bin op (Dynamic g1) (Dynamic g2)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g1 th i) (g2 th i)) + + -- 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 +replicateAs :: Columns -> Double -> Target +replicateAs xss c = let m = V.length (head xss) in V.replicate 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 +evalTree :: Columns -> Theta -> Fix SRTree -> Target evalTree xss params = cata $ \case - Var ix -> xss replicateAs xss $ params ! ix - Const c -> replicateAs xss c - Uni g t -> evalFun g t - Bin op l r -> evalOp op l r + Var ix -> xss !! ix + Param ix -> replicateAs xss $ params V.! ix + Const c -> replicateAs xss c + Y _ -> undefined + Uni g t -> evalFun g t + Bin op l r -> evalOp op l r {-# INLINE evalTree #-} -- evaluates an operator @@ -213,3 +389,5 @@ invertibles :: [Function] invertibles = [Id, Sin, Cos, Tan, Tanh, ASin, ACos, ATan, ATanh, Sqrt, Square, Log, Exp, Recip] {-# INLINE invertibles #-} + + diff --git a/src/Data/SRTree/Internal.hs b/src/Data/SRTree/Internal.hs --- a/src/Data/SRTree/Internal.hs +++ b/src/Data/SRTree/Internal.hs @@ -3,6 +3,7 @@ {-# language RankNTypes #-} {-# language OverloadedStrings #-} {-# language LambdaCase #-} +{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Internal @@ -56,22 +57,25 @@ import Text.Read (readMaybe) import qualified Data.IntMap as IntMap import Data.List ( nub ) +import GHC.Generics (Generic) +import Control.DeepSeq (NFData) -- | Tree structure to be used with Symbolic Regression algorithms. -- This structure is a fixed point of a n-ary tree. data SRTree val = - Var Int -- ^ index of the variables - | Param Int -- ^ index of the parameter - | Const Double -- ^ constant value, can be converted to a parameter + Var {-# UNPACK #-} !Int -- ^ index of the variables + | Param {-# UNPACK #-} !Int -- ^ index of the parameter + | Const {-# UNPACK #-} !Double -- ^ constant value, can be converted to a parameter + | Y {-# UNPACK #-} !Int -- ^ index of the target variable, always 0 for now -- | IConst Int -- TODO: integer constant -- | RConst Ratio -- TODO: rational constant | Uni Function val -- ^ univariate function | Bin Op val val -- ^ binary operator - deriving (Show, Eq, Ord, Functor) + deriving (Show, Eq, Ord, Functor, Generic, NFData) -- | Supported operators data Op = Add | Sub | Mul | Div | Power | PowerAbs | AQ - deriving (Show, Read, Eq, Ord, Enum) + deriving (Show, Read, Eq, Ord, Enum, Generic, NFData) -- | Supported functions data Function = @@ -98,7 +102,7 @@ | Exp | Recip | Cube - deriving (Show, Read, Eq, Ord, Enum) + deriving (Show, Read, Eq, Ord, Enum, Generic, NFData) removeProtectedOps :: Fix SRTree -> Fix SRTree removeProtectedOps = cata alg diff --git a/src/Data/SRTree/Random.hs b/src/Data/SRTree/Random.hs --- a/src/Data/SRTree/Random.hs +++ b/src/Data/SRTree/Random.hs @@ -42,9 +42,9 @@ import Data.Maybe (fromJust) import Data.SRTree.Internal import System.Random (Random (random, randomR), StdGen, mkStdGen) -import Data.Massiv.Array as MA hiding (forM_, forM, P) import Data.SRTree.Eval import Control.Monad +import qualified Data.Vector.Unboxed as V -- * Class definition of properties that a certain parameter type has. @@ -205,8 +205,8 @@ 2 -> replaceFixChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) -randomVec :: Monad m => Int -> Rng m PVector -randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1)) +randomVec :: Monad m => Int -> Rng m Theta +randomVec n = V.fromList <$> replicateM n (randomRange (-1, 1)) randomTree :: Monad m => Int -> Int -> Int -> Rng m (Fix SRTree) -> Rng m (SRTree ()) -> Bool -> Rng m (Fix SRTree) randomTree minDepth maxDepth maxSize genTerm genNonTerm grow diff --git a/src/Numeric/Optimization/NLOPT.hs b/src/Numeric/Optimization/NLOPT.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Optimization/NLOPT.hs @@ -0,0 +1,976 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : Numeric.NLOPT +Copyright : (c) Matthew Peddie 2017 +License : BSD3 +Maintainer : Matthew Peddie +Stability : provisional +Portability : GHC + +This module provides a high-level, @hmatrix@-compatible interface to +the library by +Steven G. Johnson. + +NOTE: This is an adaptation from https://hackage.haskell.org/package/hmatrix-nlopt-0.2.0.0 +that removes the dependency to hmatrix and support any Vector Storage. + += Documentation + +Most non-numerical details are documented, but for specific +information on what the optimization methods do, how constraints are +handled, etc., you should consult: + + * The + + * The + + * The + += Example program + +The following interactive session example uses the Nelder-Mead simplex +algorithm, a derivative-free local optimizer, to minimize a trivial +function with a minimum of 22.0 at @(0, 0)@. + +>>> import Numeric.LinearAlgebra ( dot, fromList ) +>>> let objf x = x `dot` x + 22 -- define objective +>>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +>>> let algorithm = NELDERMEAD objf [] Nothing -- specify algorithm +>>> let problem = LocalProblem 2 stop algorithm -- specify problem +>>> let x0 = fromList [5, 10] -- specify initial guess +>>> minimizeLocal problem x0 +Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) + +-} + +module Numeric.Optimization.NLOPT ( + -- * Specifying the objective function + Objective + , ObjectiveD + , Preconditioner + -- * Specifying the constraints + -- ** Bound constraints + , Bounds(..) + -- ** Nonlinear constraints + -- + -- $nonlinearconstraints + + -- *** Constraint functions + , ScalarConstraint + , ScalarConstraintD + , VectorConstraint + , VectorConstraintD + -- *** Constraint types + , Constraint(..) + , EqualityConstraint(..) + , InequalityConstraint(..) + -- *** Collections of constraints + , EqualityConstraints + , EqualityConstraintsD + , InequalityConstraints + , InequalityConstraintsD + -- * Stopping conditions + -- + -- $nonempty + , StoppingCondition(..) + , NonEmpty(..) + -- * Additional configuration + , RandomSeed(..) + , Population(..) + , VectorStorage(..) + , InitialStep(..) + -- * Minimization problems + -- ** Local minimization + , LocalAlgorithm(..) + , LocalProblem(..) + , minimizeLocal + -- ** Global minimization + , GlobalAlgorithm(..) + , GlobalProblem(..) + , minimizeGlobal + -- ** Minimization by augmented Lagrangian + , AugLagAlgorithm(..) + , AugLagProblem(..) + , minimizeAugLag + -- ** Results + , Solution(..) + , N.Result(..) + ) where + +import qualified Numeric.Optimization.NLOPT.Bindings as N + +import Data.List.NonEmpty (NonEmpty(..)) + +import qualified Data.Vector.Storable as V +import Data.Vector.Storable ( Vector ) + +import Control.Exception ( Exception ) +import qualified Control.Exception as Ex +import Data.Typeable ( Typeable ) +import Data.Foldable ( traverse_ ) + +import System.IO.Unsafe ( unsafePerformIO ) + +-- each element i contains a row vec +type Matrix a = [Vector a] + +flatten :: V.Storable a => Matrix a -> Vector a +flatten = V.concat +{-# INLINE flatten #-} + +{- Function wrapping for the immutable HMatrix interface -} +wrapScalarFunction :: (Vector Double -> Double) -> N.ScalarFunction () +wrapScalarFunction f params _ _ = return $ f params + +wrapScalarFunctionD :: (Vector Double -> (Double, Vector Double)) + -> N.ScalarFunction () +wrapScalarFunctionD f params grad _ = do + case grad of + Nothing -> return () + Just g -> V.copy g usergrad + return result + where + (result, usergrad) = f params + +wrapVectorFunction :: (Vector Double -> Word -> Vector Double) + -> Word -> N.VectorFunction () +wrapVectorFunction f n params vout _ _ = V.copy vout $ f params n + +wrapVectorFunctionD :: (Vector Double -> Word -> (Vector Double, Matrix Double)) + -> Word -> N.VectorFunction () +wrapVectorFunctionD f n params vout jac _ = do + V.copy vout result + case jac of + Nothing -> return () + Just j -> V.copy j (flatten userjac) + where + (result, userjac) = f params n + +wrapPreconditionerFunction :: (Vector Double -> Vector Double -> Vector Double) + -> N.PreconditionerFunction () +wrapPreconditionerFunction f params v vpre _ = V.copy vpre (f params v) + +{- Objective functions -} +-- | An objective function that calculates the objective value at the +-- given parameter vector. +type Objective + = Vector Double -- ^ Parameter vector + -> Double -- ^ Objective function value + +-- | An objective function that calculates both the objective value +-- and the gradient of the objective with respect to the input +-- parameter vector, at the given parameter vector. +type ObjectiveD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Objective function value, gradient) + +-- | A preconditioner function, which computes @vpre = H(x) v@, where +-- @H@ is the Hessian matrix: the positive semi-definite second +-- derivative at the given parameter vector @x@, or an approximation +-- thereof. +type Preconditioner + = Vector Double -- ^ Parameter vector @x@ + -> Vector Double -- ^ Vector @v@ to precondition at @x@ + -> Vector Double -- ^ Preconditioned vector @vpre@ + +data ObjectiveFunction f + = MinimumObjective f + | PreconditionedMinimumObjective Preconditioner f + +applyObjective :: N.Opt -> ObjectiveFunction Objective -> IO N.Result +applyObjective opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunction f) () +applyObjective opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunction f) + (wrapPreconditionerFunction p) () + +applyObjectiveD :: N.Opt -> ObjectiveFunction ObjectiveD -> IO N.Result +applyObjectiveD opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunctionD f) () +applyObjectiveD opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunctionD f) + (wrapPreconditionerFunction p) () + +{- Constraint functions -} +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraint + = Vector Double -- ^ Parameter vector @x@ + -> Double -- ^ Constraint violation (deviation from 0) + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the gradient of @c(x)@ with respect to @x@ at +-- that point. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraintD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Constraint violation, constraint gradient) + +-- | A constraint function which returns a vector @c(x)@ given the +-- parameter vector @x@. The constraint will enforce that @c(x) == 0@ +-- (equality constraint) or @c(x) <= 0@ (inequality constraint). +type VectorConstraint + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> Vector Double -- ^ Constraint violation vector + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the Jacobian (first derivative) matrix of +-- @c(x)@ with respect to @x@ at that point. The constraint will +-- enforce that @c(x) == 0@ (equality constraint) or @c(x) <= 0@ +-- (inequality constraint). +type VectorConstraintD + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> (Vector Double, Matrix Double) -- ^ (Constraint violation vector, + -- constraint Jacobian) + +-- $nonlinearconstraints +-- +-- Note that most NLOPT algorithms do not support nonlinear +-- constraints natively; if you need to enforce nonlinear constraints, +-- you may want to use the 'AugLagAlgorithm' family of solvers, which +-- can add nonlinear constraints to some algorithm that does not +-- support them by a principled modification of the objective +-- function. +-- +-- == Example program +-- +-- The following interactive session example enforces a scalar +-- constraint on the problem given in the beginning of the module: the +-- parameters must always sum to 1. The minimizer finds a constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = COBYLA objf [] [] [constraint] Nothing +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.500000000013028, solutionParams = [0.5000025521533521,0.49999744784664796], solutionResult = FTOL_REACHED}) + + +data Constraint s v + -- | A scalar constraint. + = Scalar s + -- | A vector constraint. + | Vector Word v + -- | A scalar constraint with an attached preconditioning function. + | Preconditioned Preconditioner s + +-- | An equality constraint, comprised of both the constraint function +-- (or functions, if a preconditioner is used) along with the desired +-- tolerance. +data EqualityConstraint s v = EqualityConstraint + { eqConstraintFunctions :: Constraint s v + , eqConstraintTolerance :: Double + } + +-- | An inequality constraint, comprised of both the constraint +-- function (or functions, if a preconditioner is used) along with the +-- desired tolerance. +data InequalityConstraint s v = InequalityConstraint + { ineqConstraintFunctions :: Constraint s v + , ineqConstraintTolerance :: Double + } + +-- | A collection of equality constraints that do not supply +-- constraint derivatives. +type EqualityConstraints = + [EqualityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of inequality constraints that do not supply +-- constraint derivatives. +type InequalityConstraints = + [InequalityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of equality constraints that supply constraint +-- derivatives. +type EqualityConstraintsD = [EqualityConstraint ScalarConstraintD VectorConstraintD] + +-- | A collection of inequality constraints that supply constraint +-- derivatives. +type InequalityConstraintsD = [InequalityConstraint ScalarConstraintD VectorConstraintD] + +class ApplyConstraint constraint where + applyConstraint :: N.Opt -> constraint -> IO N.Result + +instance ApplyConstraint (EqualityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (EqualityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +{- Bounds -} + +-- | Bound constraints are specified by vectors of the same dimension +-- as the parameter space. +-- +-- == Example program +-- +-- The following interactive session example enforces lower bounds on +-- the example from the beginning of the module. This prevents the +-- optimizer from locating the true minimum at @(0, 0)@; a slightly +-- higher constrained minimum at @(1, 1)@ is found. Note that the +-- optimizer returns 'N.XTOL_REACHED' rather than 'N.FTOL_REACHED', +-- because the bound constraint is active at the final minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +-- >>> let lowerbound = LowerBounds $ fromList [1, 1] -- specify bounds +-- >>> let algorithm = NELDERMEAD objf [lowerbound] Nothing -- specify algorithm +-- >>> let problem = LocalProblem 2 stop algorithm -- specify problem +-- >>> let x0 = fromList [5, 10] -- specify initial guess +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 24.0, solutionParams = [1.0,1.0], solutionResult = XTOL_REACHED}) +data Bounds + -- | Lower bound vector @v@ means we want @x >= v@. + = LowerBounds (Vector Double) + -- | Upper bound vector @u@ means we want @x <= u@. + | UpperBounds (Vector Double) + deriving (Eq, Show, Read) + +applyBounds :: N.Opt -> Bounds -> IO N.Result +applyBounds opt (LowerBounds lbvec) = N.set_lower_bounds opt lbvec +applyBounds opt (UpperBounds ubvec) = N.set_upper_bounds opt ubvec + +{- Stopping conditions -} + +-- | A 'StoppingCondition' tells NLOPT when to stop working on a +-- minimization problem. When multiple 'StoppingCondition's are +-- provided, the problem will stop when any one condition is met. +data StoppingCondition + -- | Stop minimizing when an objective value @J@ less than or equal + -- to the provided value is found. + = MinimumValue Double + -- | Stop minimizing when an optimization step changes the objective + -- value @J@ by less than the provided tolerance multiplied by @|J|@. + | ObjectiveRelativeTolerance Double + -- | Stop minimizing when an optimization step changes the objective + -- value by less than the provided tolerance. + | ObjectiveAbsoluteTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than @x@ scaled by the provided + -- tolerance. + | ParameterRelativeTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than the corresponding element in + -- the provided vector of tolerances values. + | ParameterAbsoluteTolerance (Vector Double) + -- | Stop when the number of evaluations of the objective function + -- exceeds the provided count. + | MaximumEvaluations Word + -- | Stop when the optimization time exceeds the provided time (in + -- seconds). This is not a precise limit. + | MaximumTime Double + deriving (Eq, Show, Read) + +-- $nonempty +-- +-- The 'NonEmpty' data type from 'Data.List.NonEmpty' is re-exported +-- here, because it is used to ensure that you always specify at least +-- one stopping condition. + +applyStoppingCondition :: N.Opt -> StoppingCondition -> IO N.Result +applyStoppingCondition opt (MinimumValue x) = N.set_stopval opt x +applyStoppingCondition opt (ObjectiveRelativeTolerance x) = N.set_ftol_rel opt x +applyStoppingCondition opt (ObjectiveAbsoluteTolerance x) = N.set_ftol_abs opt x +applyStoppingCondition opt (ParameterRelativeTolerance x) = N.set_xtol_rel opt x +applyStoppingCondition opt (ParameterAbsoluteTolerance v) = N.set_xtol_abs opt v +applyStoppingCondition opt (MaximumEvaluations n) = N.set_maxeval opt n +applyStoppingCondition opt (MaximumTime deltat) = N.set_maxtime opt deltat + +{- Random seed control -} + +-- | This specifies how to initialize the random number generator for +-- stochastic algorithms. +data RandomSeed + -- | Seed the RNG with the provided value. + = SeedValue Word + -- | Seed the RNG using the system clock. + | SeedFromTime + -- | Don't perform any explicit initialization of the RNG. + | Don'tSeed + deriving (Eq, Show, Read) + +applyRandomSeed :: RandomSeed -> IO () +applyRandomSeed Don'tSeed = return () +applyRandomSeed (SeedValue n) = N.srand n +applyRandomSeed SeedFromTime = N.srand_time + +{- Random stuff -} + +-- | This specifies the population size for algorithms that use a pool +-- of solutions. +newtype Population = Population Word deriving (Eq, Show, Read) + +applyPopulation :: N.Opt -> Population -> IO N.Result +applyPopulation opt (Population n) = N.set_population opt n + +-- | This specifies the memory size to be used by algorithms like +-- 'LBFGS' which store approximate Hessian or Jacobian matrices. +newtype VectorStorage = VectorStorage Word deriving (Eq, Show, Read) + +applyVectorStorage :: N.Opt -> VectorStorage -> IO N.Result +applyVectorStorage opt (VectorStorage n) = N.set_vector_storage opt n + +-- | This vector with the same dimension as the parameter vector @x@ +-- specifies the initial step for the optimizer to take. (This +-- applies to local gradient-free algorithms, which cannot use +-- gradients to estimate how big a step to take.) +newtype InitialStep = InitialStep (Vector Double) deriving (Eq, Show, Read) + +applyInitialStep :: N.Opt -> InitialStep -> IO N.Result +applyInitialStep opt (InitialStep v) = N.set_initial_step opt v + +{- Algorithms -} + +data GlobalProblem = GlobalProblem + { lowerBounds :: Vector Double -- ^ Lower bounds for @x@ + , upperBounds :: Vector Double -- ^ Upper bounds for @x@ + , gstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , galgorithm :: GlobalAlgorithm -- ^ Algorithm specification + } + +-- | These are the global minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'Population', you can simply specify 'Nothing' to use +-- the default behavior. +data GlobalAlgorithm + -- | DIviding RECTangles + = DIRECT Objective + -- | DIviding RECTangles, locally-biased variant + | DIRECT_L Objective + -- | DIviding RECTangles, "slightly randomized" + | DIRECT_L_RAND Objective RandomSeed + -- | DIviding RECTangles, unscaled version + | DIRECT_NOSCAL Objective + -- | DIviding RECTangles, locally-biased and unscaled + | DIRECT_L_NOSCAL Objective + -- | DIviding RECTangles, locally-biased, unscaled and "slightly + -- randomized" + | DIRECT_L_RAND_NOSCAL Objective RandomSeed + -- | DIviding RECTangles, original FORTRAN implementation + | ORIG_DIRECT Objective InequalityConstraints + -- | DIviding RECTangles, locally-biased, original FORTRAN + -- implementation + | ORIG_DIRECT_L Objective InequalityConstraints + -- | Stochastic Global Optimization. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO ObjectiveD + -- | Stochastic Global Optimization, randomized variant. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO_RAND ObjectiveD RandomSeed + -- | Controlled Random Search with Local Mutation + | CRS2_LM Objective RandomSeed (Maybe Population) + -- | Improved Stochastic Ranking Evolution Strategy + | ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population) + -- | Evolutionary Algorithm + | ESCH Objective + -- | Original Multi-Level Single-Linkage + | MLSL Objective LocalProblem (Maybe Population) + -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy + -- Sequence for starting points + | MLSL_LDS Objective LocalProblem (Maybe Population) + +algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm +algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT +algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L +algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND +algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL +algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT +algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L +algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO +algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND +algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM +algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES +algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH +algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL +algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS + +applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + + go (DIRECT o) = obj o + go (DIRECT_L o) = obj o + go (DIRECT_NOSCAL o) = obj o + go (DIRECT_L_NOSCAL o) = obj o + go (ESCH o) = obj o + go (STOGO o) = objD o + go (DIRECT_L_RAND o _) = obj o + go (DIRECT_L_RAND_NOSCAL o _) = obj o + go (ORIG_DIRECT o _) = obj o + go (ORIG_DIRECT_L o _) = obj o + go (STOGO_RAND o _) = objD o + go (CRS2_LM o _ _) = obj o + go (ISRES o _ _ _ _) = obj o + go (MLSL o _ _) = obj o + go (MLSL_LDS o _ _) = obj o + +applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalAlgorithm opt alg = do + applyGlobalObjective opt alg + go alg + where + seed = applyRandomSeed + pop = maybe (return ()) (tryTo . applyPopulation opt) + ic = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + + local lp = setupLocalProblem lp >>= N.set_local_optimizer opt + + go (DIRECT_L_RAND _ s) = seed s + go (DIRECT_L_RAND_NOSCAL _ s) = seed s + go (ORIG_DIRECT _ ineq) = ic ineq + go (ORIG_DIRECT_L _ ineq) = ic ineq + go (STOGO_RAND _ s) = seed s + go (CRS2_LM _ s p) = seed s *> pop p + go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p + go (MLSL _ lp p) = local lp *> pop p + go (MLSL_LDS _ lp p) = local lp *> pop p + go _ = return () + +tryTo :: IO N.Result -> IO () +tryTo act = do + result <- act + if (N.isSuccess result) + then return () + else Ex.throw $ NloptException result + +data NloptException = NloptException N.Result deriving (Show, Typeable) +instance Exception NloptException + +-- | Solve the specified global optimization problem. +-- +-- = Example program +-- +-- The following interactive session example uses the 'ISRES' +-- algorithm, a stochastic, derivative-free global optimizer, to +-- minimize a trivial function with a minimum of 22.0 at @(0, 0)@. +-- The search is conducted within a box from -10 to 10 in each +-- dimension. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion +-- >>> let algorithm = ISRES objf [] [] (SeedValue 22) Nothing -- specify algorithm +-- >>> let lowerbounds = fromList [-10, -10] -- specify bounds +-- >>> let upperbounds = fromList [10, 10] -- specify bounds +-- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm +-- >>> let x0 = fromList [5, 8] -- specify initial guess +-- >>> minimizeGlobal problem x0 +-- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED}) +minimizeGlobal :: GlobalProblem -- ^ Problem specification + -> Vector Double -- ^ Initial parameter guess + -> Either N.Result Solution -- ^ Optimization results +minimizeGlobal prob x0 = + unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +applyGlobalProblem :: N.Opt -> GlobalProblem -> IO () +applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do + tryTo $ applyBounds opt (LowerBounds lb) + tryTo $ applyBounds opt (UpperBounds ub) + traverse_ (tryTo . applyStoppingCondition opt) stop + applyGlobalAlgorithm opt alg + +newOpt :: N.Algorithm -> Word -> IO N.Opt +newOpt alg sz = do + opt' <- N.create alg sz + case opt' of + Nothing -> Ex.throw $ NloptException N.FAILURE + Just opt -> return opt + +setupGlobalProblem :: GlobalProblem -> IO N.Opt +setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do + opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp) + applyGlobalProblem opt gp + return opt + +solveProblem :: N.Opt -> Vector Double -> IO Solution +solveProblem opt x0 = do + (N.Output outret outcost outx nevals) <- N.optimize opt x0 + if (N.isSuccess outret) + then return $ Solution outcost outx outret nevals + else Ex.throw $ NloptException outret + +minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution +minimizeGlobal' gp x0 = do + opt <- setupGlobalProblem gp + solveProblem opt x0 + +data LocalProblem = LocalProblem + { lsize :: Word -- ^ The dimension of the + -- parameter vector. + , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification + } + +-- | These are the local minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one +-- another. Note that some local methods require you provide +-- derivatives (gradients or Jacobians) for your objective function +-- and constraint functions. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to +-- use the default behavior. +data LocalAlgorithm + -- | Limited-memory BFGS + = LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage) + -- | Limited-memory BFGS + | LBFGS ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-2 + | VAR2 ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-1 + | VAR1 ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method + | TNEWTON ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method with automatic restarting + | TNEWTON_RESTART ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method + | TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method with automatic + -- restarting + | TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage) + -- | Method of moving averages + | MMA ObjectiveD InequalityConstraintsD + -- | Sequential Least-Squares Quadratic Programming + | SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD + -- | Conservative Convex Separable Approximation + | CCSAQ ObjectiveD Preconditioner + -- | PRincipal AXIS gradient-free local optimization + | PRAXIS Objective [Bounds] (Maybe InitialStep) + -- | Constrained Optimization BY Linear Approximations + | COBYLA Objective [Bounds] InequalityConstraints EqualityConstraints + (Maybe InitialStep) + -- | Powell's NEWUOA algorithm + | NEWUOA Objective (Maybe InitialStep) + -- | Powell's NEWUOA algorithm with bounds by SGJ + | NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep) + -- | Nelder-Mead Simplex gradient-free method + | NELDERMEAD Objective [Bounds] (Maybe InitialStep) + -- | NLOPT implementation of Rowan's Subplex algorithm + | SBPLX Objective [Bounds] (Maybe InitialStep) + -- | Bounded Optimization BY Quadratic Approximations + | BOBYQA Objective [Bounds] (Maybe InitialStep) + +algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm +algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL +algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS +algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2 +algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1 +algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON +algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART +algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND +algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART +algorithmEnumOfLocal (MMA _ _) = N.LD_MMA +algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP +algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ +algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS +algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA +algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD +algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX +algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA + +applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO () +applyLocalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p + + go (LBFGS_NOCEDAL o _) = objD o + go (LBFGS o _) = objD o + go (VAR2 o _) = objD o + go (VAR1 o _) = objD o + go (TNEWTON o _) = objD o + go (TNEWTON_RESTART o _) = objD o + go (TNEWTON_PRECOND o _) = objD o + go (TNEWTON_PRECOND_RESTART o _) = objD o + go (MMA o _) = objD o + go (SLSQP o _ _ _) = objD o + go (CCSAQ o prec) = precond prec o + go (PRAXIS o _ _) = obj o + go (COBYLA o _ _ _ _) = obj o + go (NEWUOA o _) = obj o + go (NEWUOA_BOUND o _ _) = obj o + go (NELDERMEAD o _ _) = obj o + go (SBPLX o _ _) = obj o + go (BOBYQA o _ _) = obj o + +applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO () +applyLocalAlgorithm opt alg = do + applyLocalObjective opt alg + go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + ecd = traverse_ (tryTo . applyConstraint opt) + store = maybe (return ()) (tryTo . applyVectorStorage opt) + bound = traverse_ (tryTo . applyBounds opt) + step0 = maybe (return ()) (tryTo . applyInitialStep opt) + + go (LBFGS_NOCEDAL _ vs) = store vs + go (LBFGS _ vs) = store vs + go (VAR2 _ vs) = store vs + go (VAR1 _ vs) = store vs + go (TNEWTON _ vs) = store vs + go (TNEWTON_RESTART _ vs) = store vs + go (TNEWTON_PRECOND _ vs) = store vs + go (TNEWTON_PRECOND_RESTART _ vs) = store vs + go (MMA _ ineqd) = icd ineqd + go (SLSQP _ b ineqd eqd) = + bound b *> icd ineqd *> ecd eqd + go (CCSAQ _ _ ) = return () + go (PRAXIS _ b s) = bound b *> step0 s + go (COBYLA _ b ineq eq s) = + bound b *> ic ineq *> ec eq *> step0 s + go (NEWUOA _ s) = step0 s + go (NEWUOA_BOUND _ b s) = bound b *> step0 s + go (NELDERMEAD _ b s) = bound b *> step0 s + go (SBPLX _ b s) = bound b *> step0 s + go (BOBYQA _ b s) = bound b *> step0 s + +applyLocalProblem :: N.Opt -> LocalProblem -> IO () +applyLocalProblem opt (LocalProblem _ stop alg) = do + traverse_ (tryTo . applyStoppingCondition opt) stop + applyLocalAlgorithm opt alg + +setupLocalProblem :: LocalProblem -> IO N.Opt +setupLocalProblem lp@(LocalProblem sz _ alg) = do + opt <- newOpt (algorithmEnumOfLocal alg) sz + applyLocalProblem opt lp + return opt + +minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution +minimizeLocal' lp x0 = do + opt <- setupLocalProblem lp + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the SLSQP solver to find the minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale ) +-- >>> let objf x = (x `dot` x + 22, 2 `scale` x) +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1]) +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = SLSQP objf [] [] [constraint] +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED}) +minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution +minimizeLocal prob x0 = + unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +class ProblemSize c where + problemSize :: c -> Word + +instance ProblemSize LocalProblem where + problemSize = lsize + +instance ProblemSize GlobalProblem where + problemSize = fromIntegral . V.length . lowerBounds + +instance ProblemSize AugLagProblem where + problemSize (AugLagProblem _ _ alg) = case alg of + AUGLAG_LOCAL lp _ _ -> problemSize lp + AUGLAG_EQ_LOCAL lp -> problemSize lp + AUGLAG_GLOBAL gp _ _ -> problemSize gp + AUGLAG_EQ_GLOBAL gp -> problemSize gp + + +-- | __IMPORTANT NOTE__ +-- +-- For augmented lagrangian problems, you, the user, are responsible +-- for providing the appropriate type of constraint. If the +-- subsidiary problem requires an `ObjectiveD`, then you should +-- provide constraint functions with derivatives. If the subsidiary +-- problem requires an `Objective`, you should provide constraint +-- functions without derivatives. If you don't do this, you may get a +-- runtime error. +data AugLagProblem = AugLagProblem + { alEquality :: EqualityConstraints -- ^ Possibly empty set of + -- equality constraints + , alEqualityD :: EqualityConstraintsD -- ^ Possibly empty set of + -- equality constraints with + -- derivatives + , alalgorithm :: AugLagAlgorithm -- ^ Algorithm specification. + } + +-- | The Augmented Lagrangian solvers allow you to enforce nonlinear +-- constraints while using local or global algorithms that don't +-- natively support them. The subsidiary problem is used to do the +-- minimization, but the @AUGLAG@ methods modify the objective to +-- enforce the constraints. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- See the documentation for 'AugLagProblem' for an important note +-- about the constraint functions. +data AugLagAlgorithm + -- | AUGmented LAGrangian with a local subsidiary method + = AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a local subsidiary method and with + -- penalty functions only for equality constraints + | AUGLAG_EQ_LOCAL LocalProblem + -- | AUGmented LAGrangian with a global subsidiary method + | AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a global subsidiary method and with + -- penalty functions only for equality constraints. + | AUGLAG_EQ_GLOBAL GlobalProblem + +algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm +algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ +algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ + +-- | This structure is returned in the event of a successful +-- optimization. +data Solution = Solution + { solutionCost :: Double -- ^ The objective function value + -- at the minimum + , 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 () +applyAugLagAlgorithm opt alg = go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + -- AUGLAG won't work at all if you don't pass it the same + -- objective as the subproblem -- here we pull out the subproblem + -- objectives from the algorithm spec and set the same objective + -- function so the user can't mess it up. + local lp = tryTo $ do + localopt <- setupLocalProblem lp + applyLocalObjective opt (lalgorithm lp) + N.set_local_optimizer opt localopt + global gp = do + tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt + applyGlobalObjective opt (galgorithm gp) + + go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_LOCAL lp) = local lp + go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_GLOBAL gp) = global gp + +applyAugLagProblem :: N.Opt -> AugLagProblem -> IO () +applyAugLagProblem opt (AugLagProblem eq eqd alg) = do + traverse_ (tryTo . applyConstraint opt) eq + traverse_ (tryTo . applyConstraint opt) eqd + applyAugLagAlgorithm opt alg + +minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution +minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do + opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap) + applyAugLagProblem opt ap + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the augmented Lagrangian method to enforce the constraint and +-- the 'SBPLX' algorithm, which does not support nonlinear constraints +-- itself, to perform the minimization. As before, the parameters +-- must always sum to 1, and the minimizer finds the same constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let algorithm = SBPLX objf [] Nothing +-- >>> let subproblem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal subproblem x0 +-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem) +-- >>> minimizeAugLag problem x0 +-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED}) + +minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution +minimizeAugLag prob x0 = + unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode diff --git a/src/Text/ParseSR.hs b/src/Text/ParseSR.hs --- a/src/Text/ParseSR.hs +++ b/src/Text/ParseSR.hs @@ -11,7 +11,7 @@ -- Functions to parse a string representing an expression -- ----------------------------------------------------------------------------- -module Text.ParseSR ( parseSR, parsePat, parseNonTerms, showOutput, SRAlgs(..), Output(..) ) +module Text.ParseSR ( parseSR, parseNonTerms, showOutput, SRAlgs(..), Output(..) ) -- parsePat, where import Control.Applicative ((<|>)) @@ -21,7 +21,7 @@ import Data.Char (toLower) import Data.List (sortOn) import Data.SRTree -import Algorithm.EqSat.DB +--import Algorithm.EqSat.DB import qualified Data.SRTree.Print as P import qualified Data.Map.Strict as Map import Data.List.Split ( splitOn ) @@ -34,7 +34,7 @@ -- numerical values represented as `Double`. The numerical values type -- can be changed with `fmap`. type ParseTree = Parser (Fix SRTree) -type ParsePat = Parser Pattern +--type ParsePat = Parser Pattern -- * Data types and caller functions @@ -65,8 +65,8 @@ 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 +--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 @@ -325,7 +325,7 @@ ix <- decimal pure $ Fix $ Var ix "var" - +{- -- parse a pattern expression parsePatExpr :: ParsePat parsePatExpr = parsePattern (prefixOps : binOps) binFuns var @@ -387,7 +387,7 @@ getParserVar k v = (string k <|> enveloped k) >> pure (Fix $ Var v) enveloped s = (char ' ' <|> char '(') >> string s >> (char ' ' <|> char ')') >> pure "" - + -} -- * Parse the non-terminal nodes into a SRTree () value parseNonTerms :: String -> [SRTree ()] parseNonTerms = Prelude.map toNonTerm . splitOn "," diff --git a/src/Text/ParseSR/IO.hs b/src/Text/ParseSR/IO.hs --- a/src/Text/ParseSR/IO.hs +++ b/src/Text/ParseSR/IO.hs @@ -15,7 +15,7 @@ where -- import Data.SRTree.EqSat1 -import Algorithm.EqSat.Simplify ( simplifyEqSatDefault ) +--import Algorithm.EqSat.Simplify ( simplifyEqSatDefault ) import Control.Monad (forM_, unless) import qualified Data.ByteString.Char8 as B import Data.SRTree @@ -35,7 +35,7 @@ contents <- hGetLines h let myParserFun = parseSR sr (B.pack hd) param . B.pack -- myParser = if simpl then fmap simplifyEqSat . myParserFun else myParserFun - myParser = if simpl then fmap simplifyEqSatDefault . myParserFun else myParserFun + myParser = myParserFun -- if simpl then fmap simplifyEqSatDefault . myParserFun else myParserFun es = map myParser $ filter (not . null) contents unless (null fname) $ hClose h pure es diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -1,246 +1,236 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.38.1. +-- This file has been generated from package.yaml by hpack version 0.39.6. -- -- see: https://github.com/sol/hpack -name: srtree -version: 2.0.1.8 -synopsis: A general library to work with Symbolic Regression expression trees. -description: A Symbolic Regression Tree data structure to work with mathematical expressions with support to first order derivative and simplification; -category: Math, Data, Data Structures -homepage: https://github.com/folivetti/srtree#readme -bug-reports: https://github.com/folivetti/srtree/issues -author: Fabricio Olivetti de França -maintainer: fabricio.olivetti@gmail.com -copyright: 2023 Fabricio Olivetti de França -license: BSD3 -license-file: LICENSE -build-type: Simple +name: srtree +version: 3.0.0.0 +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; +license: BSD3 +license-file: LICENSE +author: Fabricio Olivetti de França +maintainer: fabricio.olivetti@gmail.com +copyright: 2023 Fabricio Olivetti de França +category: Math, Data, Data Structures +homepage: https://github.com/folivetti/srtree#readme +bug-reports: https://github.com/folivetti/srtree/issues +build-type: Simple extra-source-files: - README.md - ChangeLog.md + README.md + ChangeLog.md source-repository head - type: git - location: https://github.com/folivetti/srtree + type: git + location: https://github.com/folivetti/srtree library - exposed-modules: - Algorithm.EqSat - Algorithm.EqSat.Build - Algorithm.EqSat.DB - Algorithm.EqSat.Egraph - Algorithm.EqSat.Info - Algorithm.EqSat.Queries - Algorithm.EqSat.SearchSR - Algorithm.EqSat.SearchSRCache - Algorithm.EqSat.Simplify - Algorithm.Massiv.Utils - Algorithm.SRTree.AD - Algorithm.SRTree.ConfidenceIntervals - Algorithm.SRTree.Likelihoods - Algorithm.SRTree.ModelSelection - Algorithm.SRTree.NonlinearOpt - Algorithm.SRTree.Opt - Data.SRTree - Data.SRTree.Datasets - Data.SRTree.Derivative - Data.SRTree.Eval - Data.SRTree.Internal - Data.SRTree.Print - Data.SRTree.Random - Data.SRTree.Recursion - Numeric.Optimization.NLOPT.Bindings - Text.ParseSR - Text.ParseSR.IO - other-modules: - Paths_srtree - hs-source-dirs: - src - ghc-options: -fwarn-incomplete-patterns -threaded - extra-libraries: - nlopt - 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.9 - , 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.1 && <1.1 - , mtl >=2.2 && <2.4 - , 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 - default-language: Haskell2010 + exposed-modules: + Algorithm.EqSat + Algorithm.EqSat.Build + Algorithm.EqSat.DB + Algorithm.EqSat.Egraph + Algorithm.EqSat.Info + Algorithm.EqSat.Queries + Algorithm.EqSat.SearchSR + Algorithm.EqSat.Simplify + Algorithm.EqSat.Store + Algorithm.SRTree.AD + Algorithm.SRTree.AD.CompiledAD + Algorithm.SRTree.AD.Unboxed + Algorithm.SRTree.Compile + Algorithm.SRTree.ConfidenceIntervals + Algorithm.SRTree.Likelihoods + Algorithm.SRTree.ModelSelection + Algorithm.SRTree.NonlinearOpt + Algorithm.SRTree.Utils + Data.SRTree + Data.SRTree.Datasets + Data.SRTree.Derivative + Data.SRTree.Eval + Data.SRTree.Internal + Data.SRTree.Print + Data.SRTree.Random + Data.SRTree.Recursion + Numeric.Optimization.NLOPT + Numeric.Optimization.NLOPT.Bindings + Text.ParseSR + Text.ParseSR.IO + other-modules: + Paths_srtree + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , statistics >=0.15 && <0.17 + , time >=1.9 && <1.15 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + hs-source-dirs: + src + ghc-options: -O2 -fwarn-incomplete-patterns -fspec-constr + extra-libraries: + nlopt + default-language: Haskell2010 -executable srsimplify - main-is: Main.hs - other-modules: - Paths_srtree - hs-source-dirs: - apps/srsimplify - 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.19 && <5 - , binary >=0.8.9.1 && <0.9 - , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.9 - , 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.1 && <1.1 - , mtl >=2.2 && <2.4 - , optparse-applicative >=0.18 && <0.20 - , 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 bench + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/Bench + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -fllvm -pgmlo opt-20 -pgmlc llc-20 -optlo-O3 -optlc-mcpu=native -mavx2 -mfma -fspec-constr -fmax-simplifier-iterations=20 -fexpose-all-unfoldings + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , criterion >=1.5 && <2 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 -executable srtools - main-is: Main.hs - other-modules: - Args - IO - Report - Paths_srtree - hs-source-dirs: - apps/srtools - 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.19 && <5 - , binary >=0.8.9.1 && <0.9 - , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.9 - , 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.1 && <1.1 - , mtl >=2.2 && <2.4 - , optparse-applicative >=0.18 && <0.20 - , 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 bench-eqsat + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/BenchEqSat + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , criterion >=1.5 && <2 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 -executable tinygp - main-is: Main.hs - other-modules: - GP - Initialization - Util - Paths_srtree - hs-source-dirs: - apps/tinygp - 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.19 && <5 - , binary >=0.8.9.1 && <0.9 - , bytestring >=0.11 && <0.13 - , containers >=0.6.7 && <0.9 - , 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.1 && <1.1 - , mtl >=2.2 && <2.4 - , optparse-applicative >=0.18 && <0.20 - , 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 srtree-report + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/Report + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , optparse-applicative >=0.16 && <0.19 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 test-suite srtree-test - type: exitcode-stdio-1.0 - main-is: Spec.hs - other-modules: - Paths_srtree - hs-source-dirs: - test - ghc-options: -threaded -rtsopts -with-rtsopts=-N - build-depends: - HUnit - , ad - , 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.9 - , 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.1 && <1.1 - , mtl >=2.2 && <2.4 - , 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 + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + EqSatTests + StoreTests + Paths_srtree + hs-source-dirs: + test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + HUnit >=1.6 && <1.7 + , ad >=5.0 && <6 + , async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 diff --git a/test/EqSatTests.hs b/test/EqSatTests.hs new file mode 100644 --- /dev/null +++ b/test/EqSatTests.hs @@ -0,0 +1,630 @@ +{-# LANGUAGE OverloadedStrings #-} + +module EqSatTests where + +import Test.HUnit +import Data.SRTree +import Data.SRTree.Print (showExpr) +import qualified Data.IntSet as IntSet +import qualified Data.IntMap as IntMap +import qualified Data.Map as Map +import qualified Data.HashSet as Set +import qualified Data.Vector.Unboxed as VU +import qualified Data.Set as RangeSet +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Simplify (simplifyEqSatDefault, rewrites, rewritesParams) +import Control.Monad.State.Strict +import Control.Monad (forM_) +import Control.Monad.Identity +import Data.List (nub, sort) + +eps :: Double +eps = 1e-9 + +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 + +runEG :: EGraphST Identity a -> (a, EGraph) +runEG m = runIdentity $ runStateT m emptyGraph + +evalEG :: EGraphST Identity a -> a +evalEG m = runIdentity $ evalStateT m emptyGraph + +-- | Test 1: fromTree with a leaf (variable) +test_fromTree_var :: Test +test_fromTree_var = TestCase $ do + let tree = var 0 + (eid, eg) = runEG $ fromTree myCost tree + assertBool "fromTree var: eid should be >= 0" (eid >= 0) + assertBool "fromTree var: eclass exists" (IntMap.member eid (_eClass eg)) + let ec = _eClass eg IntMap.! eid + assertBool "fromTree var: eclass has nodes" (not $ null (_eNodes ec)) + let bestNode = head $ Set.toList (_eNodes ec) + assertEqual "fromTree var: best is Var 0" (EVar 0) bestNode + +-- | Test 2: fromTree with a binary expression +test_fromTree_bin :: Test +test_fromTree_bin = TestCase $ do + let tree = var 0 + constv 1.0 + (eid, eg) = runEG $ fromTree myCost tree + assertBool "fromTree bin: eid >= 0" (eid >= 0) + let ec = _eClass eg IntMap.! eid + assertBool "fromTree bin: eclass has nodes" (not $ null (_eNodes ec)) + +-- | Test 3: Canonical identity (an e-class should be its own canonical) +test_canonical_identity :: Test +test_canonical_identity = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0) + (canId, _) = runIdentity $ runStateT (canonical eid) eg + assertEqual "canonical of fresh id is itself" eid canId + +-- | Test 4: canonize canonizes children +test_canonize :: Test +test_canonize = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0 + constv 1.0) + (canNode, _) = runIdentity $ runStateT (do + ec <- getEClass eid + let someNode = head $ Set.toList (_eNodes ec) + canonize someNode) eg + -- All children should be canonical now + let children = eChildren canNode + forM_ children $ \c -> do + let (canC, _) = runIdentity $ runStateT (canonical c) eg + assertEqual "canonize: child is canonical" c canC + +-- | Test 5: Adding duplicate e-node returns existing e-class +test_add_duplicate :: Test +test_add_duplicate = TestCase $ do + let tree = constv 2.0 + (eid1, eg1) = runEG $ fromTree myCost tree + (eid2, eg2) = runEG' eg1 $ add myCost (EConst 2.0) + assertEqual "add duplicate returns same eclass" eid1 eid2 + where + runEG' eg m = runIdentity $ runStateT m eg + +-- | Test 6: Merge two distinct e-classes +test_merge :: Test +test_merge = TestCase $ do + let (eid1, eg1) = runEG $ fromTree myCost (var 0) + (eid2, eg2) = runIdentity $ runStateT (fromTree myCost (var 1)) eg1 + assertBool "merge: eid1 and eid2 start different" (eid1 /= eid2) + let (mergedId, eg3) = runIdentity $ runStateT (merge myCost eid1 eid2) eg2 + can1 = _canonicalMap eg3 IntMap.! eid1 + can2 = _canonicalMap eg3 IntMap.! eid2 + assertEqual "merge: canonicals are equal" can1 can2 + assertEqual "merge: leader matches canonical" mergedId can1 + +-- | Test 7: Rebuild after add +test_rebuild :: Test +test_rebuild = TestCase $ do + let tree = var 0 + constv 1.0 + eg = snd $ runEG $ do + _ <- fromTree myCost tree + rebuild myCost + assertBool "rebuild: eNodeToEClass non-empty" (not $ null (_eNodeToEClass eg)) + assertBool "rebuild: worklist empty" (null (_worklist (_eDB eg))) + assertBool "rebuild: analysis empty" (null (_analysis (_eDB eg))) + +-- | Test 8: Basic pattern matching +test_match :: Test +test_match = TestCase $ do + let tree = var 0 + constv 1.0 + pat = Fixed (Bin Add (VarPat 'x') (VarPat 'y')) + (substs, _) = runEG $ do + _ <- fromTree myCost tree + match pat + assertBool "match: should have at least one substitution" (not $ null substs) + +-- | Test 9: Extraction (getBestExpr) +test_getBestExpr :: Test +test_getBestExpr = TestCase $ do + let tree = var 0 + constv 1.0 + (extracted, _) = runEG $ do + eid <- fromTree myCost tree + getBestExpr eid + assertEqual "getBestExpr preserves structure" (showExpr tree) (showExpr extracted) + +-- | Test 10: Equality saturation with x + 0 = x +test_eqsat_x_plus_0 :: Test +test_eqsat_x_plus_0 = TestCase $ do + let tree = var 0 + constv 0.0 + rule = "a" + 0 :=> "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x+0 = x" (showExpr (var 0)) (showExpr best) + +-- | Test 11: Equality saturation with x * 1 = x +test_eqsat_x_times_1 :: Test +test_eqsat_x_times_1 = TestCase $ do + let tree = var 0 * constv 1.0 + rule = "a" * 1 :=> "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x*1 = x" (showExpr (var 0)) (showExpr best) + +-- | Test 12: Fitness and theta storage round-trip +test_fitness_theta :: Test +test_fitness_theta = TestCase $ do + let theta = [VU.fromList [1.0, 2.0]] + (mf, _) = runEG $ do + eid <- fromTree myCost (var 0) + insertFitness eid 0.5 theta + getFitness eid + case mf of + Nothing -> assertFailure "getFitness returned Nothing" + Just f -> assertBool "fitness should be ~0.5" (abs (f - 0.5) < eps) + +-- | Test 13: Insert fitness and check range tree +test_fitness_range :: Test +test_fitness_range = TestCase $ do + let (eg, _) = runEG $ do + eid1 <- fromTree myCost (var 0) + eid2 <- fromTree myCost (constv 1.0) + insertFitness eid1 (-1.0) [] + insertFitness eid2 2.0 [] + gets _eDB + rt = _fitRangeDB eg + case getGreatest rt of + Just (bestFit, _) -> assertBool "fitness range: best is 2.0" (abs (bestFit - 2.0) < eps) + Nothing -> assertFailure "fitness range: non-empty" + +-- | Test 14: getTopFitEClassWithSize +test_top_fit_size :: Test +test_top_fit_size = TestCase $ do + let (eclasses, _) = runEG $ do + eid1 <- fromTree myCost (var 0) -- size 1 + eid2 <- fromTree myCost (constv 1.0) -- size 1 + eid3 <- fromTree myCost (var 0 + constv 1.0) -- size 3 + insertFitness eid1 0.5 [] + insertFitness eid2 1.0 [] + insertFitness eid3 2.0 [] + getTopFitEClassWithSize 1 1 + assertBool "top fit size 1: should have at least one" (not $ null eclasses) + assertEqual "top fit size 1: should be 1 result" 1 (length eclasses) + +-- | Test 15: Bidirectional rule (x + 0 == x) +test_eqsat_comm :: Test +test_eqsat_comm = TestCase $ do + let tree = var 0 + constv 0.0 + rule = "a" + 0 :==: "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x+0 == x" (showExpr (var 0)) (showExpr best) + +-- | Test 16: Double negation elimination +test_eqsat_double_neg :: Test +test_eqsat_double_neg = TestCase $ do + -- var 0 - (var 0 - const 2) should simplify via x - (x - y) = y + -- but we don't have that rule. Instead use const folding: + -- (1 + 0) * x = x via x * 1 = x after const folding simplifies 1+0 to 1 + -- Actually let's use a simpler rule set + let tree = (constv 1.0 + constv 0.0) * var 0 -- (1+0)*x + rules = ["a" + 0 :=> "a", "a" * 1 :=> "a"] + (best, _) = runEG $ eqSat tree rules myCost 10 + assertEqual "eqSat: (1+0)*x = x" (showExpr (var 0)) (showExpr best) + +-- | Test 17: fromTrees builds multiple independent trees +test_fromTrees :: Test +test_fromTrees = TestCase $ do + let trees = [var 0, constv 1.0, var 0 + constv 1.0] + (eids, eg) = runEG $ fromTrees myCost trees + assertEqual "fromTrees: three trees" 3 (length eids) + -- each eid should be distinct and valid + let allDistinct = length eids == length (map (\x -> _canonicalMap eg IntMap.! x) eids) + assertBool "fromTrees: distinct eclasses" allDistinct + assertBool "fromTrees: each eid in eClass" (all (`IntMap.member` _eClass eg) eids) + +-- | Test 18: Cost function respects node types +test_cost :: Test +test_cost = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0) + cost = _cost . _info $ (_eClass eg IntMap.! eid) + assertEqual "cost of Var is 1" 1 cost + +-- | Test 19: getAllExpressionsFrom +test_get_all_expr :: Test +test_get_all_expr = TestCase $ do + let (exprs, _) = runEG $ do + eid <- fromTree myCost (var 0 + constv 1.0) + getAllExpressionsFrom eid + assertBool "getAllExpressionsFrom: non-empty" (not $ null exprs) + assertEqual "getAllExpressionsFrom: includes original" (showExpr (var 0 + constv 1.0)) (showExpr (head exprs)) + +-- | Test 20: sizeFitDB has no stale entries after refit with lower fitness +test_sizeFitDB_no_stale :: Test +test_sizeFitDB_no_stale = TestCase $ do + let (eg, _) = runEG $ do + eid <- fromTree myCost (var 0) -- size = 1 + insertFitness eid 1.0 [] -- insert higher fitness + insertFitness eid 0.5 [] -- refit with lower fitness + gets _eDB + sfd = _sizeFitDB eg + -- size 1 should have exactly 1 entry (the new fitness 0.5) + size1Entries = case IntMap.lookup 1 sfd of + Nothing -> 0 + Just rt -> length (RangeSet.toList rt) + assertEqual "sizeFitDB: size 1 should have 1 entry after refit" 1 size1Entries + -- verify the entry is the new fitness, not the old one + case IntMap.lookup 1 sfd >>= RangeSet.lookupMax of + Nothing -> assertFailure "sizeFitDB: size 1 should have an entry" + Just (f, eId) -> assertBool "sizeFitDB: fitness should be 0.5" (abs (f - 0.5) < eps) + +-- | Test 21: trie paths are canonical after merge+rebuild +-- repair never calls addToDB, so stale non-canonical keys remain in the trie. +-- This test verifies that no stale (non-canonical) keys exist after a merge. +test_trie_no_stale_keys :: Test +test_trie_no_stale_keys = TestCase $ do + let (eg, _) = runEG $ do + eid_a <- fromTree myCost (var 0) -- eclass 0 + eid_0 <- fromTree myCost (constv 0.0) -- eclass 1 + eid_t <- fromTree myCost (addZero (var 0) (constv 0.0)) -- eclass 2 (a+0) + + -- Merge a+0 (2) with a (0), so 2 → canonical 0 + mergedId <- merge myCost eid_t eid_a + rebuild myCost + + -- Add a parent (a+0)*b after the merge + eid_b <- fromTree myCost (var 1) -- eclass 3 + eid_parent <- fromTree myCost (addZero (var 0) (constv 0.0) * var 1) -- (a+0)*b + rebuild myCost + + gets id + can = _canonicalMap eg + staleKeys = getAllStaleTrieKeys can (_patDB $ _eDB eg) + assertBool ("trie: expected exactly 1 stale key (2), got: " <> show staleKeys) (staleKeys == [2]) + +-- | Helper: construct a+0 bypassing Num instance optimization that rewrites +0 to identity +addZero :: Fix SRTree -> Fix SRTree -> Fix SRTree +addZero l r = Fix (Bin Add l r) + +-- | Helper: construct a binary tree bypassing Num instance simplifications +mkBin :: Op -> Fix SRTree -> Fix SRTree -> Fix SRTree +mkBin op l r = Fix (Bin op l r) + +-- | Test 22: multi-atom match works after merge (requires toCanon in intersectAtoms) +test_match_after_merge_multi_atom :: Test +test_match_after_merge_multi_atom = TestCase $ do + let pat = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'a') (Fixed (Const 0.0)))) (VarPat 'b')) + ((substs, _, _, _, _), _) = runEG $ do + eid_a <- fromTree myCost (var 0) + eid_0 <- fromTree myCost (constv 0.0) + eid_t <- fromTree myCost (addZero (var 0) (constv 0.0)) + mergedId <- merge myCost eid_t eid_a + rebuild myCost + eid_b <- fromTree myCost (var 1) + eid_parent <- fromTree myCost (addZero (var 0) (constv 0.0) * var 1) + rebuild myCost + substs <- match pat + pure (substs, (), (), (), ()) + assertBool "match: multi-atom should work after merge" (not $ null substs) + +-- | Test 23: flattened ENAry multiset for a right-nested Add +test_enary_flatten :: Test +test_enary_flatten = TestCase $ do + let tree = mkBin Add (var 0) (mkBin Add (var 1) (var 2)) + (eid, eg) = runEG $ fromTree myCost tree + ec = _eClass eg IntMap.! eid + case _best . _info $ ec of + ENAry EAdd xs -> do + let children = expandedList xs + assertEqual "enary: 3 children" 3 (length children) + assertBool "enary: distinct children" (length (nub children) == length children) + assertBool "enary: sorted children" (children == sort children) + _ -> assertFailure "enary: best should be a 3-ary ENAry EAdd" + +-- | Test 24: commutativity is structural (a+b ≡ b+a, no rules needed) +test_enary_comm :: Test +test_enary_comm = TestCase $ do + let ((c1, c2), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (var 0) (var 1)) + eid2 <- fromTree myCost (mkBin Add (var 1) (var 0)) + a <- canonical eid1 + b <- canonical eid2 + pure (a, b) + assertEqual "comm: a+b == b+a" c1 c2 + +-- | Test 25: associativity flattens (a+b)+c ≡ a+(b+c) ≡ a+(c+b) +test_enary_assoc :: Test +test_enary_assoc = TestCase $ do + let ((c1, c2, c3), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (mkBin Add (var 0) (var 1)) (var 2)) + eid2 <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 1) (var 2))) + eid3 <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 2) (var 1))) + a <- canonical eid1 + b <- canonical eid2 + c <- canonical eid3 + pure (a, b, c) + assertEqual "assoc: (a+b)+c == a+(b+c)" c1 c2 + assertEqual "assoc: (a+b)+c == a+(c+b)" c1 c3 + +-- | Test 26: multiset semantics (x+x is distinct from x) +test_enary_multiset :: Test +test_enary_multiset = TestCase $ do + let ((cX, cXX), _) = runEG $ do + eidX <- fromTree myCost (var 0) + eidXX <- fromTree myCost (mkBin Add (var 0) (var 0)) + a <- canonical eidX + b <- canonical eidXX + pure (a, b) + assertBool "multiset: x+x /= x" (cX /= cXX) + +-- | Test 27: constants fold inside flattened nodes (2+3+x ≡ 5+x) +test_enary_fold_const :: Test +test_enary_fold_const = TestCase $ do + let ((c1, c2), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (mkBin Add (constv 2.0) (constv 3.0)) (var 0)) + eid2 <- fromTree myCost (mkBin Add (constv 5.0) (var 0)) + a <- canonical eid1 + b <- canonical eid2 + pure (a, b) + assertEqual "fold-const: 2+3+x == 5+x" c1 c2 + +-- | Test 28: direct add of an unsorted ENAry canonicalizes and folds consts +test_enary_direct_add :: Test +test_enary_direct_add = TestCase $ do + let ((c1, c2), _) = runEG $ do + e2 <- fromTree myCost (constv 2.0) + e3 <- fromTree myCost (constv 3.0) + ex <- fromTree myCost (var 0) + eid <- add myCost (ENAry EAdd (imFromList [e3, ex, e2])) + eid5x <- fromTree myCost (mkBin Add (constv 5.0) (var 0)) + a <- canonical eid + b <- canonical eid5x + pure (a, b) + assertEqual "direct add: ENAry [3,x,2] sorts and folds to 5+x" c1 c2 + +-- | Test 29: extraction of a flattened class right-folds to a binary tree +test_enary_extract :: Test +test_enary_extract = TestCase $ do + let t1 = mkBin Add (var 0) (mkBin Add (var 1) (var 2)) + (extracted, _) = runEG $ do + eid <- fromTree myCost t1 + getBestExpr eid + assertEqual "extract: flattened a+b+c == a+(b+c)" (showExpr t1) (showExpr extracted) + +-- | Test 30: merge cascade propagates through ENAry parents (a≡b -> a+c ≡ b+c) +test_enary_merge_cascade :: Test +test_enary_merge_cascade = TestCase $ do + let ((c1, c2), _) = runEG $ do + ea <- fromTree myCost (var 0) + eb <- fromTree myCost (var 1) + _ <- fromTree myCost (var 2) + eac <- fromTree myCost (mkBin Add (var 0) (var 2)) + ebc <- fromTree myCost (mkBin Add (var 1) (var 2)) + merge myCost ea eb + rebuild myCost + a <- canonical eac + b <- canonical ebc + pure (a, b) + assertEqual "cascade: after a==b, a+c == b+c" c1 c2 + +-- | Soundness: a closed 2-ary pattern (a+b) does NOT match a 3-ary multiset. +test_match_closed2_not_3ary :: Test +test_match_closed2_not_3ary = TestCase $ do + let pat = "a" + "b" + (substs, _) = runEG $ do + x <- fromTree myCost (var 0) + y <- fromTree myCost (var 1) + z <- fromTree myCost (var 2) + _ <- add myCost (ENAry EAdd (imFromList [x, y, z])) + match pat + assertBool "closed2: a+b does not match x+y+z" (null substs) + +-- | Soundness: a+a does NOT match x+x+y (only exact multisets match). +test_match_aa_not_3ary :: Test +test_match_aa_not_3ary = TestCase $ do + let pat = "a" + "a" + (substs, _) = runEG $ do + _ <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 0) (var 1))) + match pat + assertBool "aa: a+a does not match x+x+y" (null substs) + +-- | B3: 0 + x + y = x + y (n-ary open-rest rule). +test_eqsat_zero_plus_rest :: Test +test_eqsat_zero_plus_rest = TestCase $ do + let tree = addZero (constv 0.0) (addZero (var 0) (var 1)) + assertEqual "0+x+y = x+y" + (showExpr (var 0 + var 1)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B7: xy + xz + w = x(y+z) + w (n-ary factoring with a rest variable). +test_eqsat_factoring :: Test +test_eqsat_factoring = TestCase $ do + let tree = ((var 0 * var 1) + (var 0 * var 2)) + var 3 + assertEqual "xy+xz+w = x(y+z)+w" + (showExpr ((var 0 * (var 1 + var 2)) + var 3)) + (showExpr (simplifyEqSatDefault tree)) + +-- | C9 is a closed 2-ary rule: (x+y+z)^2 is NOT expanded to a binomial. +test_eqsat_binomial_closed2 :: Test +test_eqsat_binomial_closed2 = TestCase $ do + let tree = ((var 0 + var 1) + var 2) ** constv 2.0 + assertEqual "(x+y+z)^2 not expanded" + (showExpr ((var 0 + (var 1 + var 2)) ** constv 2.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | C14: sqrt(x*x) = abs x (closed 2-ary multiset). +test_eqsat_sqrt_square :: Test +test_eqsat_sqrt_square = TestCase $ do + let rule = sqrt (NAry EMul [Ch "x", Ch "x"]) :=> abs "x" + (best, _) = runEG $ eqSat (sqrt (var 0 * var 0)) [rule] myCost 5 + assertEqual "sqrt(x*x) = abs x" (showExpr (abs (var 0))) (showExpr best) + +-- | x/x = 1 and x-x = 0 (constant identities). +test_eqsat_identities :: Test +test_eqsat_identities = TestCase $ do + assertEqual "x/x = 1" (showExpr (constv 1.0)) (showExpr (simplifyEqSatDefault (var 0 / var 0))) + assertEqual "x-x = 0" (showExpr (constv 0.0)) (showExpr (simplifyEqSatDefault (var 0 - var 0))) + +-- | helper: run eqSat with the full rule set and collect every expression +-- in the root eclass (used to assert that a rule "fires" even if a cheaper +-- representative is extracted). +allExprsOf :: Fix SRTree -> [Fix SRTree] +allExprsOf t = fst $ runEG $ do + root <- fromTree myCost t + _ <- runEqSat myCost rewrites 20 + getAllExpressionsFrom root + +-- | C11 fires: log(x*y) expands to log x + log y inside the root eclass. +test_eqsat_log_distributes :: Test +test_eqsat_log_distributes = TestCase $ do + let exprs = allExprsOf (log (var 0 * var 1)) + target = showExpr (log (var 0) + log (var 1)) + assertBool "log(x*y) contains log x + log y" + (any (\e -> showExpr e == target) exprs) + +-- | C12 fires: abs(x*y) expands to abs x * abs y inside the root eclass. +test_eqsat_abs_distributes :: Test +test_eqsat_abs_distributes = TestCase $ do + let exprs = allExprsOf (abs (var 0 * var 1)) + target = showExpr (abs (var 0) * abs (var 1)) + assertBool "abs(x*y) contains abs x * abs y" + (any (\e -> showExpr e == target) exprs) + +-- | C13 fires: (x*y)^z expands to x^z * y^z inside the root eclass. +test_eqsat_pow_distributes :: Test +test_eqsat_pow_distributes = TestCase $ do + let exprs = allExprsOf ((var 0 * var 1) ** constv 2.0) + target = showExpr ((var 0 ** constv 2.0) * (var 1 ** constv 2.0)) + assertBool "(x*y)^2 contains x^2 * y^2" + (any (\e -> showExpr e == target) exprs) + +-- | B9 (a :==: rule): x^2 * x^3 = x^5. +test_eqsat_pow_mul :: Test +test_eqsat_pow_mul = TestCase $ do + let tree = (var 0 ** constv 2.0) * (var 0 ** constv 3.0) + assertEqual "x^2*x^3 = x^5" (showExpr (var 0 ** constv 5.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B11 (a :==: rule): (x^2)^3 = x^6. +test_eqsat_pow_pow :: Test +test_eqsat_pow_pow = TestCase $ do + let tree = (var 0 ** constv 2.0) ** constv 3.0 + assertEqual "(x^2)^3 = x^6" (showExpr (var 0 ** constv 6.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | x^y * x = x^(y+1): x^2 * x = x^3. +test_eqsat_pow_mul_x :: Test +test_eqsat_pow_mul_x = TestCase $ do + let tree = (var 0 ** constv 2.0) * var 0 + assertEqual "x^2*x = x^3" (showExpr (var 0 ** constv 3.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B4: (0*x)*y = 0. +test_eqsat_zero_mul :: Test +test_eqsat_zero_mul = TestCase $ do + let tree = mkBin Mul (mkBin Mul (constv 0.0) (var 0)) (var 1) + assertEqual "(0*x)*y = 0" (showExpr (constv 0.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B4 guard: (0*NaN)*x is NOT folded to 0 (NaN invalidates the rest). +test_eqsat_zero_mul_nan :: Test +test_eqsat_zero_mul_nan = TestCase $ do + let tree = mkBin Mul (mkBin Mul (constv 0.0) (constv (0/0))) (var 0) + best = simplifyEqSatDefault tree + assertBool "(0*NaN)*x /= 0" (showExpr best /= showExpr (constv 0.0)) + +-- | rewritesParams: x-x and x/x become Param 0. +test_eqsat_params :: Test +test_eqsat_params = TestCase $ do + let (b1, _) = runEG $ eqSat (var 0 - var 0) rewritesParams myCost 10 + (b2, _) = runEG $ eqSat (var 0 / var 0) rewritesParams myCost 10 + assertEqual "x-x = Param 0 (param mode)" (showExpr (param 0)) (showExpr b1) + assertEqual "x/x = Param 0 (param mode)" (showExpr (param 0)) (showExpr b2) + +-- | Soundness: x*x*y stays as a right-folded Mul, NOT x^2 (B1 is 2-ary only). +test_eqsat_xxy_sound :: Test +test_eqsat_xxy_sound = TestCase $ do + let tree = mkBin Mul (mkBin Mul (var 0) (var 0)) (var 1) + assertEqual "x*x*y stays right-folded" + (showExpr (var 0 * (var 0 * var 1))) + (showExpr (simplifyEqSatDefault tree)) + +-- | Completeness: a*b matches every Mul node inside a merged class. +test_match_complete_multinode :: Test +test_match_complete_multinode = TestCase $ do + let pat = "a" * "b" + (n, _) = runEG $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 2) + _ <- fromTree myCost (var 3) + m1 <- fromTree myCost (var 0 * var 1) + m2 <- fromTree myCost (var 2 * var 3) + _ <- merge myCost m1 m2 + rebuild myCost + s <- match pat + pure (length s) + assertBool "complete: a*b yields all substs in a merged class" (n >= 2) + +-- | helper: find all non-canonical eclass ids in the trie +getAllStaleTrieKeys :: IntMap.IntMap Int -> DB -> [EClassId] +getAllStaleTrieKeys can = concatMap goIntTrie . Map.elems + where + goIntTrie (IntTrie m) = + [k | k <- IntMap.keys m, not (isCanon k)] + ++ concatMap goIntTrie (IntMap.elems m) + isCanon eid = case IntMap.lookup eid can of + Just v -> v == eid + Nothing -> False + +prependLabel :: String -> Test -> Test +prependLabel label t = TestLabel label t + +tests :: Test +tests = TestList + [ prependLabel "fromTree-var" test_fromTree_var + , prependLabel "fromTree-bin" test_fromTree_bin + , prependLabel "canonical-identity" test_canonical_identity + , prependLabel "canonize" test_canonize + , prependLabel "add-duplicate" test_add_duplicate + , prependLabel "merge" test_merge + , prependLabel "rebuild" test_rebuild + , prependLabel "match" test_match + , prependLabel "getBestExpr" test_getBestExpr + , prependLabel "eqsat-x+0" test_eqsat_x_plus_0 + , prependLabel "eqsat-x*1" test_eqsat_x_times_1 + , prependLabel "fitness-theta" test_fitness_theta + , prependLabel "fitness-range" test_fitness_range + , prependLabel "top-fit-size" test_top_fit_size + , prependLabel "eqsat-comm" test_eqsat_comm + , prependLabel "eqsat-double-neg" test_eqsat_double_neg + , prependLabel "fromTrees" test_fromTrees + , prependLabel "cost" test_cost + , prependLabel "getAllExpressions" test_get_all_expr + , prependLabel "sizeFitDB-no-stale" test_sizeFitDB_no_stale + , prependLabel "trie-no-stale-keys" test_trie_no_stale_keys + , prependLabel "match-after-merge" test_match_after_merge_multi_atom + , prependLabel "enary-flatten" test_enary_flatten + , prependLabel "enary-comm" test_enary_comm + , prependLabel "enary-assoc" test_enary_assoc + , prependLabel "enary-multiset" test_enary_multiset + , prependLabel "enary-fold-const" test_enary_fold_const + , prependLabel "enary-direct-add" test_enary_direct_add + , prependLabel "enary-extract" test_enary_extract + , prependLabel "enary-merge-cascade" test_enary_merge_cascade + , prependLabel "match-closed2-3ary" test_match_closed2_not_3ary + , prependLabel "match-aa-not-3ary" test_match_aa_not_3ary + , prependLabel "eqsat-0+rest" test_eqsat_zero_plus_rest + , prependLabel "eqsat-factoring" test_eqsat_factoring + , prependLabel "eqsat-binomial-2ary" test_eqsat_binomial_closed2 + , prependLabel "eqsat-sqrt-square" test_eqsat_sqrt_square + , prependLabel "eqsat-identities" test_eqsat_identities + , prependLabel "eqsat-log-dist" test_eqsat_log_distributes + , prependLabel "eqsat-abs-dist" test_eqsat_abs_distributes + , prependLabel "eqsat-pow-dist" test_eqsat_pow_distributes + , prependLabel "eqsat-pow-mul" test_eqsat_pow_mul + , prependLabel "eqsat-pow-pow" test_eqsat_pow_pow + , prependLabel "eqsat-pow-mul-x" test_eqsat_pow_mul_x + , prependLabel "eqsat-0*mul" test_eqsat_zero_mul + , prependLabel "eqsat-0*mul-NaN" test_eqsat_zero_mul_nan + , prependLabel "eqsat-params" test_eqsat_params + , prependLabel "eqsat-x*x*y-sound" test_eqsat_xxy_sound + , prependLabel "match-complete" test_match_complete_multinode + ] diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,4 +1,115 @@ -import Test.HUnit +import Test.HUnit +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as VS +import Data.SRTree.Internal +import Data.SRTree.Recursion (Fix) +import Data.SRTree.Eval (compile) +import Algorithm.SRTree.AD.Unboxed (CompiledTree, compileTree, compileTreeMulti, evalGrad, evalGradVec, evalGradMulti) +import qualified EqSatTests +import qualified StoreTests +import Data.SRTree.Random (randomTree, tossBiased, randomFrom) +import System.Random (mkStdGen) +import Control.Monad.State.Strict (evalStateT) +import Data.SRTree.Datasets (loadDataset) +import Control.Monad (forM_) +-- Small epsilon compare for Doubles +eps :: Double +eps = 1e-9 + +approxEqual :: [Double] -> [Double] -> Bool +approxEqual a b = and $ zipWith (\x y -> abs (x - y) < eps) a b + +test_compile :: Test +test_compile = TestCase $ do + let xss = [VU.fromList [1.0, 2.0, 3.0]] + tree = var 0 * param 0 + param 1 + theta = VU.fromList [2.0, 0.5] + yhat = compile xss tree theta + got = VU.toList yhat + expected = [2.5, 4.5, 6.5] + assertBool ("compile produced " ++ show got ++ " expected " ++ show expected) (approxEqual got expected) + +-- Gradient correctness: the compact ctStatic layout must agree with finite +-- differences (objective) and with the row-fused `evalGrad` backend across +-- the vectorized `evalGradVec` and chunked `evalGradMulti` paths. +test_grad :: Test +test_grad = TestCase $ do + let xss = [ VU.fromList [1.0, 2.0, 3.0, 4.0] + , VU.fromList [0.5, 1.5, 2.5, 3.5] + , VU.fromList [2.0, 1.0, 0.5, 0.25] ] + y = VU.fromList [3.1, 5.2, 7.3, 9.4] + -- ((x0 + t0) * exp(x1)) / (x2 + t1) -- mixes static and dynamic subtrees + tree = (var 0 + param 0) * exp (var 1) / (var 2 + param 1) + theta = VS.fromList [1.0, 0.5] + ct = compileTree xss y Nothing tree + cts = compileTreeMulti xss y Nothing tree + (f0, g0) = evalGrad ct theta + (f1, g1) = evalGradVec ct theta + (f2, g2) = evalGradMulti cts theta + -- finite-difference gradient + h = 1e-6 + gfd = VS.toList $ VS.generate (VS.length theta) $ \i -> + let e = VS.fromList (map (\j -> if j == i then h else 0) [0 .. VS.length theta - 1]) + (fp, _) = evalGradVec ct (VS.zipWith (+) theta e) + (fm, _) = evalGradVec ct (VS.zipWith (-) theta e) + in (fp - fm) / (2 * h) + assertBool "evalGradVec objective != evalGrad" (abs (f1 - f0) < 1e-6) + assertBool "evalGradMulti objective != evalGrad" (abs (f2 - f0) < 1e-6) + assertBool "evalGradVec gradient != finite diff" + (and (zipWith (\a b -> abs (a - b) < 1e-4) (VS.toList g1) gfd)) + assertBool "evalGrad gradient != finite diff" + (and (zipWith (\a b -> abs (a - b) < 1e-4) (VS.toList g0) gfd)) + +test_benchgrad :: Test +test_benchgrad = TestCase $ do + let genTerm = do coin <- tossBiased 0.4 + if coin then randomFrom [Fix $ Var ix | ix <- [0..8]] else randomFrom [Fix $ Param ix | ix <- [0..9]] + genNonTerm = randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Uni LogAbs (), Uni SqrtAbs ()] + genMultipleTrees 0 = pure [] + genMultipleTrees n = do + t <- randomTree 5 10 150 genTerm genNonTerm False + ts <- genMultipleTrees (n-1) + pure (t:ts) + g = mkStdGen 42 + trees' <- evalStateT (genMultipleTrees 5) g + ((dataset, y, _, _), _, _, _) <- loadDataset "data.tsv" True + let thetaU = VU.fromList [1.0, 0.5, 0.2, 0.3, 0.1, 0.5, 0.9, 0.3, 0.2, 0.4] + thetaS = VS.convert thetaU + trees = map relabelParamsOrder $ filter (\t -> let v = VU.sum (compile dataset t thetaU) in not (isInfinite v || isNaN v)) trees' + h = 1e-6 + gfd :: CompiledTree -> VS.Vector Double + gfd ct = VS.generate (VS.length thetaS) $ \i -> + let e = VS.fromList (map (\j -> if j == i then h else 0) [0 .. VS.length thetaS - 1]) + (fp, _) = evalGradVec ct (VS.zipWith (+) thetaS e) + (fm, _) = evalGradVec ct (VS.zipWith (-) thetaS e) + in (fp - fm) / (2 * h) + forM_ (zip [0..] trees) $ \(i, t) -> do + let ct = compileTree dataset y Nothing t + cts = compileTreeMulti dataset y Nothing t + (f1, g1) = evalGradVec ct thetaS + (f0, g0) = evalGrad ct thetaS + (f2, g2) = evalGradMulti cts thetaS + fd = gfd ct + putStrLn ("benchgrad tree " ++ show i ++ " obj=" ++ show f1) + assertBool ("tree " ++ show i ++ " evalGradVec objective != evalGrad") (abs (f1 - f0) < 1e-6 * max 1 (abs f0)) + assertBool ("tree " ++ show i ++ " evalGradMulti objective != evalGrad") (abs (f2 - f0) < 1e-6 * max 1 (abs f0)) + assertBool ("tree " ++ show i ++ " evalGradVec gradient mismatch") (and (zipWith (\a b -> abs (a - b) < 1e-3 * max 1 (abs a)) (VS.toList g1) (VS.toList fd))) + assertBool ("tree " ++ show i ++ " evalGrad gradient mismatch") (and (zipWith (\a b -> abs (a - b) < 1e-3 * max 1 (abs a)) (VS.toList g0) (VS.toList fd))) + assertBool ("tree " ++ show i ++ " evalGradMulti gradient != evalGrad") (and (zipWith (\a b -> abs (a - b) < 1e-9 * max 1 (abs a)) (VS.toList g0) (VS.toList g2))) + main :: IO () -main = pure () +main = do + let t1 = TestLabel "compile" test_compile + t2 = TestLabel "grad" test_grad + + counts <- runTestTT $ TestList + [ t1 + , t2 + , TestLabel "benchgrad" test_benchgrad + , TestLabel "eqsat" EqSatTests.tests + , TestLabel "store" StoreTests.tests + ] + if failures counts /= 0 || errors counts /= 0 + then error "Some tests failed" + else pure () diff --git a/test/StoreTests.hs b/test/StoreTests.hs new file mode 100644 --- /dev/null +++ b/test/StoreTests.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE TupleSections #-} + +module StoreTests where + +import Test.HUnit +import Data.SRTree +import qualified Data.IntMap as IntMap +import qualified Data.HashMap.Strict as HashMap +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Store +import Control.Monad.State.Strict +import Control.Monad.Identity + +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 + +-- | run a stateful computation on a specific graph +runIn :: EGraph -> EGraphST Identity a -> (a, EGraph) +runIn g m = runIdentity $ runStateT m g + +evalIn :: EGraph -> EGraphST Identity a -> a +evalIn g m = runIdentity $ evalStateT m g + +-- | graph A: x0, x1, x0+x1 (with fitness on the sum) +buildA :: (EClassId, EGraph) +buildA = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + eidSum <- fromTree myCost (var 0 + var 1) + insertFitness eidSum 0.5 [] + pure eidSum + +-- | graph B: x1, x0+x1, (x0+x1)*x2 (shares x1 and x0+x1 with A) +buildB :: EGraph +buildB = snd $ runIn emptyGraph $ do + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 0 + var 1) + _ <- fromTree myCost ((var 0 + var 1) * var 2) + pure () + +-- | pattern (x0+x1)*x2 = (A + B) * C +prodPattern :: Pattern +prodPattern = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'A') (VarPat 'B'))) (VarPat 'C')) + +-- | Test 1: export/import round-trip preserves the rows exactly +test_roundtrip :: Test +test_roundtrip = TestCase $ do + let (_, g) = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 0 + var 1) + _ <- fromTree myCost ((var 0 + var 1) * var 2) + pure () + rows = exportEGraph g + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right g' -> do + let rows' = exportEGraph g' + assertBool "round-trip: rows differ" (rows == rows') + assertBool "round-trip: class count" (IntMap.size (_grEClasses rows) == IntMap.size (_grEClasses rows')) + assertBool "round-trip: node count" (HashMap.size (_grENodeToEClass rows) == HashMap.size (_grENodeToEClass rows')) + +-- | Test 2: round-trip preserves fitness and rebuilds the range DB +test_roundtrip_fitness :: Test +test_roundtrip_fitness = TestCase $ do + let (sumEid, g) = runIn emptyGraph $ do + eidSum <- fromTree myCost (var 0 + var 1) + insertFitness eidSum 0.42 [] + pure eidSum + rows = exportEGraph g + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right g' -> do + let fit = evalIn g' (getFitness sumEid) + assertEqual "round-trip: fitness" (Just 0.42) fit + let mx = getGreatest (_fitRangeDB (_eDB g')) + assertEqual "round-trip: fitRangeDB max" (Just (0.42, sumEid)) mx + -- a node added *after* import dedups against the loaded graph (no dup class) + let (eidNew, g'') = runIn g' $ fromTree myCost (var 0 + var 1) + nClasses = IntMap.size (_eClass g'') + assertBool "post-import dedup adds no class" (eidNew == sumEid && nClasses == IntMap.size (_eClass g')) + +-- | Test 3: import rejects inconsistent rows +test_import_invalid :: Test +test_import_invalid = TestCase $ do + let (_, g) = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + pure () + rows = exportEGraph g + bad = rows { _grENodeToEClass = HashMap.insert (EVar 0) 999 (_grENodeToEClass rows) } -- 999 not in canonical map + case importEGraph bad of + Left _ -> pure () + Right _ -> assertFailure "invalid rows should have been rejected" + +-- | Test 4: merge dedups shared structure and adds only new classes +test_merge :: Test +test_merge = TestCase $ do + let (sumEidA, gA) = buildA + gM = case mergeEGraph myCost gA buildB of + Left err -> error ("merge failed: " ++ err) + Right g -> g + nA = IntMap.size (_eClass gA) + nM = IntMap.size (_eClass gM) + assertEqual "merge: adds only classes absent from A (x2, product)" (nA + 2) nM + -- B's unique expression (x0+x1)*x2 is present and matchable + let nMatch = length $ evalIn gM (match prodPattern) + assertBool "merge: B's unique expression present" (nMatch > 0) + -- A's fitness on the shared sum class is preserved (same canonical id) + assertEqual "merge: A fitness preserved" (Just 0.5) (evalIn gM (getFitness sumEidA)) + +-- | Test 5: merge preserves round-trip +test_merge_roundtrip :: Test +test_merge_roundtrip = TestCase $ do + let (_, gA) = buildA + gM = case mergeEGraph myCost gA buildB of + Left err -> error ("merge failed: " ++ err) + Right g -> g + rows = exportEGraph gM + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right gM' -> assertBool "merge round-trip: rows differ" (exportEGraph gM' == rows) + +-- | Test 6: stale node->class entries (a node pointing at a class whose +-- canonical representative is another class) are canonicalized on import +test_import_stale_canonicalizes :: Test +test_import_stale_canonicalizes = TestCase $ do + let (keep, g) = buildA -- keep = x0+x1, a root class, has fitness + rows0 = exportEGraph g + dead = _grNextId rows0 -- a fresh id not yet in the graph + rows = rows0 { _grCanonical = IntMap.insert dead keep (_grCanonical rows0) + , _grEClasses = IntMap.insert dead + (IntMap.findWithDefault (error "keep missing") keep (_grEClasses rows0)) + (_grEClasses rows0) + , _grENodeToEClass = HashMap.insert (EBin Add 2 3) dead (_grENodeToEClass rows0) + , _grNextId = dead + 1 } + case importEGraph rows of + Left err -> assertFailure ("import of stale rows failed: " ++ err) + Right g' -> do + let canon = _grCanonical (exportEGraph g') + posts = exportEGraph g' + deadNext = IntMap.lookup dead (_grEClasses posts) + -- the dead class is gone and every node points at a canonical class + assertEqual "dead class dropped" Nothing deadNext + assertBool "all node->class values canonical" + (all (\eid -> IntMap.lookup eid canon == Just eid) (HashMap.elems (_grENodeToEClass posts))) + -- the kept class is still there with its fitness (via the fit range db) + assertEqual "kept fitness preserved" (Just 0.5) (evalIn g' (getFitness keep)) + +prependLabel :: String -> Test -> Test +prependLabel label t = TestLabel label t + +tests :: Test +tests = TestList + [ prependLabel "store-roundtrip" test_roundtrip + , prependLabel "store-roundtrip-fit" test_roundtrip_fitness + , prependLabel "store-import-invalid" test_import_invalid + , prependLabel "store-merge" test_merge + , prependLabel "store-merge-roundtrip" test_merge_roundtrip + , prependLabel "store-stale-canon" test_import_stale_canonicalizes + ]