diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,19 @@ # Changelog for srtree +## 2.0.0.0 + +- Complete refactoring of the library +- Integration of other tools such as: srtree-opt, srtree-tools, srsimplify +- Implementation of Equality Saturation and support to e-graph +- Using Massiv for performance +- Using NLOpt as the optimization library + +## 1.1.0.0 + +- Reorganization of modules +- Renaming AD functions +- Inclusion of reverse mode that calculates the diagonal of and the full Hessian matrices + ## 1.0.0.5 - Changed `base` and `mtl` versions diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,25 +1,274 @@ -# srtree: A symbolic regression expression tree structure. +# srtree: A supporting library for tree-based symbolic regression -`srtree` is a Haskell library with a data structure and supporting functions to manipulate expression trees for symbolic regression. +`srtree` is a Haskell library that implements a tree-based structure for expressions and supporting functions to be used in the context of **symbolic regression**. -The tree-like structure is defined as a fixed-point of an n-ary tree. The variables and parameters of the regression model are indexed as `Int`type and the constant values are `Double`. +The expression structure is defined as a fixed-point of a mix of unary and binary tree. This makes it easier to implement supporting functions that requires the traversal of the trees. Also, since it is a parameterized structure, we can creating partial trees to pattern math structures of interest. +This structure may contain four types of nodes: -The tree supports leaf nodes containing a variable, a free parameter, or a constant value; internal nodes that represents binary operators such as the four basic math operations, logarithm with custom base, and the power of two expressions; and unary functions specified by `Function` data type. +- `Bin Op l r` that represents a binary operator `Op` with two children. +- `Uni Function t` that represents an unary function `Function` with a single child. +- `Var Int` representing the index of a variable (i.e., x0, x1, etc.). +- `Param Int` representing the index of a adjustable parameter (i.e., theta0, theta1, etc.). +- `Const Double` representing a constant value. -The `SRTree` structure has instances for `Num, Fractional, Floating` which allows to create an expression as a valid Haskell expression such as: +The `SRTree` structure has instances for `Num, Fractional, Floating, IsString` which allows to create an expression as a valid Haskell expression such as (remember to turn on OverloadedStrings extension): + ```haskell -x = var 0) -y = var 1 -expr = x * 2 + sin(y * pi + x) :: Fix SRTree +expr = "x0" * 2 + sin("x1" * pi + "x0") :: Fix SRTree ``` -## Other features: +This library comes with support to many quality of life functions to handle this data structure. Such as: -- derivative w.r.t. a variable (`deriveByVar`) and w.r.t. a parameter (`deriveByParam`) -- evaluation (`evalTree`) -- relabel free parameters sequentially (`relabelParams`) -- gradient calculation with `forwardMode`, or optimized with `gradParams` if there is only a single occurrence of each parameter (most of the cases). +- getting the arity of a node +- getting the children of a node as a list +- count the number of nodes +- number of nodes of a specific type +- counting unique tokens +- number of variables and parameters +- relabeling the parameters from 0 to p +- converting floating point constants to parameters + +Additionally, the library provides supporting function to work with datasets, evaluating the expressions, +calculating the derivatives, printing, generating random trees, simplifying the expression, calculating overall statistics, +optimizing parameters, and model selection metrics. + +Together with this library, we provide example applications (please refer to their corresponding README files): + +- [srsimplify](apps/srsimplify/README.md): a parser and simplification tool supporting the output of many popular SR algorithms. +- [srtools](apps/srtools/README.md): a tool that can be used to evaluate symbolic regression expressions and create nice reports with confidence intervals. +- [tinygp](apps/tinygp/README.md): a simple GP implementation based on tinyGP. + +## Organization + +The library is organized as `Data`, `Algorithm`, and `Text` modules where the `Data` modules implement functions directly tied to the data structure and the `Algorithm` modules implement algorithms related to symbolic regression, finally, the `Text` modules parse string expressions from different formats and apply simplification, when requested. + +### `Data` modules + +The `Data` modules is split into $5$ submodules: + +- `Data.SRTree` contains the data strucuture and basic supporting functions. +- `Data.SRTree.Datasets` contains functions supporting loading datasets into Massiv.Arrays (aka numpy arrays). +- `Data.SRTree.Derivative` contains the symbolic derivatives of the functions and operators. +- `Data.SRTree.Eval` contains functions to evaluate the tree given a dataset and parameters. +- `Data.SRTree.Print` contains supporting functions for converting trees to different string representation. +- `Data.SRTree.Random` contains functions to generate random trees. + +#### `Data.SRTree` + +The `SRTree val` data structure is a sum type structure that can be either a variable index, a parameter index, a constant value (of type `Double`), an univariate function or a binary operator. The data type is implemented as a fixed point so all the algorithms act on `Fix SRTree`: + +```haskell +t = "x0" + "t0" * sin("x1" + "t1"**2) :: Fix SRTree +``` + +When creating the expression in a more natural notation, the variables and parameters are `String` composed of the first letter either `x`, for variables, or `t` for parameters (as in theta), and an integer corresponding to the index of the variable or parameter. The fixed point notation, allows us to implment recursive processing of a tree without many of the common boilerplate: + +```haskell +countNodes = + \case + Var _ = 1 + Const _ = 1 + Param _ = 1 + Uni _ t = 1 + t + Bin _ l r = 1 + l + r +``` + +The children are parameterized by the `val` type parameter. This allows us to create convenient partial structures, such as: + +```haskell +-- + operator pointing to some structure +-- with index 1 and 2 +Bin Add 1 2 + +-- canonical representation of + operator +Bin Add () () +``` + +The main functions of this module are: + +- `arity`: returns the arity of an operator. +- `getChildren`: returns the children of a `Fix SRTree` as a list +- `countNodes`: returns the number of nodes +- `countOccurrences`: counts the occurence of a given variable +- `countVars`: returns the number of unique variables appearing the expression +- `relabelParams`: relabels the parameters from the left leaves to the right +- `constsToParams`: replace `Const` nodes with `Param` nodes. + +#### `Data.SRTree.Datasets` module + +This module exports only the `loadDataset` function which takes a filename and +returns the training and test sets together with the column labels. +The filename must follow the format: + +`filename.ext:start_row:end_row:target:features` + +where each ':' field is optional. The fields are: + +- **start_row:end_row** is the range of the training rows (default 0:nrows-1). + every other row not included in this range will be used as validation +- **target** is either the name of the PVector (if the datafile has headers) or the index + 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. + +Example of valid names: `dataset.csv`, `mydata.tsv`, `dataset.csv:20:100`, `dataset.tsv:20:100:price:m2,rooms,neighborhood`, `dataset.csv:::5:0,1,2`. + +#### `Data.SRTree.Derivative` module + +Calculates symbolic derivatives of the expression w.r.t. the variables or the parameters. +The main functions of this module are: + +- `deriveBy`: returns the symbolic derivative w.r.t. a certain variable or a certain parameter. +- `deriveByVar`: shortcut to `deriveBy` to derive by a variable. +- `deriveByParam`: shortcut to `deriveBy` to derive by a parameter. + +#### `Data.SRTree.Eval` module + +Evaluates an expression given a dataset. +The main functions of this module are: + +- `evalTree`: given a data matrix and a vector of parameters, evaluates the expression tree. +- `evalInverse`: evaluates the inverse of a function. +- `invright`: evaluates the right inverse of an operator. +- `invleft`: evaluates the left inverse of an operator. + +#### `Data.SRTree.Print` module + +Support functions to convert an expression tree into a `String`. +The main functions of this module are: + +- `showExpr` and `printExpr`: converts/print the expression into math notation . +- `showPython` and `printPython`: converts/print to a numpy notation. +- `showLatex` and `printLatex`: converts/print to a LaTeX notation. +- `showTikz` and `printTikz`: converts/print to a TikZ notation. + +#### `Data.SRTree.Random` module + +Auxiliary functions to create random trees. +The main functions of this module are: + +- `randomTree`: creates a random tree with a certain number of nodes. +- `randomTreeBalanced`: creates a (almost) balanced random tree with a certain number of nodes. + +### `Text` modules + +The `Text` module is split into $2$ modules: + +- `Text.ParseSR`: contains the main parsers for different SR algorithms output. +- `Text.ParseSR.IO`: auxiliary functions to handle files containing many expressions. + +#### `Text.ParseSR` module + +The only important function of this module is `parseSR` that parses an string expression from a given algorithm to a certain output. It also converts variable names to x0, x1,... + +#### `Text.ParseSR.IO` module + +The two main functions of this module are: + +- `withInput`: that reads the stdin or a text file and parse all expressions +- `withOutput`: that writes the parsed expression into stdout or a file with one of the choices of output format. + +These functions handle any errors with an `Either` type and they can be safely pipelined together. Any invalid expression will be printed as "invalid expression ". + +### `Algorithm` modules + +The `Algorithm` modules are split into $5$ submodules: + +- `Algorithm.SRTree.AD` contains automatic differentiation functions. +- `Algorithm.SRTree.ConfidenceIntervals` contains functions to calculate the confidence intervals of parameters and predictions of a symbolic expression using Laplace approximation or profile likelihood. +- `Algorithm.SRTree.Likelihood` contains functions support different likelihood functions and their derivatives (gradient and hessian). +- `Algorithm.SRTree.ModelSelection` implements different model selection criteria such as AIC, BIC, MDL. +- `Algorithm.SRTree.Opt` implements functions to optimize the parameters of an expression supporting different likelihood functions. + +#### `Algorithm.SRTree.AD` module + +The main functions of this module are: + +- `forwardMode`: returns the prediction errors vector multiplied by the Jacobian matrix using forward mode AD. +- `forwardModeUnique`: same as above, but assuming each parameter index appear only once in the tree. +- `reverseModeUnique`: same as above, but using reverse mode +- `forwardModeUniqueJac`: same as `forwardModeUnique` but returns the Jacobian (does not mutiply by the error). + +#### `Algorithm.SRTree.Likelihood` module + +The main functions of this module are: + +- `sse, mse, rmse`: calculates the sum-of-square, mean squared, root of mean squared errors. +- `nll`: returns the negative log-likelihood given a distribution and the associated error (`Nothing` if unknown) +- `gradNLL`: returns the gradient of the negative log-likelihood. +- `gradNLLNonUnique`: same as above but assumes non-unique parameters +- `hessianNLL`: returns the hessian of the neg log-likelihood. + +#### `Algorithm.SRTree.Opt` module + +The main functions of this module are: + +- `minimizeNLL`: minimizes the negative log-likelihood of a distribution. +- `minimizeNLLNonUnique`: same as above but assumes repeated occurrences of parameters. +- `minimizeNLLWithFixedParam`: minimizes the neg log-likelihood but fixing the value of a single parameter. +- `minimizeGaussian`, `minimizePoisson`, `minimizeBinomial`: shortcut to minimize these three distributions. + +#### `Algorithm.SRTree.ModelSelection` module + +The main functions of this module are: + +- `bic`: Bayesian Information Criteria +- `aic`: Akaike Information Criteria +- `mdl`: Minimum Description Length as described in Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023) +- `mdlLattice`: 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. +- `mdlFreq` : MDL weighted by the frequency of occurrence of functions + +#### `Algorithm.SRTree.ConfidenceIntervals` module + +The main functions of this module are: + +- `paramCI`: calculates the parameters confidence intervals. +- `predictionCI`: calculates the predictions confidence intervals + +### `EqSat` modules + +The `EqSat` modules are split into $4$ submodules: + +- `Algorithm.EqSat.Simplify` contains function supporting algebraic simplification with equality saturation. +- `Algorithm.EqSat` contains the main equality saturation function. +- `Algorithm.EqSat.EGraph` contains the e-graph data structure and supporting functions. +- `Algorithm.EqSat.EqSatDB` contains supporting functions to pattern matching and insert equivalent expressions into an e-graph. + +#### `Algorithm.EqSat` module + +The main functions of this module are: + +- `eqSat` : runs equality saturation over a single expression. +- `getBest` : returns the best expression given the cost function used to generate the e-graph +- `recalculateBest` : recalculates the cost of each e-class using a new cost function +- `runEqSat` : runs equality saturation inside `EGraphST` monad. Use this if you want to return the e-graph. + +#### `Algorithm.EqSat.EGraph` module + +The main functions of this module are: + +- `fromTree` : creates an e-graph from an expression tree. +- `fromTrees` : creates an e-graph from multiple expressions +- `fromTreeWith` : inserts a new expression into the e-graph +- `findRootClasses` : returns the roots of the e-graph, if any . +- `getExpressionFrom` : returns a single expression from a given e-class always picking the first e-node as the path +- `getAllExpressionsFrom` : returns all expressions from the given e-class +- `getRndExpressionFrom` : returns a random expression from this e-class + + +#### `Algorithm.EqSat.EqSatDB` module + +The main functions of this module are: + +- TODO: create auxiliary functions to apply substution rules inside an EGraphST monad . + +#### `Algorithm.EqSat.Simplify` module + +The main functions of this module are: + +- `simplifyEqSatDefault` : simplifies an expression using the default parameters +- `simplifyEqSat` : simplifies with custom parameters ## TODO: diff --git a/apps/egraphGP/Main.hs b/apps/egraphGP/Main.hs new file mode 100644 --- /dev/null +++ b/apps/egraphGP/Main.hs @@ -0,0 +1,598 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE BangPatterns #-} + +module Main where + +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Simplify +import Algorithm.EqSat.Build +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Info +import Algorithm.EqSat.DB +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.Opt +import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.)) +import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_) +import Control.Monad.State.Strict +import qualified Data.IntMap.Strict as IM +import Data.Massiv.Array as MA hiding (forM_, forM) +import Data.Maybe (fromJust, isNothing, isJust) +import Data.SRTree +import Data.SRTree.Datasets +import Data.SRTree.Eval +import Data.SRTree.Random (randomTree) +import Data.SRTree.Print +import Options.Applicative as Opt hiding (Const) +import Random +import System.Random +import qualified Data.HashSet as Set +import Data.List ( sort, maximumBy, intercalate, sortOn ) +import Data.IntSet (IntSet) +import qualified Data.IntSet as IntSet +import qualified Data.Sequence as FingerTree +import Data.Function ( on ) +import qualified Data.Foldable as Foldable +import qualified Data.IntMap as IntMap +import List.Shuffle ( shuffle ) + +import Debug.Trace +import Algorithm.EqSat (runEqSat) + +-- Insert random expression +-- Evaluate random subtree +-- Insert new random parent eNode + +type RndEGraph a = EGraphST (StateT StdGen IO) a + +io = lift . lift +{-# INLINE io #-} +rnd = lift +{-# INLINE rnd #-} + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin op l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +data Alg = OnlyRandom | BestFirst deriving (Show, Read, Eq) + +-- experiment 1 80/30 +fitnessFun :: SRMatrix -> PVector -> SRMatrix -> PVector -> Fix SRTree -> RndEGraph (Double, PVector) +fitnessFun x y x_val y_val _tree = do + let tree = relabelParams _tree + nParams = countParams tree + thetaOrig <- rnd $ randomVec nParams -- = MA.replicate Seq nParams 1.0 + let (theta, fit) = minimizeNLL Gaussian Nothing 50 x y tree thetaOrig + tr = negate . mse x y tree $ if nParams == 0 then thetaOrig else theta + val = negate . mse x_val y_val tree $ if nParams == 0 then thetaOrig else theta + -- val = r2 x y tree $ if nParams == 0 then thetaOrig else theta + pure $ if isNaN val || isNaN tr + then (-1/0, theta) -- infinity + else (min tr val, theta) +{-# INLINE fitnessFun #-} + +fitnessFunRep :: SRMatrix -> PVector -> SRMatrix -> PVector -> Fix SRTree -> RndEGraph (Double, PVector) +fitnessFunRep x y x_val y_val _tree = do + fits <- replicateM 1 (fitnessFun x y x_val y_val _tree) + pure (maximumBy (compare `on` fst) fits) +{-# INLINE fitnessFunRep #-} + +-- helper query functions +fitnessIs p = p . _fitness . _info +{-# INLINE fitnessIs #-} + +getFitness :: EClassId -> RndEGraph (Maybe Double) +getFitness c = gets (_fitness . _info . (IM.! c) . _eClass) +{-# INLINE getFitness #-} +getSize :: EClassId -> RndEGraph Int +getSize c = gets (_size . _info . (IM.! c) . _eClass) +{-# INLINE getSize #-} +getSizeOf :: (Int -> Bool) -> [EClassId] -> RndEGraph [EClassId] +getSizeOf p = filterM (getSize >=> (pure . p)) +{-# INLINE getSizeOf #-} + +(&&&) p1 p2 x = p1 x && p2 x +{-# INLINE (&&&) #-} + +isValidFitness = fitnessIs (isJust &&& (not . isNaN . fromJust) &&& (not . isInfinite . fromJust)) +{-# INLINE isValidFitness #-} + +evaluated = fitnessIs isJust +{-# INLINE evaluated #-} +unevaluated' = fitnessIs isNothing +{-# INLINE unevaluated' #-} + +isSizeOf p = p . _size . _info +{-# INLINE isSizeOf #-} + +funDoesNotExistWith node = Prelude.any (not . (`sameFunAs` node) . snd) . _parents + where sameFunAs (Uni f _) (Uni g _) = f == g + sameFunAs _ _ = False +{-# INLINE funDoesNotExistWith #-} + +opDoesNotExistWith :: (SRTree ()) -> EClassId -> EClass -> Bool +opDoesNotExistWith node ecId = Prelude.any (not . (`sameOpAs` node) . snd) . _parents + where sameOpAs (Bin op1 l _) (Bin op2 _ _) = op1 == op2 && ecId == l + sameOpAs _ _ = False +{-# INLINE opDoesNotExistWith #-} + +rewriteBasic2 :: [Rule] +rewriteBasic2 = + [ + "x" * "y" :=> "y" * "x" + , "x" + "y" :=> "y" + "x" + , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x" + , ("x" + "y") + "z" :=> "x" + ("y" + "z") + , ("x" * "y") * "z" :=> "x" * ("y" * "z") + , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z") + , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z" + ] + +egraphSearch alg x y x_val y_val x_te y_te terms nEvals maxSize = do + ec <- insertRndExpr maxSize + updateIfNothing ec + insertTerms + evaluateUnevaluated + runEqSat myCost rewriteBasic2 1 + + while (numberOfEvalClasses nEvals) 1 $ + \radius -> + do + --nEvs <- gets (FingerTree.size . _fitRangeDB . _eDB) + nCls <- gets (IM.size . _eClass) + nUnev <- gets (IntSet.size . _unevaluated . _eDB) + let nEvs = nCls - nUnev + --io . print $ (nCls, nEvs) + bestF <- getBestFitness + + (ecN, b) <- case alg of + OnlyRandom -> do let ratio = fromIntegral nEvs / fromIntegral nCls + b <- rnd (tossBiased ratio) + ec <- if b && ratio > 0.99 then insertRndExpr maxSize >>= canonical else evaluateRndUnevaluated >>= canonical + pure (ec, False) + BestFirst -> do + ecsPareto <- getParetoEcsUpTo radius + ecsBest <- getBestEcs (isSizeOf (<=maxSize)) radius + + ecPareto <- combineFrom ecsPareto + curFitPareto <- getFitness ecPareto + + if isNothing curFitPareto + then pure (ecPareto, False) + else do ecBest <- combineFrom ecsBest + curFitBest <- getFitness ecBest + if isNothing curFitBest + then pure (ecBest, False) + else do ee <- evalRndSubTree + case ee of + Nothing -> do ec <- insertRndExpr maxSize >>= canonical + pure (ec, True) + Just c -> pure (c, False) + + upd <- updateIfNothing ecN + when (upd) + do runEqSat myCost rewriteBasic2 1 + cleanDB + pure () + if b then pure (min 20 $ radius+1) else pure (max 1 $ radius-1) + eclasses <- gets (IntMap.toList . _eClass) + -- forM_ eclasses $ \(_, v) -> (io.print) (Set.size (_eNodes v), Set.size (_parents v)) + paretoFront + --ft <- gets (_fitRangeDB . _eDB) + --io . print $ Foldable.toList ft + + where + numberOfEvalClasses :: Monad m => Int -> EGraphST m Bool + numberOfEvalClasses nEvs = + (subtract <$> gets (IntSet.size . _unevaluated . _eDB) <*> gets (IM.size . _eClass)) + >>= \n -> pure (n do + t <- getBest ec + (f, p) <- fitnessFunRep x y x_val y_val t + insertFitness ec f p + pure True + Just _ -> pure False + + getBestFitness = do + bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical) + gets (_fitness . _info . (IM.! bec) . _eClass) + + evalRndSubTree :: RndEGraph (Maybe EClassId) + evalRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB) + if not (null ecIds) + then do rndId <- rnd $ randomFrom ecIds + Just <$> canonical rndId + else pure Nothing + + + combineFrom ecs = do + nt <- rnd rndNonTerm + p1 <- rnd (randomFrom ecs) + p2 <- rnd (randomFrom ecs) + l1 <- rnd (randomFrom [1..maxSize-2]) + l2 <- rnd (randomFrom [1..(maxSize - l1 - 1)]) + e1 <- randomChildFrom p1 l1 + ml <- gets (_size . _info . (IM.! e1) . _eClass) + e2 <- randomChildFrom p2 l2 + case nt of + Uni Id () -> canonical e1 + Uni f () -> add myCost (Uni f e1) >>= canonical + Bin op () () -> do b <- rnd toss + if b + then add myCost (Bin op e1 e2) >>= canonical + else add myCost (Bin op e2 e1) >>= canonical + + getParetoEcsUpTo n = concat <$> (forM [1..maxSize] $ \i -> getBestEcsOfSize i n) + + getBestEcsOfSize i n = do + ecs <- getTopECLassWithSize i n + Prelude.mapM canonical (Prelude.take n ecs) + + getBestEcs p n = do + ecs <- getTopECLassThat n p + --fits <- Prelude.mapM getFitness ecs + --let sorted = sort $ Prelude.zip (Prelude.map (fmap negate) fits) ecs + Prelude.mapM canonical (Prelude.take n ecs) + + randomChildFrom ec maxL = do + p <- rnd toss -- whether to go deeper or return this level + l <- gets (_size . _info . (IM.! ec) . _eClass ) + + if p || l >= maxL + then do enodes <- gets (_eNodes . (IM.! ec) . _eClass) + enode <- gets (_best . _info . (IM.! ec) . _eClass) -- we should return the best otherwise we may build larger + case enode of + Uni _ eci -> randomChildFrom eci maxL + Bin _ ecl ecr -> do coin <- rnd toss + if coin + then randomChildFrom ecl maxL + else randomChildFrom ecr maxL + _ -> pure ec + else pure ec + + nonTerms = [ Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () () + , Bin PowerAbs () (), Uni Recip (), Uni LogAbs (), Uni Exp (), Uni Sin (), Uni SqrtAbs ()] + rndTerm = Random.randomFrom terms + rndNonTerm = Random.randomFrom $ (Uni Id ()) : nonTerms + rndNonTerm2 = Random.randomFrom nonTerms + + insertTerms = + forM terms $ \t -> do fromTree myCost t >>= canonical + + insertRndExpr :: Int -> RndEGraph EClassId + insertRndExpr maxSize = + do grow <- rnd toss + t <- rnd $ Random.randomTree 2 8 maxSize rndTerm rndNonTerm2 grow + fromTree myCost t >>= canonical + + insertBestExpr :: RndEGraph EClassId + insertBestExpr = do --let t = "t0" / (recip ("t1" - "x0") + powabs "t2" "x0") + let t = ((("t0" + (powabs "t0" "x0")) / "t0") * "x0") + ecId <- fromTree myCost t >>= canonical + (f, p) <- fitnessFunRep x y x_val y_val t + insertFitness ecId f p + io . putStrLn $ "Best fit global: " <> show f + pure ecId + where powabs l r = Fix (Bin PowerAbs l r) + + getBestEclassThat p = + do ecIds <- getTopECLassThat 1 p -- isValidFitness + --bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds + --ecIds' <- getEClassesThat (fitnessIs (== Just bestFit)) + Prelude.mapM canonical $ Prelude.take 1 ecIds + + getBestExprWithSize n = + do ec <- getTopECLassWithSize n 1 + if (not (null ec)) + then do + bestFit <- getFitness $ head ec + bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass) + (:[]) . (,bestP) . (,bestFit) <$> getBest (head ec) + else pure [] + + getBestExprThat p = + do ec <- getBestEclassThat p + if (not (null ec)) + then do + bestFit <- getFitness $ head ec + (:[]) . (,bestFit) <$> getBest (head ec) + else pure [] + + printAll = do + ecs <- gets (IM.keys . _eClass) + forM_ ecs $ \ec -> + do t <- getBest ec + f <- gets (_fitness . _info . (IM.! ec) . _eClass) + io . putStrLn $ showExpr t <> " " <> show f + + paretoFront = go 1 (-1.0/0.0) + where + go n f + | n > maxSize = pure () + | otherwise = do + ecList <- getBestExprWithSize n + if (not (null ecList)) + then do let ((best, mf), mtheta) = head ecList + best' = relabelParams best + x_tot <- MA.computeAs MA.S <$> (MA.concatOuterM $ Prelude.map MA.toLoadArray [x, x_val]) + y_tot <- MA.computeAs MA.S <$> (MA.concatOuterM $ Prelude.map MA.toLoadArray [y, y_val]) + + --(fit_tr, theta) <- fitnessFunRep x_tot y_tot x_tot y_tot best' + let fit = fromJust mf + fit_tr = fit + theta = fromJust mtheta + fit_te = mse x_te y_te best' theta + str_th = intercalate ";" $ Prelude.map show $ MA.toList theta + + when (fit > f) do + io . putStrLn $ showExpr best <> "," <> str_th <> "," <> show (negate fit) <> "," <> show (negate fit_tr) <> "," <> show fit_te + go (n+1) (max fit f) + else go (n+1) f + + evaluateUnevaluated = do + ec <- gets (IntSet.toList . _unevaluated . _eDB) + forM_ ec $ \c -> do + t <- getBest c + (f, p) <- fitnessFun x y x_val y_val t + insertFitness c f p + + evaluateRndUnevaluated = do + ec <- gets (IntSet.toList . _unevaluated . _eDB) + c <- rnd . randomFrom $ ec + t <- getBest c + (f, p) <- fitnessFun x y x_val y_val t + insertFitness c f p + pure c + +while p arg prog = do b <- p + when b do arg' <- prog arg + while p arg' prog + + {- +egraphGP :: SRMatrix -> PVector -> [Fix SRTree] -> Int -> RndEGraph (Fix SRTree, Double) +egraphGP x y terms nEvals = do + replicateM_ 200 insertRndExpr + getBestExpr + runEqSat myCost rewrites 50 + evaluateUnevaluated + paretoFront + getBestExpr + where + paretoFront = do + forM_ [1..10] $ \i -> + do (best, fit) <- getBestExprThat (evaluated &&& isSizeOf (==i)) + io . putStrLn $ showExpr best <> " " <> show fit + + evaluateUnevaluated = do + ec <- getEClassesThat unevaluated + forM_ ec $ \c -> do + t <- getBest c + f <- fitnessFun x y t + updateFitness f c + + + + + + + + + + + + + + + + +------------------- GARBAGE CODE ------------------- + go i = do n <- getEClassesThat isValidFitness + unless (length n >= nEvals) + do gpStep + when (i `mod` 1000 == 0) (getBestExpr >>= (io . print . snd)) + when (i `mod` 1000000 == 0) $ do + n <- gets (IM.size . _eClass) + --io $ putStrLn ("before: " <> show n) + applyMergeOnlyDftl myCost + n1 <- gets (IM.size . _eClass) + when (n1 < n) $ io $ print (n,n1) + --io $ putStrLn ("after: " <> show n1) + go (i+1) + + rndTerm = Random.randomFrom terms + rndNonTerm = Random.randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Bin Div () () + , Bin PowerAbs () (), Uni Recip ()] + + getBestExpr :: RndEGraph (Fix SRTree, Double) + getBestExpr = do ecIds <- getEClassesThat evaluated -- isValidFitness + nc <- gets (IM.size . _eClass) + io . putStrLn $ "Evaluated expressions: " <> show (length ecIds) <> " / " <> show nc + bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds + ecIds' <- getEClassesThat (fitnessIs (== Just bestFit)) + (,bestFit) <$> getBest (head ecIds') + + getBestExprThat p = + do ecIds <- getEClassesThat p -- isValidFitness + nc <- gets (IM.size . _eClass) + bestFit <- foldM (\acc -> getFitness >=> (pure . max acc . fromJust)) ((-1.0)/0.0) ecIds + ecIds' <- getEClassesThat (fitnessIs (== Just bestFit)) + (,bestFit) <$> getBest (head ecIds') + + insertRndExpr :: RndEGraph () + insertRndExpr = do grow <- rnd toss + t <- rnd $ Random.randomTree 2 6 10 rndTerm rndNonTerm grow + f <- fitnessFun x y t + ecId <- fromTree myCost t >>= canonical + -- io $ print ('i', showExpr t, f) + updateFitness f ecId + + evalRndSubTree :: RndEGraph () + evalRndSubTree = do ecIds <- getEClassesThat unevaluated + unless (null ecIds) do + rndId <- rnd $ randomFrom ecIds + rndId' <- canonical rndId + t <- getBest rndId' + f <- fitnessFun x y t + -- io $ print ('e', showExpr t, f) + updateFitness f rndId' + + tournament :: Int -> [EClassId] -> RndEGraph EClassId + tournament n ecIds = do + (c0:cs) <- replicateM n (rnd (randomFrom ecIds)) + f0 <- gets (_fitness . _info . (IM.! c0) . _eClass) + snd <$> foldM (\(facc, acc) c -> gets (_fitness . _info . (IM.! c) . _eClass) + >>= \f -> if f > facc + then pure (f, c) + else pure (facc, acc)) (f0, c0) cs + + insertRndParent :: RndEGraph () + insertRndParent = do nt <- rnd rndNonTerm + meId <- case nt of + Uni f _ -> do ecIds <- getEClassesThat (isSizeOf (<10) &&& isValidFitness &&& funDoesNotExistWith nt) + if null ecIds + then pure Nothing + else do rndId <- tournament 5 ecIds + sz <- getSize rndId + if sz < 10 + then do node <- canonize (Uni f rndId) + Just <$> add myCost node + else pure Nothing + Bin op _ _ -> do ecIds <- getEClassesThat (isSizeOf (<9) &&& isValidFitness) + if null ecIds + then pure Nothing + else do rndIdLeft <- rnd $ randomFrom ecIds + sz1 <- getSize rndIdLeft + ecIds' <- getEClassesThat (isSizeOf (< (10 - sz1)) &&& isValidFitness &&& opDoesNotExistWith nt rndIdLeft) + if null ecIds' + then pure Nothing + else do rndIdRight <- rnd $ randomFrom ecIds' + rndIdRight <- tournament 5 ecIds' + sz2 <- getSize rndIdRight + if sz1 + sz2 < 10 + then Just <$> (canonize (Bin op rndIdLeft rndIdRight) + >>= add myCost) + else pure Nothing + when (isJust meId) do + let eId = fromJust meId + eId' <- canonical eId + curFit <- gets (_fitness . _info . (IM.! eId') . _eClass) + when (isNothing curFit) do + t <- getBest eId' + f <- fitnessFun x y t + updateFitness f eId' + -- io $ print ('p', showExpr t, f) + + gpStep :: RndEGraph () + gpStep = do choice <- rnd $ randomFrom [2,2,3,3,3] + if | choice == 1 -> insertRndExpr + | choice == 2 -> insertRndParent + | otherwise -> evalRndSubTree + rebuild myCost + -} +data Args = Args + { dataset :: String, + gens :: Int, + _alg :: Alg, + _maxSize :: Int, + _split :: Int + } + deriving (Show) + +-- parser of command line arguments +opt :: Parser Args +opt = Args + <$> strOption + ( long "dataset" + <> short 'd' + <> metavar "INPUT-FILE" + <> help "CSV dataset." ) + <*> option auto + ( long "generations" + <> short 'g' + <> metavar "GENS" + <> showDefault + <> value 100 + <> help "Number of generations." ) + <*> option auto + ( long "algorithm" + <> short 'a' + <> metavar "ALG" + <> help "Algorithm." ) + <*> option auto + ( long "maxSize" + <> short 's' + <> help "max-size." ) + <*> option auto + ( long "split" + <> short 'k' + <> help "k-split ratio training-test") + +chunksOf :: Int -> [e] -> [[e]] +chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls)) + where + splitter :: [e] -> ([e] -> a -> a) -> a -> a + splitter [] _ n = n + splitter l c n = l `c` splitter (Prelude.drop i l) c n + build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a] + build g = g (:) [] + +splitData :: SRMatrix -> PVector -> Int -> State StdGen (SRMatrix, SRMatrix, PVector, PVector) +splitData x y k = do if k == 1 + then pure (x, x, y, y) + else do + ixs' <- (state . shuffle) [0 .. sz-1] + let ixs = chunksOf k ixs' -- $ sortOn (\ix -> y MA.! ix) [0 .. sz-1] + --ixs <- forM sortedIxs $ \is -> state (shuffle is) + let xl = MA.toLists x :: [MA.ListItem MA.Ix2 Double] + x_tr = MA.fromLists' comp_x [xl !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i] + x_te = MA.fromLists' comp_x [xl !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i] + y_tr = MA.fromList comp_y [y MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i] + y_te = MA.fromList comp_y [y MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i] + + {- + ixs <- state (shuffle [0 .. sz-1]) + let ixs_tr = sort $ Prelude.take qty_tr ixs + ixs_te = sort $ Prelude.drop qty_tr ixs + + x_tr = MA.fromLists' comp_x [xl !! ix | ix <- ixs_tr] + x_te = MA.fromLists' comp_x [xl !! ix | ix <- ixs_te] + y_tr = MA.fromList comp_y [y MA.! ix | ix <- ixs_tr] + y_te = MA.fromList comp_y [y MA.! ix | ix <- ixs_te] + -} + pure (x_tr, x_te, y_tr, y_te) + where + (MA.Sz sz) = MA.size y + --qty_tr = round (thr * fromIntegral sz) + --qty_te = sz - qty_tr + comp_x = MA.getComp x + comp_y = MA.getComp y + +main :: IO () +main = do + --args <- pure (Args "nikuradse_2.csv" 100) -- execParser opts + args <- execParser opts + g <- getStdGen + ((x, y, _, _), _, _) <- loadDataset (dataset args) True + let ((x', x_te, y', y_te),g') = runState (splitData x y $ _split args) g + ((x_tr, x_val, y_tr, y_val),g'') = runState (splitData x' y' 2) g' + let (Sz2 _ nFeats) = MA.size x + terms = [var ix | ix <- [0 .. nFeats-1]] <> [param 0] -- [param ix | ix <- [0 .. 5]] + alg = evalStateT (egraphSearch (_alg args) x_tr y_tr x_val y_val x_te y_te terms (gens args) (_maxSize args)) emptyGraph + --(bestExpr, fit) <- evalStateT alg g + --printExpr bestExpr + --print fit + evalStateT alg g'' + + where + opts = Opt.info (opt <**> helper) + ( fullDesc <> progDesc "Very simple example of GP using SRTree." + <> header "tinyGP - a very simple example of GP using SRTRee." ) diff --git a/apps/egraphGP/Random.hs b/apps/egraphGP/Random.hs new file mode 100644 --- /dev/null +++ b/apps/egraphGP/Random.hs @@ -0,0 +1,54 @@ +module Random where + +import System.Random +import Control.Monad.State.Strict +import Control.Monad +import Data.SRTree +import Data.SRTree.Eval +import Data.Massiv.Array as MA + +type Rng a = StateT StdGen IO a + +toss :: Rng Bool +toss = state random +{-# INLINE toss #-} + +tossBiased :: Double -> Rng Bool +tossBiased p = do r <- state random + pure (r < p) + +randomVal :: Rng Double +randomVal = state random + +randomRange :: (Ord val, Random val) => (val, val) -> Rng val +randomRange rng = state (randomR rng) +{-# INLINE randomRange #-} + +randomFrom :: [a] -> Rng a +randomFrom funs = do n <- randomRange (0, length funs - 1) + pure $ funs !! n +{-# INLINE randomFrom #-} + +randomVec :: Int -> Rng PVector +randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-3.0, 3.0)) + +randomTree :: Int -> Int -> Int -> Rng (Fix SRTree) -> Rng (SRTree ()) -> Bool -> Rng (Fix SRTree) +randomTree minDepth maxDepth maxSize genTerm genNonTerm grow + | noSpaceLeft = genTerm + | needNonTerm = genRecursion + | otherwise = do r <- toss + if r + then genTerm + else genRecursion + where + noSpaceLeft = maxDepth <= 1 || maxSize <= 2 + needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2 + + genRecursion = do + node <- genNonTerm + case node of + Uni f _ -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow + Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 2) genTerm genNonTerm grow + r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow + pure . Fix $ Bin op l r +{-# INLINE randomTree #-} diff --git a/apps/eqsatrepr/Main.hs b/apps/eqsatrepr/Main.hs new file mode 100644 --- /dev/null +++ b/apps/eqsatrepr/Main.hs @@ -0,0 +1,302 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main where + +import Data.SRTree +import Algorithm.EqSat.Egraph +import Data.SRTree.Print +import qualified Data.Map as Map +import qualified Data.IntMap as IM +import Control.Monad.State.Strict +import System.Random +import Data.SRTree.Recursion ( cata ) +import Control.Monad +import Control.Monad.Reader +import qualified Data.SRTree.Random as RT +import Data.List ( nub ) +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Build +import Algorithm.EqSat.Queries +import Algorithm.EqSat + +isConstPt :: Pattern -> Map.Map ClassOrVar ClassOrVar -> EGraph -> Bool +isConstPt (VarPat c) subst eg = + let cid = getInt $ subst Map.! (Right $ fromEnum c) + in case (_consts . _info) (_eClass eg IM.! cid) of + ConstVal x -> True + _ -> False +isConstPt _ _ _ = False + +notZero (VarPat c) subst eg = + let cid = getInt $ subst Map.! (Right $ fromEnum c) + in case (_consts . _info) (_eClass eg IM.! cid) of + ConstVal x -> x /= 0 + _ -> True +notZero _ _ _ = True + +rewriteBasic = + [ + "x" * "x" :=> "x" ** 2 + , "x" * "y" :=> "y" * "x" + , "x" + "y" :=> "y" + "x" + , ("x" ** "y") * "x" :=> "x" ** ("y" + 1) :| isConstPt "y" + -- , ("x" * "y") / "x" :=> "y" + , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") + , ("x" + "y") + "z" :=> "x" + ("y" + "z") + , ("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" + , "x" - ("y" - "z") :=> ("x" - "y") + "z" + , ("x" * "y") / "z" :=> ("x" / "z") * "y" + , (("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y") :| isConstPt "w") :| isConstPt "z" + , ((("x" * "y") + ("z" * "w")) :=> "x" * ("y" + ("z" / "x") * "w") :| isConstPt "x") :| isConstPt "z" + , ((("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x") :| isConstPt "z" + , ((("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x") :| isConstPt "z" + ] + +rewritesFun = + [ + log (sqrt "x") :=> 0.5 * log "x" + , log (exp "x") :=> "x" + , exp (log "x") :=> "x" + , "x" ** (1/2) :=> sqrt "x" + , log ("x" * "y") :=> log "x" + log "y" + , log ("x" / "y") :=> log "x" - log "y" + , log ("x" ** "y") :=> "y" * log "x" + , sqrt ("y" * "x") :=> sqrt "y" * sqrt "x" + , sqrt ("y" / "x") :=> sqrt "y" / sqrt "x" + , abs ("x" * "y") :=> abs "x" * abs "y" + , sqrt ("z" * ("x" - "y")) :=> sqrt (negate "z") * sqrt ("y" - "x") + , sqrt ("z" * ("x" + "y")) :=> sqrt "z" * sqrt ("x" + "y") + ] + +-- Rules that reduces redundant parameters +constReduction = + [ + 0 + "x" :=> "x" + , "x" - 0 :=> "x" + , 1 * "x" :=> "x" + , 0 * "x" :=> 0 + , 0 / "x" :=> 0 :| notZero "x" + , "x" - "x" :=> 0 + , "x" / "x" :=> 1 :| notZero "x" + , "x" ** 1 :=> "x" + , 0 ** "x" :=> 0 + , 1 ** "x" :=> 1 + , "x" * (1 / "x") :=> 1 + , 0 - "x" :=> negate "x" + , "x" + negate "y" :=> "x" - "y" + , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x" + ] + + +x0 = var 0 +x1 = var 1 +x2 = var 2 +x3 = var 3 +x4 = var 4 +x5 = var 5 +x6 = var 6 +x7 = var 7 +x8 = var 8 + +trees :: [Fix SRTree] +trees = [ (4.059e-3 + (0.988153 * (((1.923901 * x1) * ((-1.228652 * x0) * (-0.278891 * x2))) * ((((((-0.35119 * x5) + (-0.354523 * x3)) - (-0.369148 * x6)) + ((0.342012 * x4) + (2.054e-2 * x2))) - ((0.349297 * x7) - (0.336081 * x8))))))) + , ((14.316036 * (((((0.975231 * x4)) * (1.259663 * x3)) * (0.314221 * x0)) * (0.178249 * x2)))) + , (1.2 * (3.4 * x1 * 4.2 * x0) / ((3.2 * x2) * ((1.1 * x3) * (3.5 * x4)))) + , ((1.002563 * (((0.428416 * x1) * (2.554566 * x0)) / (((2.53743 * x2) * (2.327917 * x3)) * (2.320736 * x3))))) + , (2.82238 + (3.092415 * (sin(log(abs(0.0))) * ((-0.162842 * x2) - (0.116404 * x1))))) + , log(0.0) * ((1.2 * x2) - (0.116404 * x1)) + , ((x0 - x0) * x0) + , (1 + 1) - 1 + , (x0 + x0) - x0 + , (x0/x0 + 1) - 1 + , (x0 * x0) / x0 + , sin(log(0.0)) + , -1 * exp(log(abs(-1.3 * (x1 - 1.2 * x2)))) + , -1 * exp(log(abs((1.3 * x1 + 1.56 * x2)))) + , -1 * exp(log(abs((-1.3 * x1 + 1.56 * x2)))) + , -1 * exp(log(abs(((0.256 * x3) + (-0.2561 * x2))))) + , log(abs(-1.199026) * abs((x2 + (1.191617 * x3)))) + , log(abs((1.199026 * x2) + (-1.191617 * x3))) + , 1 * x0 + , x0 * 1 + , x0 + x1 + , x1 + x0 + , x0 + sin(x1) + , x0 * (1 + x1) + , (1 + x1) * x0 + , x0 + x0 * x1 + , (x0 + x1) + 2 + , x0 + (x1 + 2) + , x0 + (2 + x1) + , log(abs(0)) + x0 + , abs(((1.3 * x1) + (-1.56 * x2))) * (-1.0) + ] + + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin op l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +rewrites = rewriteBasic <> constReduction <> rewritesFun + +testEqSat :: Fix SRTree -> IO () +testEqSat t = do + let e = eqSat t rewrites myCost 30 `evalState` emptyGraph + putStr $ (showExpr t) <> " == " <> (showExpr e) <> "\n" + +testEqSats :: IO () +testEqSats = mapM_ testEqSat trees + + + +initialPop :: HyperParams -> Rng [Fix SRTree] +initialPop hyperparams = do + let depths = [3 .. _maxDepth hyperparams] + pop <- forM depths $ \md -> + do let m = _popSize hyperparams `div` (_maxDepth hyperparams - 3 + 1) + g = take m $ cycle [True, False] + mapM (randomTree hyperparams{ _maxDepth = md}) g + pure (concat pop) +{-# INLINE initialPop #-} + +data Method = Grow | Full + +type Rng a = StateT StdGen IO a +type GenUni = Fix SRTree -> Fix SRTree +type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree + +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 #-} + +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 +{-# INLINE randomTree #-} + +data HyperParams = + HP { _minDepth :: Int + , _maxDepth :: Int + , _maxSize :: Int + , _popSize :: Int + , _tournSize :: Int + , _pc :: Double + , _pm :: Double + , _term :: [Fix SRTree] + , _nonterm :: [Either GenUni GenBin] + } + + +countSubTrees = do ecs <- gets (IM.keys . _eClass) + subs <- mapM (\ec -> getAllExpressionsFrom ec >>= pure . length) ecs + pure $ sum subs +countRootTrees rs = do subs <- mapM (\ec -> getAllExpressionsFrom ec >>= pure . length) rs + pure $ sum subs + +terms = [var 0, var 1, var 2, param 0, param 1, param 2, param 3] +nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> abs l ** r), Left (1/)] + +calcRedundancy :: Int -> IO () +calcRedundancy nPop = do + let hp = HP 2 4 10 nPop 2 1.0 0.25 terms nonterms + p = RT.P [0, 1, 2, 3, 4, 5] (0, 3) (1, 3) [Log] + g <- getStdGen + pop <- (`evalStateT` g) <$> replicateM nPop $ runReaderT (RT.randomTree 10) p + let nSubsSingle = sum $ map (\p -> (fromTrees myCost [p] >> countSubTrees) `evalState` emptyGraph) pop + myEqPop = do rs <- fromTrees myCost pop + let rsN = nub rs + cnt <- countSubTrees + pure (cnt, rsN) + (nSubs, rsN) = myEqPop `evalState` emptyGraph + putStr "Ratio of subtrees: " + putStrLn $ show nSubsSingle <> "/" <> show nSubs <> " = " <> show (fromIntegral nSubsSingle / fromIntegral nSubs) + let nSubsR = sum $ map (\p -> (fromTree myCost p >>= \r -> countRootTrees [r]) `evalState` emptyGraph) pop + nSubsSingleR = (fromTrees myCost pop >>= countRootTrees) `evalState` emptyGraph + putStr "Ratio of rooted trees: " + putStrLn $ show nSubsSingleR <> "/" <> show nSubsR <> " = " <> show (fromIntegral nSubsSingleR / fromIntegral nSubsR) + +main :: IO () +main = do + let t1 = var 0 + 12.0 + t2 = 3.2 * var 0 + t3 = 3.2 * var 0 / (var 0 + 12.0) + t4 = var 0 + sin (var 0) + t5 = 1.5 + exp 5.2 + egraphRun :: EGraphST IO () + egraphRun = do v <- fromTrees myCost [t3,t1,t2,t4] + roots <- findRootClasses + ecId <- gets ((Map.! (Var 0)) . _eNodeToEClass) + calculateHeights + h <- gets (map _height . IM.elems . _eClass) + v <- gets (map (_consts . _info) . IM.elems . _eClass) + c <- gets (map (_cost . _info) . IM.elems . _eClass) + parents <- gets (_parents . (IM.! ecId) . _eClass) + exprs <- mapM getExpressionFrom roots + exprs' <- gets (IM.keys . _eClass) >>= mapM getExpressionFrom + + lift $ do putStr "Parents of x0: " + print parents + putStrLn "\nexpressions from root: " + mapM_ (putStrLn . showExpr) exprs + putStrLn "\nexpressions from each e-class: " + mapM_ (putStrLn . showExpr) exprs' + putStrLn "heights: " + mapM_ print h -- (print . _height) (IM.elems $ _eClass eg') + putStrLn "values: " + mapM_ print v -- (print . _consts . _info) (IM.elems $ _eClass eg') + putStrLn "costs: " + mapM_ print c -- (print . _cost . _info) (IM.elems $ _eClass eg') + nPop = 10000 + hp = HP 3 7 100 nPop 2 1.0 0.25 terms nonterms + p = RT.P [0] (-3, 3) (-3, 3) [] + egraphRun `evalStateT` emptyGraph + g <- getStdGen + pop <- evalStateT (initialPop hp) g + mapM_ (\nP -> putStr "pop " >> print nP >> calcRedundancy nP >> putStrLn "") [100, 200, 500, 1000, 5000, 10000, 20000, 100000] diff --git a/apps/ieeexplore/Main.hs b/apps/ieeexplore/Main.hs new file mode 100644 --- /dev/null +++ b/apps/ieeexplore/Main.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE BlockArguments #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Simplify +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.Opt +import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.)) +import Control.Monad (foldM, forM_, forM, when, unless, filterM, (>=>), replicateM, replicateM_) +import Control.Monad.State +import qualified Data.IntMap as IM +import Data.Massiv.Array as MA hiding (forM_, forM) +import Data.Maybe (fromJust, isNothing, isJust) +import Data.SRTree +import Data.SRTree.Datasets +import Data.SRTree.Eval +import Data.SRTree.Random (randomTree) +import Data.SRTree.Print +import Options.Applicative as Opt hiding (Const) +import System.Random +import qualified Data.Set as Set +import Data.List ( sort ) + +import Debug.Trace +import Algorithm.EqSat (runEqSat) + + +type RangedTree a = a + +data QueryDB = QDB { _mseDB :: RangedTree Double + , _maeDB :: RangedTree Double + , _r2DB :: RangedTree Double + , _mdlDB :: RangedTree Double + , _bicDB :: RangedTree Double + , _aicDB :: RangedTree Double + , _lenDB :: RangedTree Int + , _ptDB :: RangedTree Int -- DB + } + +data ExtraInfo = ExtraInfo { _thetaMap :: IM.IntMap PVector + , _qdb :: QueryDB + } + +type InfoEGraph a = EGraphST (StateT ExtraInfo IO) a + +--io = lift . lift +--{-# INLINE io #-} +--ext = lift +--{-# INLINE ext #-} + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin op l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + + + +data Args = Args + { dataset :: String, + gens :: Int, + _maxSize :: Int + } + deriving (Show) + +-- parser of command line arguments +opt :: Parser Args +opt = Args + <$> strOption + ( long "dataset" + <> short 'd' + <> metavar "INPUT-FILE" + <> help "CSV dataset." ) + <*> option auto + ( long "generations" + <> short 'g' + <> metavar "GENS" + <> showDefault + <> value 100 + <> help "Number of generations." ) + <*> option auto + ( long "maxSize" + <> short 's' + <> help "max-size." ) + +main :: IO () +main = do + --args <- pure (Args "nikuradse_2.csv" 100) -- execParser opts + args <- execParser opts + ((x, y, _, _), _, _) <- loadDataset (dataset args) True + print "opa" + + where + opts = Opt.info (opt <**> helper) + ( fullDesc <> progDesc "Very simple example of GP using SRTree." + <> header "tinyGP - a very simple example of GP using SRTRee." ) diff --git a/apps/srsimplify/Main.hs b/apps/srsimplify/Main.hs new file mode 100644 --- /dev/null +++ b/apps/srsimplify/Main.hs @@ -0,0 +1,103 @@ +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 new file mode 100644 --- /dev/null +++ b/apps/srtools/Args.hs @@ -0,0 +1,178 @@ +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 + , msErr :: Maybe Double + , restart :: Bool + , rseed :: Int + , toScreen :: Bool + , useProfile :: Bool + , 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:vars where start and end \ + \ corresponds to the range of rows that should be used for fitting,\ + \ target is the column index (or name) of the target variable and cols\ + \ is a comma separated list of column indeces or names of the variables\ + \ in the same order as used by the symbolic model." ) + <*> 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." + ) + <*> option s2Reader + ( long "sErr" + <> metavar "Serr" + <> showDefault + <> value Nothing + <> help "Estimated standard error of the data.\ + \ If not passed, it uses the model MSE.") + <*> switch + ( long "restart" + <> help "If set, it samples the initial values of\ + \ 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." ) + <*> 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 : map toLower cs diff --git a/apps/srtools/IO.hs b/apps/srtools/IO.hs new file mode 100644 --- /dev/null +++ b/apps/srtools/IO.hs @@ -0,0 +1,181 @@ +{-# 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.Opt ( estimateSErr ) +import Algorithm.SRTree.Likelihoods ( Distribution (..) ) +import Algorithm.SRTree.ConfidenceIntervals ( printCI, BasicStats(_stdErr, _corr), CI ) +import qualified Data.SRTree.Print as P +import Data.SRTree.Eval ( compMode ) + +import Args ( Args(outfile, alpha,msErr,dist,niter) ) +import Report +import Data.SRTree.Recursion ( cata ) + +import Debug.Trace ( trace, traceShow ) + +-- Header of CSV file +csvHeader :: String +csvHeader = intercalate "," (basicFields <> optFields <> modelFields) +{-# inline csvHeader #-} + +-- 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 + mSErr' = case dist args of + Gaussian -> estimateSErr Gaussian (msErr args) (_xTr dset) (_yTr dset) (A.fromList compMode theta0) tree (niter args) + _ -> Nothing + args' = args{ msErr = mSErr' } + basic = getBasicStats args' seed dset tree theta0 ix + 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 tree + sseOpt = getSSE dset (_expr basic) + info = getInfo args' dset (_expr basic) treeVal + cis = getCI args' dset basic (alpha args') + +-- 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) + +-- 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)) + ] + 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) + +-- 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 new file mode 100644 --- /dev/null +++ b/apps/srtools/Main.hs @@ -0,0 +1,31 @@ +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 screne + 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 new file mode 100644 --- /dev/null +++ b/apps/srtools/Report.hs @@ -0,0 +1,271 @@ +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 ) +import Data.Random.Normal ( normals ) + +import Data.SRTree ( SRTree, Fix (..), floatConstsToParam, paramsToConst, countNodes ) +import Data.SRTree.Eval +import Algorithm.SRTree.AD ( reverseModeUnique, forwardModeUniqueJac ) +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.ModelSelection ( aic, bic, evidence, logFunctional, logParameters, mdl, mdlFreq, mdlLatt ) +import Algorithm.SRTree.ConfidenceIntervals +import Algorithm.SRTree.Opt (minimizeNLLWithFixedParam, minimizeNLL, minimizeNLLNonUnique) +import 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 + } + +-- basic fields name +basicFields :: [String] +basicFields = [ "Index" + , "Filename" + , "Expression" + , "Number_of_nodes" + , "Number_of_parameters" + , "Parameters" + ] + +-- basic information about the tree +data BasicInfo = Basic { _index :: Int + , _fname :: String + , _expr :: Fix SRTree + , _nNodes :: Int + , _nParams :: Int + , _params :: [Double] + } + +-- 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), varnames, tgname) <- loadDataset (dataset args) (hasHeader args) + let (A.Sz m) = A.size yVal + let (mXVal, mYVal) = if m == 0 + then (Nothing, Nothing) + else (Just xVal, Just yVal) + (mXTe, mYTe) <- if null (test args) + then pure (Nothing, Nothing) + else do ((xTe, yTe, _, _), _, _) <- loadDataset (test args) (hasHeader args) + pure (Just xTe, Just yTe) + pure (DS xTr yTr mXVal mYVal mXTe mYTe, varnames, tgname) + +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 + where + -- (tree', theta0) = floatConstsToParam tree + thetas = if restart args + then A.fromList compMode $ take nParams (normals seed) + else A.fromList compMode theta0 + t = fst $ minimizeNLL (dist args) (msErr args) (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' msErr' xTr yTr thetaOpt' tOpt + , _bicVal = bicVal + , _aic = aic dist' msErr' xTr yTr thetaOpt' tOpt + , _aicVal = aicVal + , _evidence = evidence dist' msErr' xTr yTr thetaOpt' tOpt + , _evidenceVal = evidenceVal + , _mdl = mdl dist' msErr' xTr yTr thetaOpt' tOpt + , _mdlFreq = mdlFreq dist' msErr' xTr yTr thetaOpt' tOpt + , _mdlLatt = mdlLatt dist' msErr' xTr yTr thetaOpt' tOpt + , _mdlVal = mdlVal + , _mdlFreqVal = mdlFreqVal + , _mdlLattVal = mdlLattVal + , _nllTr = nllTr + , _nllVal = nllVal + , _nllTe = nllTe + , _cc = logFunctional tOpt + , _cp = logParameters dist' msErr' xTr yTr thetaOpt' tOpt + , _fisher = A.toList $ fisherNLL dist' (msErr args) xTr yTr tOpt thetaOpt' + } + 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) = floatConstsToParam tree + thetaOpt' = A.fromList compMode thetaOpt + + (tOptVal, thetaOptVal) = floatConstsToParam treeVal + thetaOptVal' = A.fromList compMode thetaOptVal + + dist' = dist args + msErr' = msErr args + nllTr = nll dist' msErr' (_xTr dset) (_yTr dset) tOpt (A.fromList compMode thetaOpt) + bicVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> bic dist' msErr' xVal yVal thetaOptVal' tOptVal + aicVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> aic dist' msErr' xVal yVal thetaOptVal' tOptVal + evidenceVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> evidence dist' msErr' xVal yVal thetaOptVal' tOptVal + nllVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> nll dist' msErr' xVal yVal tOptVal (A.fromList compMode thetaOptVal) + mdlVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> mdl dist' msErr' xVal yVal thetaOptVal' tOptVal + mdlFreqVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> mdlFreq dist' msErr' xVal yVal thetaOptVal' tOptVal + mdlLattVal = case (_xVal dset, _yVal dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + _ -> mdlLatt dist' msErr' xVal yVal thetaOptVal' tOptVal + nllTe = case (_xTe dset, _yTe dset) of + (Nothing, _) -> 0.0 + (_, Nothing) -> 0.0 + (Just xTe, Just yTe) -> nll dist' msErr' xTe yTe tOpt (A.fromList compMode thetaOpt) + +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 + msErr' = msErr args + stats' = getStatsFromModel dist' msErr' xTr yTr tree (A.fromList compMode theta) + profiles = getAllProfiles (ptype args) dist' msErr' xTr yTr tree (A.fromList compMode theta) (_stdErr stats') estCIs alpha' + method = if useProfile args + then Profile stats' profiles + else Laplace stats' + predFun = A.computeAs A.S . predict dist' tree (A.fromList compMode theta) + + prof estPi th t = + let (thOpt, _) = minimizeNLLNonUnique dist' (Just 1) 100 xTr yTr t th + ssr = sse xTr yTr t thOpt + est = sqrt $ ssr / fromIntegral (n - _nParams basic) + stdErr = _stdErr stats' A.! 0 + fun = case ptype args of + Bates -> getProfile dist' (Just est) xTr yTr t thOpt stdErr tau_max 0 + ODE -> getProfileODE dist' (Just est) xTr yTr t thOpt stdErr estPi tau_max 0 + Constrained -> getProfileCnstr dist' (Just est) xTr yTr t thOpt stdErr tau_max' 0 + 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 new file mode 100644 --- /dev/null +++ b/apps/tinygp/GP.hs @@ -0,0 +1,222 @@ +{-# 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 +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 ) + +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 -> 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 +{-# INLINE randomTree #-} + +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 +{-# INLINE tournament #-} + +randomIndividual :: HyperParams -> FitFun -> Bool -> Rng Individual +randomIndividual hyperparams fitFun grow = do + t <- randomTree hyperparams grow + let p = countParams t + theta' <- replicateM p (randomRange (-1,1)) + let ind = fitFun $ Individual t 0.0 (M.fromList compMode theta' :: PVector) + pure ind + --if isInfinite (_fit ind) + -- then randomIndividual hyperparams fitFun grow + -- else pure ind +{-# INLINE randomIndividual #-} + +initialPop :: HyperParams -> FitFun -> Rng (V.Vector Individual) +initialPop hyperparams fitFun = do + 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) +{-# INLINE initialPop #-} + +fitness :: SRMatrix -> PVector -> Individual -> Individual +fitness x y ind = + let + tree = relabelParams $ _tree ind + thetaOrig = _params ind + (theta, fit) = minimizeNLL Gaussian (Just 1) 10 x y tree thetaOrig + --theta = _params ind + fit' = negate $ mse x y tree thetaOrig -- nll Gaussian (Just 1) x y (relabelParams $ _tree ind) (_params ind) + -- (fit, g) = gradNLL Gaussian Nothing x y (_tree ind) (_params ind) + in if M.isNull (_params ind) + then ind{_fit=fit'} + else ind{_fit = negate (mse x y tree theta), _params = theta} + --in ind{_fit = fit, _params = theta} +{-# INLINE fitness #-} + +isAbs (Fix (Uni Abs _)) = True +isAbs _ = False +{-# INLINE isAbs #-} + +isInv (Fix (Bin Div (Fix (Const 1.0)) _)) = True +isInv _ = False +{-# INLINE isInv #-} + +mutate :: HyperParams -> Individual -> Rng Individual +mutate hp ind = do + let sz = countNodes' (_tree ind) + (t, b) <- go sz (_pm hp) (_tree ind) + if b + then pure $ Individual t 0.0 M.empty + else pure ind + where + go s p t + | isAbs t = do let [x] = getChildren t + (t', b) <- go s p x + pure (Fix $ replaceChildren [t'] $ unfix t, b) + | otherwise = do + v <- state random + if v < p + then do let sz2 = countNodes' t + maxSz = _maxSize hp - s + sz2 - 2 + (, True) <$> randomTree hp{_maxSize = maxSz} True + else case arity t of + 0 -> pure (t, False) + 1 -> do let [x] = getChildren t + (t', b) <- go s p x + pure (Fix $ replaceChildren [t'] $ unfix t, b) + 2 -> do let [l,r] = getChildren t + (l', b) <- if isInv t + then pure (l, False) + else go s p l + if b + then pure (Fix $ replaceChildren [l', r] $ unfix t, b) + else do (r', b') <- go s p r + pure (Fix $ replaceChildren [l, r'] $ unfix t, b') + +crossover :: HyperParams -> Individual -> Individual -> Rng Individual +crossover hp ind1 ind2 = pure ind1 + +evolve :: HyperParams -> FitFun -> V.Vector Individual -> Rng Individual +evolve hp fitFun pop = do + parent1 <- tournament hp pop + parent2 <- tournament hp pop + child <- crossover hp parent1 parent2 + child' <- mutate hp child + let p = countParams (_tree child') + theta' <- M.fromList compMode <$> replicateM p (randomRange (-1,1)) + pure $ fitFun child'{_params = theta'} +{-# INLINE evolve #-} + +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 (M.toList $ _params ind) +{-# INLINE report #-} + +evolution :: Int -> HyperParams -> FitFun -> Rng (V.Vector Individual) +evolution gen hp fitFun = do + pop <- initialPop hp fitFun + liftIO $ report (-1) pop + go gen pop + where + go 0 pop = pure pop + go n pop = do + let best = V.maximumOn _fit $ V.filter (not.isNaN._fit) pop + pop' <- V.cons best <$> V.replicateM (_popSize hp - 1) (evolve hp fitFun pop) + liftIO $ report (gen-n) pop' + go (n-1) pop' diff --git a/apps/tinygp/Initialization.hs b/apps/tinygp/Initialization.hs new file mode 100644 --- /dev/null +++ b/apps/tinygp/Initialization.hs @@ -0,0 +1,50 @@ +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 new file mode 100644 --- /dev/null +++ b/apps/tinygp/Main.hs @@ -0,0 +1,72 @@ +module Main (main) where + +import GP ( HyperParams(HP), fitness, evolution ) +import Data.SRTree ( param, var ) +import System.Random ( getStdGen ) +import Control.Monad.State ( evalStateT ) +import Data.SRTree.Datasets ( loadDataset ) +import Options.Applicative +import Data.Massiv.Array + +-- Data type to store command line arguments +data Args = Args + { dataset :: String, + popSize :: Int, + gens :: Int, + pc :: Double, + pm :: Double + } + deriving (Show) + +-- parser of command line arguments +opt :: Parser Args +opt = Args + <$> strOption + ( long "dataset" + <> short 'd' + <> metavar "INPUT-FILE" + <> help "CSV dataset." ) + <*> option auto + ( long "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 "probCx" + <> metavar "PC" + <> showDefault + <> value 0.9 + <> help "Crossover probability." ) + <*> option auto + ( long "probMut" + <> metavar "PM" + <> showDefault + <> value 0.3 + <> help "Mutation probability." ) + +nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> abs l ** r), Left (1/)] +--nonterms = [Right (+), Right (-), Right (*)] + +main :: IO () +main = do + args <- execParser opts + g <- getStdGen + ((x, y, _, _), _, _) <- loadDataset (dataset args) True + let hp = HP 2 4 25 (popSize args) 2 (pc args) (pm args) terms nonterms + (Sz2 _ nFeats) = size x + terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]] + pop <- evalStateT (evolution (gens args) hp (fitness x y)) g + putStrLn "Fin" + 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/src/Algorithm/EqSat.hs b/src/Algorithm/EqSat.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE TupleSections #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Equality Saturation for SRTree +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- + +module Algorithm.EqSat where + +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Build +import Control.Lens (element, makeLenses, over, (&), (+~), (-~), (.~), (^.)) +import Control.Monad.State +import Data.Function (on) +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap +import Data.List (intercalate, minimumBy) +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 + +-- | 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 expr rules costFun maxIt = + do root <- fromTree costFun expr + (end, it) <- runEqSat costFun rules maxIt + best <- getBest root + --info <- gets ((IntMap.! root) . _eClass) + --info2 <- gets ((IntMap.! 9) . _eClass) + --traceShow (info, info2) $ + if not end -- if had an early stop + then do eqSat best rules costFun it -- reapplies eqsat on the best so far + else pure best + +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 costFun eid = + do classes <- gets _eClass + let costs = fillUpCosts classes Map.empty + eid' <- canonical eid + pure $ snd $ costs Map.! eid' + where + nodeCost :: CostMap -> ENode -> Maybe (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 + + 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 + + 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) + +-- | 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 + 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 + + go it sch = do eNodes <- gets _eNodeToEClass + eClasses <- gets _eClass + --createDB + --db <- gets (_patDB . _eDB) -- createDB -- creates the DB + + -- step 1: match the rules + let matchSch = matchWithScheduler it + matchAll = zipWithM matchSch [0..] + (rules, sch') = runState (matchAll rules') 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 + + -- recalculate heights + --calculateHeights + eNodes' <- gets _eNodeToEClass + eClasses' <- gets _eClass + + -- if nothing changed, return + if it == 1 || (eNodes' == eNodes && eClasses' == eClasses) + then pure (True, it) + else if IntMap.size eClasses' > 500 -- maximum allowed number of e-classes. TODO: customize + then pure (False, it) + else go (it-1) sch' + +-- | apply a single step of merge-only equality saturation +applySingleMergeOnlyEqSat :: Monad m => CostFun -> [Rule] -> EGraphST m () +applySingleMergeOnlyEqSat costFun rules = + do db <- gets (_patDB . _eDB) -- createDB + let matchSch = matchWithScheduler 10 + matchAll = zipWithM matchSch [0..] + (rules, sch') = runState (matchAll rules') IntMap.empty + matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rules + mapM_ (uncurry (applyMergeOnlyMatch costFun)) $ concat matches + 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 + +-- | matches the rules given a scheduler +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 + then pure [] + else do -- let matches = match db (source rule) + modify (IntMap.insert ruleNumber (it+5)) + pure [rule] -- $ map (rule,) matches diff --git a/src/Algorithm/EqSat/Build.hs b/src/Algorithm/EqSat/Build.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Build.hs @@ -0,0 +1,460 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE BangPatterns #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Build +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Functions related to building and maintaining e-graphs +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- + +module Algorithm.EqSat.Build where + +import System.Random (Random (randomR), StdGen) +import Control.Lens ( over ) +import Control.Monad ( forM_, when, foldM, forM ) +import Data.Maybe ( fromMaybe, catMaybes ) +import Data.SRTree +import Algorithm.EqSat.Egraph +--import Algorithm.EqSat.Info +import Algorithm.EqSat.DB +import qualified Data.IntMap.Strict as IntMap +import Data.Map.Strict ( Map ) +import qualified Data.Map.Strict as Map +import qualified Data.HashSet as Set +import Control.Monad.State.Strict +import Data.SRTree.Recursion (cataM) +import Algorithm.EqSat.Info +import qualified Data.IntSet as IntSet +import Data.Maybe +import Data.Sequence (Seq(..), (><)) + +import Debug.Trace (trace, traceShow) + +-- | 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'' + let enode' = case constEnode of + ConstVal x -> Const x + ParamIx x -> Param x + _ -> enode'' + maybeEid <- gets ((Map.!? enode') . _eNodeToEClass) -- check if canonical e-node exists + 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 + --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) + modify' $ over (eDB . unevaluated) (IntSet.insert curId) + pure curId + where + addParents :: Monad 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') + +-- | rebuilds the e-graph after inserting or merging +-- e-classes +rebuild :: Monad m => CostFun -> EGraphST m () +rebuild costFun = + do wl <- gets (_worklist . _eDB) + al <- gets (_analysis . _eDB) + modify' $ over (eDB . worklist) (const Set.empty) + . over (eDB . analysis) (const Set.empty) + forM_ wl (uncurry (repair costFun)) + forM_ al (uncurry (repairAnalysis costFun)) + +-- | 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 costFun ecId enode = + do modify' $ over eNodeToEClass (Map.delete enode) + enode' <- canonize enode + ecId' <- canonical ecId + doExist <- gets ((Map.!? enode') . _eNodeToEClass) + 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') + + +-- | repair the analysis of the e-class +-- considering the new added e-node +repairAnalysis :: Monad m => CostFun -> EClassId -> ENode -> EGraphST m () +repairAnalysis costFun ecId enode = + do ecId' <- canonical ecId + enode' <- canonize enode + eclass <- getEClass ecId' + info <- makeAnalysis costFun enode' + 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') + _ <- modifyEClass costFun ecId' + pure () + +-- | merge to equivalent e-classes +merge :: Monad m => CostFun -> EClassId -> EClassId -> EGraphST m EClassId +merge costFun c1 c2 = + do c1' <- canonical c1 + c2' <- canonical c2 + if c1' == c2' -- if they are already merged, return canonical + then pure c1' + 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 led ledC ledO sub subC subO = + do modify' $ over canonicalMap (IntMap.insert sub 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 + when (_info newC /= _info subC) + $ modify' $ over (eDB . analysis) (_parents subC <>) + updateDBs newC led ledC ledO sub subC subO + modifyEClass costFun led + pure led + + getLeaderSub c1 c1O c2 c2O = + do ec1 <- getEClass c1 + ec2 <- getEClass c2 + let n1 = length (_parents ec1) + n2 = length (_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 newC led ledC ledO sub subC subO = do + updateFitnessDB newC led ledC ledO sub subC subO + updateSizeDB newC led ledC ledO sub subC subO + + updateSizeDB :: Monad m => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m () + updateSizeDB newC led ledC ledO sub subC subO = do + let sz = (_size . _info) newC + szL = (_size . _info) ledC + szS = (_size . _info) subC + fun = IntMap.adjust (IntSet.insert led) sz . IntMap.adjust (IntSet.delete led . IntSet.delete ledO) szL . IntMap.adjust (IntSet.delete sub . IntSet.delete subO) szS + modify' $ over (eDB . sizeDB) fun + + 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) + where + fitNew = (_fitness . _info) newC + fitLed = (_fitness . _info) ledC + fitSub = (_fitness . _info) subC + szNew = (_size . _info) newC + szLed = (_size . _info) ledC + 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 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}) + 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}) + -- TODO: what happen to the orphans? + case maybeEid of + Nothing -> pure ecId + Just eid' -> trace "merge" $ merge costFun eid' ecId + -} + _ -> pure ecId + + where + isTerm (Var _) = True + isTerm (Const _) = True + isTerm (Param _) = True + isTerm _ = False + + toConst (Param ix) = ParamIx ix + toConst (Const 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) + +-- | `addToDB` adds an e-node and e-class id to the database +addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB () +addToDB enode eid = do + 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 + 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 + +-- | 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)) + +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) + 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) + +applyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m () +applyMatch costFun rule match' = + do let conds = getConditions rule + match <- canonizeMap match' + validHeight <- isValidHeight match + validConds <- mapM (`isValidConditions` match) conds + when (validHeight && and validConds) $ + do new_eclass <- reprPrat costFun (fst match) (target rule) + merge costFun (getInt (snd match)) new_eclass + pure () + +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 () + +-- | 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) + case maybeEid of + Nothing -> pure Nothing + Just eid -> Just <$> canonical eid +classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (Const x) +classOfENode costFun subst (Fixed target) = do newChildren <- mapM (classOfENode costFun subst) (getElems target) + case sequence newChildren of + Nothing -> pure Nothing + Just cs -> do let new_enode = replaceChildren cs target + cs' <- mapM canonical cs + areConsts <- mapM isConst cs' + if and areConsts + then do eid <- add costFun new_enode + rebuild costFun -- eid new_enode + pure (Just eid) + else gets ((Map.!? new_enode) . _eNodeToEClass) + +-- | adds the target of the rule into the e-graph +reprPrat :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m EClassId +reprPrat costFun subst (VarPat c) = canonical $ getInt $ subst Map.! Right (fromEnum c) +reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target) + add costFun (replaceChildren newChildren target) + +isValidHeight :: Monad m => (Map ClassOrVar ClassOrVar, 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 + +-- | 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) + +-- * 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) + +-- | Builds an e-graph from multiple independent trees +fromTrees :: Monad m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId] +fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) [] + + +-- | gets the best expression given the default cost function +getBest :: Monad m => EClassId -> EGraphST m (Fix SRTree) +getBest eid = do eid' <- canonical eid + best <- gets (_best . _info . (IntMap.! eid') . _eClass) + childs <- mapM getBest $ childrenOf best + pure . Fix $ replaceChildren childs best + +-- | returns one expression rooted at e-class `eId` +-- TODO: avoid loopings +getExpressionFrom :: Monad 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 + +-- | returns all expressions rooted at e-class `eId` +-- TODO: check for infinite list +getAllExpressionsFrom :: Monad 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] + 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] + ts <- go ns + pure (t:ts) + +-- | 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) + 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 + where + randomRange rng = state (randomR rng) + randomFrom xs = do n <- randomRange (0, length xs - 1) + pure $ xs !! n + +cleanMaps :: Monad 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 + +forceState :: Monad m => StateT s m () +forceState = get >>= \ !_ -> return () diff --git a/src/Algorithm/EqSat/DB.hs b/src/Algorithm/EqSat/DB.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/DB.hs @@ -0,0 +1,351 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE FlexibleContexts #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.EqSatDB +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Pattern matching and rule application functions +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- +module Algorithm.EqSat.DB where + +import Algorithm.EqSat.Egraph +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 Data.Map (Map) +import qualified Data.Map as Map +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 Data.String (IsString (..)) + +import Debug.Trace + +-- A Pattern is either a fixed-point of a tree or an +-- index to a pattern variable. The pattern variable matches anything. +data Pattern = Fixed (SRTree Pattern) | VarPat Char deriving Show -- Fixed structure of a pattern or a variable that matches anything + +-- 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, +-- and it will translate to `Fixed (Bin Add (VarPat 120) (VarPat 121)`. +instance IsString Pattern where + fromString [] = error "empty string in VarPat" + fromString [c] | n >= 65 && n <= 122 = VarPat c where n = fromEnum c + fromString s = error $ "invalid string in VarPat: " <> s + +-- A rule is either a directional rule where pat1 can be replaced by pat2, a bidirectional rule +-- where pat1 can be replaced or replace pat2, or a pattern with a conditional function +-- describing when to apply the rule +data Rule = Pattern :=> Pattern | Pattern :==: Pattern | Rule :| Condition + +infix 3 :=> +infix 3 :==: +infixl 2 :| + +instance Show Rule where + show (a :=> b) = show a <> " => " <> show b + show (a :==: b) = show a <> " == " <> show b + show (a :| b) = show a <> " | " + +-- 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 + +-- 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 + +unFixPat :: Pattern -> SRTree Pattern +unFixPat (Fixed p) = p +{-# INLINE unFixPat #-} + + +instance Num Pattern where + l + r = Fixed $ Bin Add l r + {-# INLINE (+) #-} + l - r = Fixed $ Bin Sub l r + {-# INLINE (-) #-} + l * r = Fixed $ Bin Mul l r + {-# INLINE (*) #-} + + abs = Fixed . Uni Abs + {-# INLINE abs #-} + + negate t = Fixed (Const (-1)) * t + {-# INLINE negate #-} + + signum t = case t of + Fixed (Const x) -> Fixed . Const $ signum x + _ -> Fixed (Const 0) + fromInteger x = Fixed $ Const (fromInteger x) + {-# INLINE fromInteger #-} + +instance Fractional Pattern where + l / r = Fixed $ Bin Div l r + {-# INLINE (/) #-} + + fromRational = Fixed . Const . fromRational + {-# INLINE fromRational #-} + +instance Floating Pattern where + pi = Fixed $ Const pi + {-# INLINE pi #-} + exp = Fixed . Uni Exp + {-# INLINE exp #-} + log = Fixed . Uni Log + {-# INLINE log #-} + sqrt = Fixed . Uni Sqrt + {-# INLINE sqrt #-} + sin = Fixed . Uni Sin + {-# INLINE sin #-} + cos = Fixed . Uni Cos + {-# INLINE cos #-} + tan = Fixed . Uni Tan + {-# INLINE tan #-} + asin = Fixed . Uni ASin + {-# INLINE asin #-} + acos = Fixed . Uni ACos + {-# INLINE acos #-} + atan = Fixed . Uni ATan + {-# INLINE atan #-} + sinh = Fixed . Uni Sinh + {-# INLINE sinh #-} + cosh = Fixed . Uni Cosh + {-# INLINE cosh #-} + tanh = Fixed . Uni Tanh + {-# INLINE tanh #-} + asinh = Fixed . Uni ASinh + {-# INLINE asinh #-} + acosh = Fixed . Uni ACosh + {-# INLINE acosh #-} + atanh = Fixed . Uni ATanh + {-# INLINE atanh #-} + + l ** r = Fixed $ Bin Power l r + {-# INLINE (**) #-} + + logBase l r = log l / log r + {-# INLINE logBase #-} + +target :: Rule -> Pattern +target (r :| _) = target r +target (_ :=> t) = t +target (_ :==: t) = t + +source :: Rule -> Pattern +source (r :| _) = source r +source (s :=> _) = s +source (s :==: _) = s + +getConditions :: Rule -> [Condition] +getConditions (r :| c) = c : getConditions r +getConditions _ = [] + + +cleanDB :: Monad m => EGraphST m () +cleanDB = modify' $ over (eDB. patDB) (const Map.empty) + +-- | 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] + +-- | Returns a Query (list of atoms) of a pattern +compileToQuery :: Pattern -> (Query, ClassOrVar) +compileToQuery pat = evalState (processPat pat) 256 -- returns (atoms, root) + where + -- 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) + +-- get the value from the Either Int Int +getInt :: ClassOrVar -> Int +getInt (Left a) = a +getInt (Right a) = a + +-- | returns the list of the children values +getElems :: SRTree a -> [a] +getElems (Bin _ l r) = [l,r] +getElems (Uni _ t) = [t] +getElems _ = [] + +-- | 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 + 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 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 + pure (concat maps) + + + -- [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 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 + + --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 _ [] root = pure [] +intersectAtoms var (a:atoms) root = do + a0 <- go a + Set.toList <$> (foldM (\acc atom -> Set.intersection acc <$> go atom) 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 + +-- | searches for the intersection of e-class ids that +-- matches each part of the query. +-- Returns Nothing if the intersection is empty. +-- +-- var is the current variable being investigated +-- xs is the map of ids being investigated and their corresponding e-class id +-- 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 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 + 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 + else Just $ IntMap.foldrWithKey (\k v acc -> + case intersectTries var (Map.insert i 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 + Nothing -> acc + Just s -> Set.union acc s + ) Set.empty (_trie trie) + +-- | updates all occurrence of var with the new id x +updateVar :: ClassOrVar -> ClassOrVar -> Query -> Query +updateVar var x = map replace + where + replace (Atom r t) = let children = [if c == var then x else c | c <- getElems t] + t' = replaceChildren children t + in Atom (if r == var then x else r) t' + +-- | checks whether two ClassOrVar are different +-- only check if it is a pattern variable, else returns true +isDiffFrom :: Int -> ClassOrVar -> Bool +isDiffFrom x y = case y of + Left _ -> False + Right z -> x /= z + +-- | checks if v is an element of an atom +elemOfAtom :: ClassOrVar -> Atom -> Bool +elemOfAtom v (Atom root tree) = + case root of + Left _ -> v `elem` getElems tree + Right x -> Right x == v || v `elem` getElems tree + +-- | sorts the variables in a query by the most frequently occurring +orderedVars :: Query -> [ClassOrVar] +orderedVars atoms = sortBy (comparing varCost) $ nub [a | atom <- atoms, a <- getIdsFrom atom, isRight a] + where + getIdsFrom (Atom r t) = r : getElems t + isRight (Right _) = True + isRight _ = False + + varCost :: ClassOrVar -> Int + varCost var = foldr (\a acc -> if elemOfAtom var a then acc - 100 + atomLen a else acc) 0 atoms + + atomLen (Atom _ t) = 1 + length (getElems t) diff --git a/src/Algorithm/EqSat/Egraph.hs b/src/Algorithm/EqSat/Egraph.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Egraph.hs @@ -0,0 +1,270 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE StrictData #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Egraph +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Equality Graph data structure +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- + +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.State.Strict +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.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 Data.SRTree +import Data.SRTree.Eval +import Data.Hashable + +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) +type EGraphST m a = StateT EGraph m a +type Cost = Int +type CostFun = SRTree Cost -> Cost + +instance Hashable ENode where + hashWithSalt n enode = hashWithSalt n (encodeEnode enode) + +type RangeTree a = Seq (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 #-} + +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 #-} + +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 + +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 + + +-- 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 + + +getSmallest :: Ord a => RangeTree a -> (a, EClassId) +getSmallest rt = case rt of + Empty -> error "empty finger" + x :<| t -> x +getGreatest :: Ord a => RangeTree a -> (a, EClassId) +getGreatest rt = case rt of + Empty -> error "empty finger" + t :|> x -> x + + +data EGraph = EGraph { _canonicalMap :: ClassIdMap EClassId -- maps an e-class id to its canonical form + , _eNodeToEClass :: Map ENode EClassId -- maps an e-node to its e-class id + , _eClass :: ClassIdMap EClass -- maps an e-class id to its e-class data + , _eDB :: EGraphDB + } deriving Show + +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 + , _patDB :: DB -- database of patterns + , _fitRangeDB :: RangeTree Double -- database of valid fitness + , _sizeDB :: IntMap IntSet -- database of model sizes + , _sizeFitDB :: IntMap (RangeTree Double) -- hacky! Size x Fitness DB + , _unevaluated :: IntSet -- set of not-evaluated e-classes + , _nextId :: Int -- next available id + } deriving Show + +data EClass = EClass { _eClassId :: Int -- e-class id (maybe we don't need that here) + , _eNodes :: HashSet ENodeEnc -- set of e-nodes inside this e-class + , _parents :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s + , _height :: Int -- height + , _info :: EClassData -- data + } deriving (Show, Eq) + +data Consts = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq) +data Property = Positive | Negative | NonZero | Real deriving (Show, Eq) -- TODO: incorporate properties + +data EClassData = EData { _cost :: Cost + , _best :: ENode + , _consts :: Consts + , _fitness :: Maybe Double -- NOTE: this cannot be NaN + , _theta :: Maybe PVector + , _size :: Int + -- , _properties :: Property + -- TODO: include evaluation of expression from this e-class + } deriving (Show) + +instance Eq EClassData where + EData c1 b1 cs1 ft1 _ s1 == EData c2 b2 cs2 ft2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && s1==s2 + +-- The database maps a symbol to an IntTrie +-- The IntTrie stores the possible paths from a certain e-class +-- that matches a pattern +type DB = Map (SRTree ()) IntTrie +-- The IntTrie is composed of the set of available keys (for convenience) +-- and an IntMap that maps one e-class id to the first child IntTrie, +-- the first child IntTrie will point to the next child and so on +data IntTrie = IntTrie { _keys :: HashSet EClassId, _trie :: IntMap IntTrie } -- deriving Show + +-- 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 <> "}" + +makeLenses ''EGraph +makeLenses ''EClass +makeLenses ''EClassData +makeLenses ''EGraphDB + +-- * E-Graph basic supporting functions + +-- | returns an empty e-graph +emptyGraph :: EGraph +emptyGraph = EGraph IntMap.empty Map.empty IntMap.empty emptyDB + +-- | returns an empty e-graph DB +emptyDB :: EGraphDB +emptyDB = EDB Set.empty Set.empty Map.empty FingerTree.empty IntMap.empty IntMap.empty IntSet.empty 0 + +-- | 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 +{-# 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 +{-# INLINE canonical #-} + +-- | canonize the e-node children +canonize :: Monad m => ENode -> EGraphST m ENode +canonize = mapM canonical -- applies canonical to the children +{-# INLINE canonize #-} + +-- | gets an e-class with id `c` +getEClass :: Monad m => EClassId -> EGraphST m EClass +getEClass c = gets ((IntMap.! c) . _eClass) +{-# INLINE getEClass #-} + +-- | Creates a singleton trie from an e-class id +trie :: EClassId -> IntMap IntTrie -> IntTrie +trie eid = IntTrie (Set.singleton eid) + +-- | Check whether an e-class is a constant value +isConst :: Monad m => EClassId -> EGraphST m Bool +isConst eid = do ec <- gets ((IntMap.! eid) . _eClass) + case (_consts . _info) ec of + ConstVal _ -> pure True + _ -> pure False +{-# INLINE isConst #-} diff --git a/src/Algorithm/EqSat/Info.hs b/src/Algorithm/EqSat/Info.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Info.hs @@ -0,0 +1,168 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Info +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Functions related to info/data calculation in Equality Graph data structure +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- + +module Algorithm.EqSat.Info where + +import Control.Lens ( over ) +import Control.Monad --(forM, forM_, when, foldM, void) +import Control.Monad.State +import Data.AEq (AEq ((~==))) +import Data.IntMap (IntMap) -- , delete, empty, insert, toList) +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.HashSet (HashSet) +import qualified Data.HashSet as Set +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 +-- TODO: instead of folding, just do not apply rules +-- list of values instead of single value +joinData :: EClassData -> EClassData -> EClassData +joinData (EData c1 b1 cn1 fit1 p1 sz1) (EData c2 b2 cn2 fit2 p2 sz2) = + EData (min c1 c2) b (combineConsts cn1 cn2) (minMaybe fit1 fit2) (bestParam p1 p2 fit1 fit2) (min sz1 sz2) + where + minMaybe Nothing x = x + minMaybe x Nothing = x + minMaybe x y = min x y + + bestParam Nothing x _ _ = x + bestParam x Nothing _ _ = x + bestParam x y (Just f1) (Just f2) = if f1 < f2 then x else y + + b = if c1 <= c2 then b1 else b2 + combineConsts (ConstVal x) (ConstVal y) + | abs (x-y) < 1e-7 = ConstVal $ (x+y)/2 + | isNaN x || isInfinite x = ConstVal y + | isNaN y || isInfinite y = ConstVal x + | isNaN x && isNaN y = ConstVal x + | x ~== y = ConstVal $ (x+y)/2 + | abs (x / y) < 1 + 1e-6 || abs (y / x) < 1 + 1e-6 = ConstVal $ min x y + | isInfinite x && isInfinite y = ConstVal x + | isInfinite x && isNaN y = ConstVal y + | isNaN x && isInfinite y = ConstVal x + | otherwise = error $ "Combining different values: " <> show x <> " " <> show y <> " " <> show (x/y) + combineConsts (ParamIx ix) (ParamIx iy) = ParamIx (min ix iy) + combineConsts NotConst x = x + combineConsts x NotConst = x + combineConsts x y = error (show x <> " " <> show y) + +-- | Calculate e-node data (constant values and cost) +makeAnalysis :: Monad m => CostFun -> ENode -> EGraphST m EClassData +makeAnalysis costFun enode = + do consts <- calculateConsts enode + 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) + +getChildrenMinHeight :: Monad 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 + +-- | update the heights of each e-class +-- won't work if there's no root +calculateHeights :: Monad m => EGraphST m () +calculateHeights = + do queue <- findRootClasses + classes <- gets (Prelude.map fst . IntMap.toList . _eClass) + 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 x eId' = + do eId <- canonical eId' + ec <- getEClass eId + let ec' = over height (const x) ec + modify' $ over eClass (IntMap.insert eId ec') + + 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] + + go [] _ _ = pure () + go qs tabu h = + do childrenOf <- (TrueSet.\\ tabu) . TrueSet.fromList . concat <$> forM qs getChildrenEC -- rerieve all unvisited children + let childrenL = TrueSet.toList childrenOf + forM_ childrenL (setMinHeight h) -- set the height of the children as the minimum between current and h + 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 + costs <- traverse (fmap (_cost . _info) . getEClass) cs + pure . f $ replaceChildren costs t + +-- | 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 + consts <- traverse (fmap (_consts . _info) . getEClass) cs + case combineConsts $ replaceChildren consts t of + ConstVal x | isNaN x -> pure (ConstVal x) + a -> pure a + +combineConsts :: SRTree Consts -> Consts +combineConsts (Const x) = ConstVal x +combineConsts (Param ix) = ParamIx ix +combineConsts (Var _) = NotConst +combineConsts (Uni f t) = case t of + ConstVal x -> ConstVal $ evalFun f x + _ -> t +combineConsts (Bin op l r) = evalOp' l r + where + evalOp' (ParamIx ix) (ParamIx iy) = ParamIx (min ix iy) + 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 + ec <- gets ((IntMap.! eId) . _eClass) + let oldFit = _fitness . _info $ ec + newInfo = (_info ec){_fitness = Just fit, _theta = Just 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) + else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit)) diff --git a/src/Algorithm/EqSat/Queries.hs b/src/Algorithm/EqSat/Queries.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Queries.hs @@ -0,0 +1,92 @@ +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE BangPatterns #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Queries +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Query functions for e-graphs +-- Heavily based on hegg (https://github.com/alt-romes/hegg by alt-romes) +-- +----------------------------------------------------------------------------- + +module Algorithm.EqSat.Queries where + +import Algorithm.EqSat.Egraph +import qualified Data.IntMap as IntMap +import qualified Data.Map as Map +import qualified Data.HashSet as Set +import Control.Monad.State ( gets, modify' ) +import Control.Monad ( filterM ) +import Control.Lens ( over ) +import Data.Maybe +import Data.Sequence ( Seq(..) ) + +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 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' + +updateFitness :: Monad m => Double -> EClassId -> EGraphST m () +updateFitness f ecId = do + ec <- gets ((IntMap.! ecId) . _eClass) + let info = _info ec + modify' $ over eClass (IntMap.insert ecId 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) + 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 => Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopECLassThat n p = do + gets (_fitRangeDB . _eDB) + >>= go n [] + where + go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go 0 bests rt = pure bests + go m bests rt = case rt of + Empty -> pure bests + t :|> y -> do let x = snd y + ecId <- canonical x + ec <- gets ((IntMap.! ecId) . _eClass) + if (isInfinite . fromJust . _fitness . _info $ ec) + then pure bests + else if p ec + then go (m-1) (x:bests) t + else go m bests t +getTopECLassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId] +getTopECLassWithSize sz n = do + gets ((IntMap.!? sz) . _sizeFitDB . _eDB) + >>= go n [] + where + go :: Monad m => Int -> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId] + go _ bests Nothing = pure [] + go 0 bests (Just rt) = pure bests + go m bests (Just 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 pure bests + else go (m-1) (x:bests) (Just t) diff --git a/src/Algorithm/EqSat/Simplify.hs b/src/Algorithm/EqSat/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Simplify.hs @@ -0,0 +1,214 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE LambdaCase #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.EqSat.Simplify +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : +-- +-- Module containing the algebraic rules and simplification function. +-- +----------------------------------------------------------------------------- +module Algorithm.EqSat.Simplify ( Rule(..), simplifyEqSatDefault, applyMergeOnlyDftl, rewrites, rewriteBasic, rewritesFun ) where + +import Algorithm.EqSat (eqSat, applySingleMergeOnlyEqSat) +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.DB + ( ClassOrVar, + Pattern (Fixed, VarPat), + Rule (..), + getInt, + ) +import Control.Monad.State.Strict (evalState) +import Data.IntMap (IntMap) +import qualified Data.IntMap as IM +import Data.Map (Map) +import qualified Data.Map as Map +import Data.SRTree + +type ConstrFun = Pattern -> Map ClassOrVar ClassOrVar -> EGraph -> Bool + +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 + +-- TODO: aux functions to avoid repeated pattern in constraint creation +-- +-- check if a matched pattern contains constant +isConstPt :: ConstrFun +isConstPt = constrainOnVal $ + \case + ConstVal _ -> True + _ -> False + +-- check if the matched pattern is a positive constant +isConstPos :: ConstrFun +isConstPos = constrainOnVal $ + \case + ConstVal x -> x > 0 + _ -> False + +isNotParam :: ConstrFun +isNotParam = constrainOnVal $ + \case + ParamIx _ -> False + _ -> True + +-- check if the matched pattern is nonzero +isNotZero :: ConstrFun +isNotZero = constrainOnVal $ + \case + ConstVal x -> abs x < 1e-9 + _ -> True + +-- check if the matched pattern is even +isEven :: ConstrFun +isEven = constrainOnVal $ + \case + ConstVal x -> ceiling x == floor x && even (round x) + _ -> True + +-- check if the matched pattern is integer +isInteger :: ConstrFun +isInteger = constrainOnVal $ + \case + ConstVal x -> ceiling x == floor x + _ -> True + +-- check if the matched pattern is positive +isPositive :: ConstrFun +isPositive = constrainOnVal $ + \case + ConstVal x -> x > 0 + _ -> True + +-- check if the matched pattern is valid +isValid :: ConstrFun +isValid = constrainOnVal $ + \case + ConstVal x -> not (isNaN x || isInfinite x) + _ -> True + +-- basic algebraic rules +rewriteBasic :: [Rule] +rewriteBasic = + [ + "x" * "x" :=> "x" ** 2 + , "x" * "y" :=> "y" * "x" + , "x" + "y" :=> "y" + "x" + , ("x" ** "y") * "x" :=> "x" ** ("y" + 1) :| isConstPt "y" + , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x" + , ("x" + "y") + "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" + , (("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x" + , (("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x" :| isConstPt "z" + -- , "x" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL + -- , "x" + "y" * "z" :=> "y" * ("x" * "y" ** (-1) + "z") :| isNotZero "y" -- GABRIEL + ] + +-- rules for nonlinear functions +rewritesFun :: [Rule] +rewritesFun = + [ + log (sqrt "x") :=> 0.5 * log "x" :| isNotParam "x" + , log (exp "x") :==: exp (log "x") + , log (exp "x") :=> "x" + -- , exp (log "x") :=> "x" -- :| isPositive "x" ??? exp(log(x)), x, log(exp(0)) + , "x" ** (1/2) :==: sqrt "x" -- <==> + , "x" ** (1/3) :==: Fixed (Uni Cbrt "x") + , log ("x" * "y") :=> log "x" + log "y" :| isConstPos "x" :| isConstPos "y" + -- , log ("x" / "y") :=> log "x" - log "y" :| isConstPos "x" :| isConstPos "y" + , log ("x" ** "y") :=> "y" * log "x" + --, 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" + , 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 + , abs "x" ** "y" :=> "x" ** "y" :| isEven "y" + ] + +-- Rules that reduces redundant parameters +constReduction :: [Rule] +constReduction = + [ + 0 + "x" :=> "x" + -- , "x" - 0 :=> "x" + , 1 * "x" :=> "x" + , 0 * "x" :=> 0 :| isValid "x" -- :| isNotParam "x" + -- , 0 / "x" :=> 0 :| isNotZero "x" + --, "x" - "x" :=> 0 :| isNotParam "x" + --, "x" / "x" :=> 1 :| isNotZero "x" :| isNotParam "x" + , "x" ** 1 :=> "x" + , 0 ** "x" :=> 0 :| isPositive "x" + , 1 ** "x" :=> 1 + -- , "x" * (1 / "x") :=> 1 :| isNotParam "x" :| isNotZero "x" + , 0 - "x" :=> negate "x" + , "x" + negate "y" :==: "x" - "y" + -- , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x" + , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isPositive "x" + , "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isPositive "x" + , ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isPositive "x" + , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isPositive "x" :| isPositive "y" + + , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isInteger "y" :| isNotZero "x" + , "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x" + , ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x" + , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isInteger "z" :| isNotZero "x" :| isNotZero "y" + + ] + +-- | default cost function for simplification +-- TODO: +-- num_params: +-- length: +-- terminal < nonterminal: +-- symbol comparison (constants, parameters, variables x0, x10, x2) +-- op priorities (+, -, *, inv_div, pow, abs, exp, log, log10, sqrt) +-- univariates +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin op l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +-- all rewrite rules +rewrites :: [Rule] +rewrites = rewriteBasic <> constReduction <> rewritesFun + +-- | simplify using the default parameters +simplifyEqSatDefault :: Fix SRTree -> Fix SRTree +simplifyEqSatDefault t = eqSat t rewrites myCost 30 `evalState` emptyGraph + +-- | 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 costFun = applySingleMergeOnlyEqSat costFun rewrites diff --git a/src/Algorithm/Massiv/Utils.hs b/src/Algorithm/Massiv/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/Massiv/Utils.hs @@ -0,0 +1,278 @@ +{-# 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 new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD.hs @@ -0,0 +1,323 @@ +{-# language FlexibleInstances, DeriveFunctor #-} +{-# language ScopedTypeVariables #-} +{-# language RankNTypes #-} +{-# language ViewPatterns #-} +{-# language FlexibleContexts #-} +{-# language BangPatterns #-} +----------------------------------------------------------------------------- +-- | +-- 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 + ( forwardMode + , forwardModeUnique + , reverseModeUnique + , forwardModeUniqueJac + ) where + +import Control.Monad (forM_) +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) + +applyUni :: (Index ix, Source r e, Floating e, Floating b) => Function -> Either (Array r ix e) b -> Either (Array D ix e) b +applyUni f (Left t) = + Left $ M.map (evalFun f) t +applyUni f (Right t) = + Right $ evalFun f t +{-# INLINE applyUni #-} + +applyDer :: (Index ix, Source r e, Floating e, Floating b) => Function -> Either (Array r ix e) b -> Either (Array D ix e) b +applyDer f (Left t) = + Left $ M.map (derivative f) t +applyDer f (Right t) = + Right $ derivative f t +{-# INLINE applyDer #-} + +negate' :: (Index ix, Source r e, Num e, Num b) => Either (Array r ix e) b -> Either (Array D ix e) b +negate' (Left t) = Left $ M.map negate t +negate' (Right t) = Right $ negate t +{-# INLINE negate' #-} + +applyBin :: (Index ix, Floating b) => Op -> Either (Array D ix b) b -> Either (Array D ix b) b -> Either (Array D ix b) b +applyBin op (Left ly) (Left ry) = + Left $ case op of + Add -> ly !+! ry + Sub -> ly !-! ry + Mul -> ly !*! ry + Div -> ly !/! ry + Power -> ly .** ry + PowerAbs -> M.map abs (ly .** ry) + AQ -> ly !/! (M.map sqrt (M.map (+1) (ry !*! ry))) + +applyBin op (Left ly) (Right ry) = + Left $ unsafeLiftArray (\ x -> evalOp op x ry) ly +applyBin op (Right ly) (Left ry) = + Left $ unsafeLiftArray (\ x -> evalOp op ly x) ry +applyBin op (Right ly) (Right ry) = + Right $ evalOp op ly ry +{-# INLINE applyBin #-} + +-- | get the value of a certain index if it is an array (Left) +-- or returns the value itself if it is a scalar. +(!??) :: (Manifest r e, Index ix) => Either (Array r ix e) e -> ix -> e +(Left y) !?? ix = y ! ix +(Right y) !?? ix = y +{-# INLINE (!??) #-} + +-- | Calculates the results of the error vector multiplied by the Jacobian of an expression using forward mode +-- provided a vector of variable values `xss`, a vector of parameter values `theta` and +-- a function that changes a Double value to the type of the variable values. +-- uses unsafe operations to use mutable array instead of a tape +forwardMode :: Array S Ix2 Double -> Array S Ix1 Double -> SRVector -> Fix SRTree -> (Array D Ix1 Double, Array S Ix1 Double) +forwardMode xss theta err tree = let (yhat, jacob) = runST $ cataM lToR alg tree + in (fromEither yhat, computeAs S err >>= UMA.unsafeFreeze cmp + pure (Left (xss >= UMA.unsafeFreeze cmp + pure (Right c, tape) + + -- if it is a parameter, returns the value of the parameter and the jacobian with a one in the corresponding column + alg (Param ix) = do tape <- M.makeMArrayS (Sz2 m p) (\(i :. j) -> pure $ if j==ix then 1 else 0) + >>= UMA.unsafeFreeze cmp + pure (Right (theta ! ix), tape) + + -- 1. applies the derivative of f in the evaluated child + -- 2. replaces the value of the Jacobian at (i, j) with yi * J[i, j] + alg (Uni f (t, tape')) = do let y = computeAs S . fromEither $ applyDer f t + tape <- UMA.unsafeThaw tape' + forM_ [0 .. m-1] $ \i -> do + let yi = y ! i + forM_ [0 .. p-1] $ \j -> do + v <- UMA.unsafeRead tape (i :. j) + UMA.unsafeWrite tape (i :. j) (yi * v) + tapeF <- UMA.unsafeFreeze cmp tape + pure (applyUni f t, tapeF) + -- li, ri are the corresponding values of the evaluated left and right children + -- vl, vr are the corresponding value of the Jacobian at (i, j) + -- applies the corresponding derivative of each binary operator + alg (Bin op (l, tl') (r, tr')) = do + tl <- UMA.unsafeThaw tl' + tr <- UMA.unsafeThaw tr' + let l' = case l of + Left y -> Left $ computeAs S y + Right v -> Right v + r' = case r of + Left y -> Left $ computeAs S y + Right v -> Right v + forM_ [0 .. m-1] $ \i -> do + let li = l' !?? i + ri = r' !?? i + forM_ [0 .. p-1] $ \j -> do + vl <- UMA.unsafeRead tl (i :. j) + vr <- UMA.unsafeRead tr (i :. j) + UMA.unsafeWrite tl (i :. j) $ case op of + Add -> (vl+vr) + Sub -> (vl-vr) + Mul -> (vl * ri + vr * li) + Div -> ((vl * ri - vr * li) / ri^2) + Power -> (li ** (ri - 1) * (ri * vl + li * log li * vr)) + PowerAbs -> (abs li ** ri) * (vr * log (abs li) + ri * vl / li) + AQ -> ((1 + ri*ri) * vl - li * ri * vr) / (1 + ri*ri) ** 1.5 + tlF <- UMA.unsafeFreeze cmp tl + pure (applyBin op l r, tlF) + + + lToR (Var ix) = pure (Var ix) + lToR (Param ix) = pure (Param ix) + lToR (Const c) = pure (Const c) + lToR (Uni f mt) = Uni f <$> mt + lToR (Bin op ml mr) = Bin op <$> ml <*> mr + +-- | The function `forwardModeUnique` calculates the numerical gradient of the tree and evaluates the tree at the same time. It assumes that each parameter has a unique occurrence in the expression. This should be significantly faster than `forwardMode`. +forwardModeUnique :: SRMatrix -> PVector -> SRVector -> Fix SRTree -> (SRVector, Array S Ix1 Double) +forwardModeUnique xss theta err = second (toGrad . DL.toList) . cata alg + where + (Sz n) = M.size theta + one = replicateAs xss 1 + toGrad grad = M.fromList (getComp xss) [g !.! err | g <- grad] + + alg (Var ix) = (xss PVector + -> SRVector + -> (SRVector -> SRVector) + -> Fix SRTree + -> (Array D Ix1 Double, Array S Ix1 Double) +reverseModeUnique xss theta ys f t = unsafePerformIO $ + do jacob <- M.newMArray (Sz p) 0 + let !_ = accu reverse (combine jacob) t ((Right 1), fwdMode) + j <- freezeS jacob + pure (v, j) + where + fwdMode = cata forward t + v = fromEither $ getTop fwdMode + err = f v - ys + (Sz2 m _) = M.size xss + p = countParams t + fromEither (Left x) = x + fromEither (Right x) = M.replicate (getComp xss) (Sz1 m) x + + oneTpl x = Fix $ Single x + tuple x y = Fix $ T x y + branch x y z = Fix $ Branch x y z + + getTop (Fix (Single x)) = x + getTop (Fix (T x y)) = x + getTop (Fix (Branch x y z)) = x + + unCons (Fix (T x y)) = y + getBranches (Fix (Branch x y z)) = (y,z) + + -- forward just creates a new tree with the partial + -- evaluation of the nodes + forward (Var ix) = oneTpl (Left $ xss v) = + let g' = applyDer f (getTop v) + in Uni f (t, ( applyBin Mul dx g', v )) + reverse (Bin op l r) (dx, getBranches -> (vl, vr)) = + let (dxl, dxr) = diff op dx (getTop vl) (getTop vr) + in Bin op (l, (dxl, vl)) (r, (dxr, vr)) + + -- dx is the current derivative so far + -- fx is the evaluation of the left branch + -- gx is the evaluation of the right branch + -- + -- this should return a tuple, where the left element is + -- dx * d op(f(x), g(x)) / d f(x) and + -- the right branch dx * d op (f(x), g(x)) / d g(x) + diff Add dx fx gy = (dx, dx) + diff Sub dx fx gy = (dx, negate' dx) + diff Mul dx fx gy = (applyBin Mul dx gy, applyBin Mul dx fx) + diff Div dx fx gy = (applyBin Div dx gy, applyBin Mul dx (applyBin Div (negate' fx) (applyBin Mul gy gy))) + diff Power dx fx gy = let dxl = applyBin Mul dx (applyBin Power fx (applyBin Sub gy (Right 1))) + dv2 = applyBin Mul fx (applyUni Log fx) + in (applyBin Mul dxl gy, applyBin Mul dxl dv2) + diff PowerAbs dx fx gy = let dxl = applyBin Mul (applyBin Mul gy fx) (applyBin PowerAbs fx (applyBin Sub gy (Right 2))) + dxr = applyBin Mul (applyUni LogAbs fx) (applyBin PowerAbs fx gy) + in (applyBin Mul dxl dx, applyBin Mul dxr dx) + diff AQ dx fx gy = let dxl = applyUni Recip (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1))) + dxy = applyBin Div (applyBin Mul fx gy) (applyUni Cube (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1)))) + in (applyBin Mul dxl dx, applyBin Mul dxy dx) + + + -- once we reach a leaf with a parameter, we return a singleton + -- with that derivative upwards until the root + --combine :: (forall s . MArray (PrimState (ST s)) S Int Double) -> SRTree () -> (Either SRVector Double, a) -> () + combine j (Var ix) s = 0 + combine j (Const _) s = 0 + combine j (Param ix) s = unsafePerformIO $ do + case fst s of + Left v -> do v' <- dotM v err + UMA.unsafeWrite j ix v' + Right v -> UMA.unsafeWrite j ix $ M.foldrS (\x acc -> x*v + acc) 0 err + UMA.unsafeRead j ix + combine j (Uni f gs) s = gs + combine j (Bin op l r) s = l+r + + +-- | 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 + + alg (Var ix) = (xss Double + , _theta2tau :: Double -> Double + } + +-- 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) +printCI :: Int -> CI -> IO () +printCI n = putStrLn . showCI n + +-- | 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 + 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 + +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 + +-- | 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] +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 + + getResStdError row = sqrt $ (A.!.!) row $ A.fromList compMode $ map (row A.!.!) covs + resStdErr = map getResStdError jac + +predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat $ take 10 xss' + where + yhat = A.toList $ predFun xss + theta' = A.toStorableVector theta + + t = sqrt $ quantile (fDistribution k (fromIntegral $ n - k)) (1 - alpha) + (A.Sz k) = A.size theta + n = length yhat + + theta0 = calcTheta0 dist tree + xss' = A.toList $ A.outerSlices xss + + 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) + +-- inverse function of the distributions +inverseDist :: Floating p => Distribution -> p -> p +inverseDist Gaussian y = y +inverseDist Bernoulli y = log (y/(1-y)) +inverseDist Poisson y = log y + +-- 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 (Bin op l r) = Fix $ Bin op l r + +evalVar :: PVector -> 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 (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?" + 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 (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 + +-- calculate the profile likelihood of every parameter +getAllProfiles :: PType -> Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> PVector -> [CI] -> Double -> [ProfileT] +getAllProfiles ptype dist mSErr xss ys tree theta stdErr estCIs alpha = reverse (getAll 0 []) + 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) + + profFun ix = case ptype of + Bates -> getProfile dist mSErr xss ys tree theta (stdErr A.! ix) tau_max ix + ODE -> getProfileODE dist mSErr xss ys tree theta (stdErr A.! ix) (estCIs !! ix) tau_max ix + Constrained -> getProfileCnstr dist mSErr xss ys tree theta (stdErr A.! ix) tau_max' ix + + getAll ix acc | ix == k = acc + | otherwise = case profFun ix of + Left t -> getAllProfiles ptype dist mSErr xss ys tree t stdErr estCIs alpha + Right p -> getAll (ix + 1) (p : acc) + +-- calculates the profile likelihood of a single parameter +getProfile :: Distribution + -> Maybe Double + -> SRMatrix + -> PVector + -> Fix SRTree + -> PVector + -> Double + -> Double + -> Int + -> Either PVector ProfileT +getProfile dist mSErr xss ys tree theta stdErr_i tau_max ix + | stdErr_i == 0.0 = pure $ ProfileT (A.fromList compMode [-tau_max, tau_max]) (A.fromLists' compMode [theta', theta']) (theta A.! ix) (const (theta A.! ix)) (const tau_max) + | otherwise = + do negDelta <- go kmax (-stdErr_i / 8) 0 1 mempty + 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 + pure $ ProfileT taus thetas optTh tau2theta theta2tau + where + theta' = A.toList theta + p0 = ([0], [theta_opt]) + kmax = 300 + nll_opt = nll dist mSErr xss ys tree theta_opt + theta_opt = fst $ minimizeNLLNonUnique dist mSErr 100 xss ys tree theta + optTh = theta_opt A.! ix + minimizer = minimizeNLLWithFixedParam dist mSErr 100 xss ys tree 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 + | otherwise = go (k-1) delta (t + inv_slope) inv_slope' acc' + where + t_delta = (theta_opt A.! ix) + delta * (t + inv_slope) + theta_delta = updateS theta_opt [(ix, t_delta)] + theta_t = minimizer theta_delta + zv = A.computeAs A.S (snd $ gradNLL dist mSErr xss ys tree theta_t) A.! ix + zvs = snd $ gradNLL dist mSErr xss ys tree theta_t + inv_slope' = min 4.0 . max 0.0625 . abs $ (tau / (stdErr_i * zv)) + nll_cond = nll dist mSErr xss ys tree theta_t + acc' = if nll_cond == nll_opt || ( (not.null) taus && tau == head taus ) || isNaN tau + then acc + else (tau:taus, theta_t:thetas) + 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 Double + -> SRMatrix + -> PVector + -> Fix SRTree + -> PVector + -> Double -> Double + -> Int + -> Either PVector ProfileT +getProfileCnstr dist mSErr xss ys tree theta stdErr_i tau_max ix + | stdErr_i == 0.0 = pure $ ProfileT taus thetas theta_i (const theta_i) (const tau_max) + | otherwise = pure $ ProfileT taus thetas theta_i tau2theta (const tau_max) + where + 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 mSErr xss ys tree theta tau_max ix + leftPt = getPoint True + rightPt = getPoint False + tau2theta tau = if tau < 0 then leftPt else rightPt + +getEndPoint :: Distribution -> Maybe Double -> A.Array A.S Ix2 Double -> A.Array A.S A.Ix1 Double -> Fix SRTree -> A.Array A.S A.Ix1 Double -> Double -> Int -> Bool -> Double +getEndPoint dist mSErr xss ys tree theta tau_max ix isLeft = + case minimizeAugLag problem (A.toStorableVector theta_opt) of + Right sol -> solutionParams sol VS.! ix + Left e -> traceShow e $ theta_opt A.! ix + where + (A.Sz1 n) = A.size theta + + theta_opt = fst $ minimizeNLLNonUnique dist mSErr 100 xss ys tree theta + nll_opt = nll dist mSErr xss ys tree theta_opt + loss_crit = nll_opt + tau_max + + loss = subtract loss_crit . nll dist mSErr xss ys tree . A.fromStorableVector compMode + obj = (if isLeft then id else negate) . (VS.! ix) + + stop = ObjectiveRelativeTolerance 1e-4 :| [] + localAlg = NELDERMEAD obj [] Nothing + local = LocalProblem (fromIntegral n) stop localAlg + constraint = InequalityConstraint (Scalar loss) 1e-6 + + problem = AugLagProblem [] [] (AUGLAG_LOCAL local [constraint] []) +{-# INLINE getEndPoint #-} + +-- 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 Double + -> SRMatrix + -> PVector + -> Fix SRTree + -> PVector + -> Double + -> CI + -> Double + -> Int + -> Either PVector ProfileT +getProfileODE dist mSErr xss ys tree theta stdErr_i estCI tau_max ix + | stdErr_i == 0.0 = pure dflt + | otherwise = let (A.fromList compMode -> taus, A.fromLists' compMode . map A.toList -> thetas) = solLeft <> ([0], [theta_opt]) <> solRight + (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix + in pure $ ProfileT taus thetas optTh tau2theta theta2tau + where + dflt = ProfileT (A.fromList compMode [-tau_max, tau_max]) (A.fromLists' compMode [theta', theta']) (theta A.! ix) (const (theta A.! ix)) (const tau_max) + minimizer = fst . minimizeNLLNonUnique dist mSErr 100 xss ys tree + grader = snd . gradNLLNonUnique dist mSErr xss ys tree + theta_opt = minimizer theta + theta' = A.toList theta + nll_opt = nll dist mSErr xss ys tree theta_opt + optTh = theta_opt A.! ix + p' = p+1 + (A.Sz1 p) = A.size theta + sErr = fromMaybe 1 mSErr + getHess = hessianNLL dist mSErr xss ys tree + + odeFun gamma _ u = + let grad = grader u + w = hessianNLL dist mSErr 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 + 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) + 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 mSErr 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]) + +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))) + 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 +{-# INLINE rk #-} + +-- tau0, tau1 theta0, thetaX = tau1 theta0 / tau0 +getStatsFromModel :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> BasicStats +getStatsFromModel dist mSErr xss ys tree theta = MkStats cov corr stdErr + where + (A.Sz1 k) = A.size theta + (A.Sz1 n) = A.size ys + nParams = fromIntegral k + ssr = sse xss ys tree theta + ident = A.computeAs A.S $ identityMatrix nParams + + -- only for gaussian + sErr = sqrt $ ssr / fromIntegral (n - k) + + hess = hessianNLL dist mSErr 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 + + stdErr = A.makeArray compMode (A.Sz1 k) (\ix -> sqrt $ cov A.! (ix :. ix)) + stdErrSq = case outer stdErr stdErr of + Left _ -> error "stdErr size mismatch?" + Right v -> v + + corr = A.computeAs A.S $ A.zipWith (/) cov stdErrSq + +-- Create splines for profile-t +createSplines :: PVector -> SRMatrix -> 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)]) + | otherwise = (tau2theta, theta2tau) + where + (A.Sz n) = A.size 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 + +getCol :: Int -> SRMatrix -> PVector +getCol ix mtx = getCols mtx A.! ix +{-# inline getCol #-} + +sortOnFirst :: PVector -> PVector -> [(Double, Double)] +sortOnFirst xs ys = sortOn fst $ zip (A.toList xs) (A.toList ys) +{-# inline sortOnFirst #-} + +splinesSketches :: Double -> PVector -> PVector -> (Double -> Double) -> (Double -> Double) +splinesSketches tauScale (A.toList -> tau) (A.toList -> theta) theta2tau + | length tau < 2 = id + | otherwise = genSplineFun gpq + where + 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) + (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 + + 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) + + -- generate the points of the curve + go 100 = [] + go ix = (p, q) : go (ix+1) + where + ai = 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 diff --git a/src/Algorithm/SRTree/Likelihoods.hs b/src/Algorithm/SRTree/Likelihoods.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Likelihoods.hs @@ -0,0 +1,260 @@ +{-# LANGUAGE ViewPatterns #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.SRTree.Likelihoods +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to calculate different likelihood functions, their gradient, and Hessian matrices. +-- +----------------------------------------------------------------------------- +module Algorithm.SRTree.Likelihoods + ( Distribution (..) + , PVector + , SRMatrix + , sse + , mse + , rmse + , r2 + , nll + , predict + , gradNLL + , gradNLLNonUnique + , fisherNLL + , getSErr + , hessianNLL + ) + where + +import Algorithm.SRTree.AD ( forwardMode, reverseModeUnique ) -- ( reverseModeUnique ) +import Data.Massiv.Array hiding (all, map, read, replicate, tail, take, zip) +import qualified Data.Massiv.Array as M +import Data.Maybe (fromMaybe) +import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams) +import Data.SRTree.Derivative (deriveByParam) +import Data.SRTree.Eval (PVector, SRMatrix, SRVector, compMode, evalTree) + +-- | Supported distributions for negative log-likelihood +data Distribution = Gaussian | Bernoulli | Poisson + deriving (Show, Read, Enum, Bounded) + +-- | 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) + +-- | 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 + +-- | Root of the mean squared errors +rmse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double +rmse xss ys tree = sqrt . mse xss ys tree + +-- | 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 + +-- | logistic function +logistic :: Floating a => a -> a +logistic x = 1 / (1 + exp (-x)) +{-# inline logistic #-} + +-- | get the standard error from a Maybe Double +-- if it is Nothing, estimate from the ssr, otherwise use the current value +-- For distributions other than Gaussian, it defaults to a constant 1 +getSErr :: Num a => Distribution -> a -> Maybe a -> a +getSErr Gaussian est = fromMaybe est +getSErr _ _ = const 1 +{-# inline getSErr #-} + +-- negation of the sum of values in a vector +negSum :: PVector -> Double +negSum = negate . M.sum +{-# inline negSum #-} + +-- | Negative log-likelihood +nll :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> Double + +-- | Gaussian distribution +nll Gaussian msErr xss ys t theta = 0.5*(ssr/s2 + m*log (2*pi*s2)) + where + (Sz m') = M.size ys + (Sz p') = M.size theta + m = fromIntegral m' + p = fromIntegral p' + ssr = sse xss ys t theta + mse' = mse xss ys t theta + est = sqrt (m - p) -- $ ssr / (m - p) + sErr = getSErr Gaussian est msErr + s2 = sErr ^ 2 + +-- | Bernoulli distribution of f(x; theta) is, given phi = 1 / (1 + exp (-f(x; theta))), +-- y log phi + (1-y) log (1 - phi), assuming y \in {0,1} +nll Bernoulli _ xss ys tree theta + | notValid ys = error "For Bernoulli distribution the output must be either 0 or 1." + | otherwise = negate . M.sum $ delay ys * yhat - log (1 + exp yhat) + 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' :: Distribution -> Double -> SRVector -> SRVector -> Double +nll' Gaussian sErr yhat ys = 0.5*(ssr/s2 + m*log (2*pi*s2)) + where + (Sz m') = M.size ys + m = fromIntegral m' + ssr = M.sum $ (ys - yhat)^2 + s2 = sErr ^ 2 +nll' Bernoulli _ yhat ys = negate . M.sum $ ys * yhat - log (1 + exp yhat) +nll' Poisson _ yhat ys = negate . M.sum $ ys * yhat - ys * log ys - exp yhat +{-# INLINE nll' #-} + +-- | Prediction for different distributions +predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector +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 + +-- | Gradient of the negative log-likelihood +gradNLL :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector) +gradNLL Gaussian msErr xss ys tree theta = + (nll' Gaussian sErr yhat ys', delay grad ./ (sErr * sErr)) + where + (Sz m) = M.size ys + (Sz p) = M.size theta + ys' = delay ys + (yhat, grad) = reverseModeUnique xss theta ys' id tree + -- err = yhat - delay ys + ssr = sse xss ys tree theta + est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p) + sErr = getSErr Gaussian est msErr + +gradNLL Bernoulli _ xss (delay -> ys) tree theta + | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1." + | otherwise = (nll' Bernoulli 1.0 yhat ys, delay grad) + where + (yhat, grad) = reverseModeUnique xss theta ys logistic tree + grad' = M.map nanTo0 grad + --err = logistic yhat - ys + nanTo0 x = if isNaN x then 0 else x + +gradNLL Poisson _ xss (delay -> ys) tree theta + | M.any (<0) ys = error "For Poisson distribution the output must be non-negative." + -- | M.any isNaN grad = error $ "NaN gradient " <> show grad + | otherwise = (nll' Poisson 1.0 yhat ys, delay grad) + where + (yhat, grad) = reverseModeUnique xss theta ys exp tree + --err = exp yhat - ys + +-- | Gradient of the negative log-likelihood +gradNLLNonUnique :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector) +gradNLLNonUnique Gaussian msErr xss ys tree theta = + (nll' Gaussian sErr yhat ys', delay grad ./ (sErr * sErr)) + where + (Sz m) = M.size ys + (Sz p) = M.size theta + ys' = delay ys + (yhat, grad) = forwardMode xss theta err tree + err = predict Gaussian tree theta xss - ys' + ssr = sse xss ys tree theta + est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p) + sErr = getSErr Gaussian est msErr + +gradNLLNonUnique Bernoulli _ xss (delay -> ys) tree theta + | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1." + | otherwise = (nll' Bernoulli 1.0 yhat ys, delay grad) + where + (yhat, grad) = forwardMode xss theta err tree + grad' = M.map nanTo0 grad + err = predict Bernoulli tree theta xss - delay ys + nanTo0 x = if isNaN x then 0 else x + +gradNLLNonUnique Poisson _ xss (delay -> ys) tree theta + | M.any (<0) ys = error "For Poisson distribution the output must be non-negative." + -- | M.any isNaN grad = error $ "NaN gradient " <> show grad + | otherwise = (nll' Poisson 1.0 yhat ys, delay grad) + where + (yhat, grad) = forwardMode xss theta err tree + err = predict Poisson tree theta xss - delay ys + +-- | Fisher information of negative log-likelihood +fisherNLL :: Distribution -> Maybe Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector +fisherNLL dist msErr xss ys tree theta = makeArray cmp (Sz p) build + where + build ix = let dtdix = deriveByParam ix t' + d2tdix2 = deriveByParam ix dtdix + f' = eval dtdix + f'' = eval d2tdix2 + in (/sErr^2) . M.sum $ phi' * f'^2 - res * f'' + cmp = getComp xss + (Sz m) = M.size ys + (Sz p) = M.size theta + t' = fst $ floatConstsToParam tree + eval = evalTree xss theta + ssr = sse xss ys tree theta + sErr = getSErr dist est msErr + est = sqrt $ fromIntegral (m-p) -- $ ssr / fromIntegral (m - p) + yhat = eval t' + res = delay ys - phi + + (phi, phi') = case dist of + Gaussian -> (yhat, M.replicate compMode (Sz m) 1) + Bernoulli -> (logistic yhat, phi*(M.replicate compMode (Sz 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 Double -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRMatrix +hessianNLL dist msErr xss ys tree theta = makeArray cmp (Sz (p :. p)) build + 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 (/sErr^2) . M.sum $ phi' * fx * fy - res * fxy + cmp = getComp xss + (Sz m) = M.size ys + (Sz p) = M.size theta + t' = tree -- relabelParams tree -- $ floatConstsToParam tree + eval = evalTree xss theta + ssr = sse xss ys tree theta + sErr = getSErr dist est msErr + est = sqrt $ fromIntegral (m - p) -- $ ssr / fromIntegral (m - p) + yhat = eval t' + res = delay ys - phi + + (phi, phi') = case dist of + Gaussian -> (yhat, M.replicate cmp (Sz m) 1) + Bernoulli -> (logistic yhat, phi*(M.replicate cmp (Sz m) 1 - phi)) + Poisson -> (exp yhat, phi) + diff --git a/src/Algorithm/SRTree/ModelSelection.hs b/src/Algorithm/SRTree/ModelSelection.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/ModelSelection.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.SRTree.ModelSelection +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Helper functions for model selection criteria +-- +----------------------------------------------------------------------------- + +module Algorithm.SRTree.ModelSelection where + +import Algorithm.Massiv.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 +import Data.SRTree +import Data.SRTree.Eval (evalTree) +import Data.SRTree.Recursion (cata) +import qualified Data.Vector.Storable as VS + + +-- | Bayesian information criterion +bic :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +bic dist mSErr xss ys theta tree = (p + 1) * log n + 2 * nll dist mSErr xss ys tree theta + where + (A.Sz (fromIntegral -> p)) = A.size theta + (A.Sz (fromIntegral -> n)) = A.size ys +{-# INLINE bic #-} + +-- | Akaike information criterion +aic :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +aic dist mSErr xss ys theta tree = 2 * (p + 1) + 2 * nll dist mSErr xss ys tree theta + where + (A.Sz (fromIntegral -> p)) = A.size theta + (A.Sz (fromIntegral -> n)) = A.size ys +{-# INLINE aic #-} + +-- | Evidence +evidence :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +evidence dist mSErr xss ys theta tree = (1 - b) * nll dist mSErr xss ys tree theta - p / 2 * log b + where + (A.Sz (fromIntegral -> p)) = A.size theta + (A.Sz (fromIntegral -> n)) = A.size ys + b = 1 / sqrt n +{-# INLINE evidence #-} + +-- | MDL as described in +-- Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023). +mdl :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +mdl dist mSErr xss ys theta tree = nll' dist mSErr xss ys theta' tree + + logFunctional tree + + logParameters dist mSErr xss ys theta tree + where + fisher = fisherNLL dist mSErr xss ys tree theta + theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher + isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +{-# INLINE mdl #-} + +-- | MDL Lattice as described in +-- Bartlett, Deaglan, Harry Desmond, and Pedro Ferreira. "Priors for symbolic regression." Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 2023. +mdlLatt :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +mdlLatt dist mSErr xss ys theta tree = nll' dist mSErr xss ys theta' tree + + logFunctional tree + + logParametersLatt dist mSErr xss ys theta tree + where + fisher = fisherNLL dist mSErr xss ys tree theta + theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher + isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +{-# INLINE mdlLatt #-} + +-- | same as `mdl` but weighting the functional structure by frequency calculated using a wiki information of +-- physics and engineering functions +mdlFreq :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +mdlFreq dist mSErr xss ys theta tree = nll dist mSErr xss ys tree theta + + logFunctionalFreq tree + + logParameters dist mSErr xss ys theta tree +{-# INLINE mdlFreq #-} + +-- 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 + where + 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) + where + tree' = fst $ floatConstsToParam tree + consts = getIntConsts tree +{-# INLINE logFunctionalFreq #-} + +-- log of the parameters complexity +logParameters :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +logParameters dist mSErr xss ys theta tree = -(p / 2) * log 3 + 0.5 * logFisher + logTheta + where + -- p = fromIntegral $ VS.length theta + fisher = fisherNLL dist mSErr 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 Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +logParametersLatt dist mSErr xss ys theta tree = 0.5 * p * (1 - log 3) + 0.5 * log detFisher + where + fisher = fisherNLL dist mSErr xss ys tree theta + detFisher = det $ hessianNLL dist mSErr xss ys tree theta + + (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 Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double +nll' dist mSErr xss ys theta tree = nll dist mSErr 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 + where + + opToNat :: Op -> Double + opToNat Add = 2.500842464597881 + opToNat Sub = 2.500842464597881 + opToNat Mul = 1.720356134912558 + opToNat Div = 2.60436883851265 + opToNat Power = 2.527957363394847 + opToNat PowerAbs = 2.527957363394847 + opToNat AQ = 2.60436883851265 + + 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 Sinh = 8.038963823353235 + funToNat Cosh = 8.262107374667444 + funToNat Tanh = 7.85664226655928 + funToNat Tan = 8.262107374667444 + funToNat _ = 8.262107374667444 + --funToNat Factorial = 7.702491586732021 +{-# INLINE treeToNat #-} diff --git a/src/Algorithm/SRTree/NonlinearOpt.hs b/src/Algorithm/SRTree/NonlinearOpt.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/NonlinearOpt.hs @@ -0,0 +1,974 @@ +{-# 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) <- N.optimize opt x0 + if (N.isSuccess outret) + then return $ Solution outcost outx outret + 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 + } 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/Algorithm/SRTree/Opt.hs b/src/Algorithm/SRTree/Opt.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Opt.hs @@ -0,0 +1,109 @@ +----------------------------------------------------------------------------- +-- | +-- 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) +import Data.SRTree.Eval (evalTree, compMode) +import qualified Data.Vector.Storable as VS + +-- | minimizes the negative log-likelihood of the expression +minimizeNLL :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double) +minimizeNLL dist msErr niter xss ys tree t0 + | niter == 0 = (t0, f) + | n == 0 = (t0, f) + | otherwise = (fromStorableVector compMode t_opt, f) + where + tree' = relabelParams tree -- $ fst $ floatConstsToParam tree + t0' = toStorableVector t0 + (Sz n) = size t0 + (Sz m) = size ys + funAndGrad = second (toStorableVector . computeAs S) . gradNLL dist msErr xss ys tree' . fromStorableVector compMode + (f, _) = gradNLL dist msErr xss ys tree t0 -- if there's no parameter or no iterations + + algorithm = LBFGS funAndGrad Nothing + stop = ObjectiveRelativeTolerance 1e-10 :| [MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + t_opt = case minimizeLocal problem t0' of + Right sol -> solutionParams sol + Left e -> t0' + +-- | minimizes the likelihood assuming repeating parameters in the expression +minimizeNLLNonUnique :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double) +minimizeNLLNonUnique dist msErr niter xss ys tree t0 + | niter == 0 = (t0, f) + | n == 0 = (t0, f) + | otherwise = (fromStorableVector compMode t_opt, f) + where + t0' = toStorableVector t0 + (Sz n) = size t0 + (Sz m) = size ys + funAndGrad = second (toStorableVector . computeAs S) . gradNLLNonUnique dist msErr xss ys tree . fromStorableVector compMode + (f, _) = gradNLLNonUnique dist msErr xss ys tree t0 -- if there's no parameter or no iterations + + algorithm = LBFGS funAndGrad Nothing + stop = ObjectiveRelativeTolerance 1e-5 :| [MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + t_opt = case minimizeLocal problem t0' of + Right sol -> solutionParams sol + Left e -> t0' + +-- | minimizes the function while keeping the parameter ix fixed (used to calculate the profile) +minimizeNLLWithFixedParam :: Distribution -> Maybe Double -> Int -> SRMatrix -> PVector -> Fix SRTree -> Int -> PVector -> PVector +minimizeNLLWithFixedParam dist msErr niter xss ys tree ix t0 + | niter == 0 = t0 + | n == 0 = t0 + | n > m = t0 + | otherwise = fromStorableVector compMode t_opt + where + t0' = toStorableVector t0 + (Sz n) = size t0 + (Sz m) = size ys + setTo0 = (VS.// [(ix, 0.0)]) + funAndGrad = second (setTo0 . toStorableVector . computeAs S). gradNLLNonUnique dist msErr xss ys tree . fromStorableVector compMode + (f, _) = gradNLLNonUnique dist msErr xss ys tree t0 -- if there's no parameter or no iterations + + algorithm = LBFGS funAndGrad Nothing + stop = ObjectiveRelativeTolerance 1e-5 :| [MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + t_opt = case minimizeLocal problem t0' of + Right sol -> solutionParams sol + Left e -> t0' + +-- | minimizes using Gaussian likelihood +minimizeGaussian :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double) +minimizeGaussian = minimizeNLL Gaussian Nothing + +-- | minimizes using Binomial likelihood +minimizeBinomial :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double) +minimizeBinomial = minimizeNLL Bernoulli Nothing + +-- | minimizes using Poisson likelihood +minimizePoisson :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double) +minimizePoisson = minimizeNLL Poisson Nothing + +-- estimates the standard error if not provided +estimateSErr :: Distribution -> Maybe Double -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Int -> Maybe Double +estimateSErr Gaussian Nothing xss ys theta0 t nIter = Just err + where + theta = fst $ minimizeNLL Gaussian (Just 1) nIter xss ys t theta0 + (Sz m) = size ys + (Sz p) = size theta + ssr = sse xss ys t theta + err = sqrt $ ssr / fromIntegral (m - p) +estimateSErr _ (Just s) _ _ _ _ _ = Just s +estimateSErr _ _ _ _ _ _ _ = Nothing diff --git a/src/Data/SRTree.hs b/src/Data/SRTree.hs --- a/src/Data/SRTree.hs +++ b/src/Data/SRTree.hs @@ -1,7 +1,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -16,25 +16,22 @@ , Op(..) , param , var + , constv , arity , getChildren + , childrenOf + , replaceChildren + , getOperator , countNodes , countVarNodes , countConsts , countParams , countOccurrences - , deriveBy - , deriveByVar - , deriveByParam - , derivative - , forwardMode - , gradParamsFwd - , gradParamsRev - , evalFun - , evalOp - , inverseFunc - , evalTree + , countUniqueTokens + , numberOfVars + , getIntConsts , relabelParams + , relabelVars , constsToParam , floatConstsToParam , paramsToConst @@ -48,25 +45,22 @@ , Op(..) , param , var + , constv , arity , getChildren + , childrenOf + , replaceChildren + , getOperator , countNodes , countVarNodes , countConsts , countParams , countOccurrences - , deriveBy - , deriveByVar - , deriveByParam - , derivative - , forwardMode - , gradParamsFwd - , gradParamsRev - , evalFun - , evalOp - , inverseFunc - , evalTree + , countUniqueTokens + , numberOfVars + , getIntConsts , relabelParams + , relabelVars , constsToParam , floatConstsToParam , paramsToConst diff --git a/src/Data/SRTree/Datasets.hs b/src/Data/SRTree/Datasets.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Datasets.hs @@ -0,0 +1,214 @@ +{-# language ImportQualifiedPost #-} +{-# language ViewPatterns #-} +{-# language OverloadedStrings #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.Datasets +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables, ConstraintKinds +-- +-- Utility library to handle regression datasets +-- this module exports only the `loadDataset` function. +-- +----------------------------------------------------------------------------- +module Data.SRTree.Datasets ( loadDataset ) + 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)) +{-# INLINE loadMtx #-} + +-- | Returns true if the extension is .gz +isGZip :: FilePath -> Bool +isGZip = (== ".gz") . takeExtension +{-# INLINE isGZip #-} + +-- | Detects the separator automatically by +-- checking whether the use of each separator generates +-- the same amount of SRMatrix in every row and at least two SRMatrix. +-- +-- >>> detectSep ["x1,x2,x3,x4"] +-- ',' +detectSep :: [B.ByteString] -> Char +detectSep xss = go seps + where + seps = [' ','\t','|',':',';',','] + xss' = map B.strip xss + + -- consistency check whether all rows have the same + -- number of columns when spliting by this sep + allSameLen [] = True + allSameLen (y:ys) = y /= 1 && all (==y) ys + + go [] = error $ "CSV parsing error: unsupported separator. Supporter separators are " + <> intercalate "," (map show seps) + go (c:cs) = if allSameLen $ map (length . B.split c) xss' + then c + else go cs +{-# INLINE detectSep #-} + +-- | reads a file and returns a list of list of `ByteString` +-- corresponding to each element of the matrix. +-- The first row can be a header. +readFileToLines :: FilePath -> IO [[B.ByteString]] +readFileToLines filename = do + content <- removeBEmpty . toLines . toChar8 . unzip <$> BS.readFile filename + let sep = getSep content + pure . removeEmpty . map (B.split sep) $ content + where + getSep = detectSep . take 100 -- use only first 100 rows to detect separator + removeBEmpty = filter (not . B.null) + removeEmpty = filter (not . null) + toLines = B.split '\n' + unzip = if isGZip filename then decompress else id + toChar8 = B.pack . map (toEnum . fromEnum) . BS.unpack +{-# INLINE readFileToLines #-} + +-- | Splits the parameters from the filename +-- the expected format of the filename is *filename.ext:p1:p2:p3:p4* +-- where p1 and p2 is the starting and end rows for the training data, +-- by default p1 = 0 and p2 = number of rows - 1 +-- p3 is the target PVector, it can be a string corresponding to the header +-- or an index. +-- p4 is a comma separated list of SRMatrix (either index or name) to be used as +-- input variables. These will be renamed internally as x0, x1, ... in the order +-- of this list. +splitFileNameParams :: FilePath -> (FilePath, [B.ByteString]) +splitFileNameParams (B.pack -> filename) = (B.unpack fname, take 4 params) + where + (fname : params') = B.split ':' filename + -- fill up the empty parameters with an empty string + params = params' <> replicate (4 - min 4 (length params')) B.empty +{-# inline splitFileNameParams #-} + +-- | Tries to parse a string into an int +parseVal :: String -> Either String Int +parseVal xs = case readMaybe xs of + Nothing -> Left xs + Just x -> Right x +{-# inline parseVal #-} + +-- | Given a map between PVector name and indeces, +-- the target PVector and the variables SRMatrix, +-- returns the indices of the variables SRMatrix and the target +getColumns :: [(B.ByteString, Int)] -> B.ByteString -> B.ByteString -> ([Int], Int) +getColumns headerMap target columns = (ixs, iy) + where + n_cols = length headerMap + getIx c = case parseVal c of + -- if the PVector is a name, retrive the index + Left name -> case find ((== B.pack name) . fst) headerMap of + Nothing -> error $ "PVector name " <> name <> " does not exist." + Just v -> snd v + -- if it is an int, check if it is within range + Right v -> if v >= 0 && v < n_cols + then v + else error $ "PVector index " <> show v <> " out of range." + -- if the input variables SRMatrix are ommitted, use + -- every PVector except for iy + ixs = if B.null columns + then delete iy [0 .. n_cols - 1] + else map (getIx . B.unpack) $ B.split ',' columns + -- if the target PVector is ommitted, use the last one + iy = if B.null target + then n_cols - 1 + else getIx $ B.unpack target +{-# inline getColumns #-} + +-- | Given the start and end rows, it returns the +-- hmatrix extractors for the training and validation data +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) + where + st_ix = if null start + then 0 + else case readMaybe start of + Nothing -> error $ "Invalid starting row " <> start <> "." + Just x -> if x < 0 || x >= nRows + then error $ "Invalid starting row " <> show x <> "." + else x + end_ix = if null end + then nRows - 1 + else case readMaybe end of + Nothing -> error $ "Invalid end row " <> end <> "." + Just x -> if x < 0 || x >= nRows + then error $ "Invalid end row " <> show x <> "." + else x +{-# inline getRows #-} + +-- | `loadDataset` loads a dataset with a filename in the format: +-- filename.ext:start_row:end_row:target:features +-- it returns the X_train, y_train, X_test, y_test, varnames, target name +-- where varnames are a comma separated list of the name of the vars +-- and target name is the name of the target +-- +-- where +-- +-- **start_row:end_row** is the range of the training rows (default 0:nrows-1). +-- every other row not included in this range will be used as validation +-- **target** is either the name of the PVector (if the datafile has headers) or the index +-- of the target variable +-- **features** is a comma separated list of SRMatrix names or indices to be used as +-- input variables of the regression model. +loadDataset :: FilePath -> Bool -> IO ((SRMatrix, PVector, SRMatrix, PVector), String, String) +loadDataset filename hasHeader = do + csv <- readFileToLines fname + pure $ processData csv params hasHeader + where + (fname, params) = splitFileNameParams filename + +-- support function that does everything for loadDataset +processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> ((SRMatrix, PVector, SRMatrix, PVector), String, String) +processData csv params hasHeader = ((x_train, y_train, x_val, y_val) , varnames, targetname) + where + ncols = length $ head csv + nrows = length csv - fromEnum hasHeader + (header, content) = if hasHeader + then (zip (map B.strip $ head csv) [0..], tail csv) + else (map (\i -> (B.pack ('x' : show i), i)) [0 .. ncols-1], csv) + varnames = intercalate "," [B.unpack v | c <- ixs + , let v = fst . fromJust $ find ((==c).snd) header + ] + targetname = if hasHeader then (B.unpack . fst . fromJust . find ((==iy).snd) $ header) else "y" + -- get rows and SRMatrix indices + (st, end) = getRows (params !! 0) (params !! 1) nrows + (ixs, iy) = getColumns header (params !! 2) (params !! 3) + + -- load data and split sets + datum = loadMtx content + p = length ixs + + x = M.computeAs S $ M.throwEither $ M.stackInnerSlicesM $ map (datum >> showExpr . deriveBy False 0 $ 2 * "x0" * "x1" +-- "(2.0 * x1)" +-- >>> showExpr . deriveBy True 1 $ 2 * "x0" * "t0" - sqrt ("t1" * "x0") +-- "(-1.0 * ((1.0 / (2.0 * Sqrt((t1 * x0)))) * x0))" +deriveBy :: Bool -> Int -> Fix SRTree -> Fix SRTree +deriveBy p dx = fst (mutu alg1 alg2) + where + alg1 (Var ix) = if not p && ix == dx then 1 else 0 + alg1 (Param ix) = if p && ix == dx then 1 else 0 + alg1 (Const _) = 0 + alg1 (Uni f t) = derivative f (snd t) * fst t + alg1 (Bin Add l r) = fst l + fst r + alg1 (Bin Sub l r) = fst l - fst r + alg1 (Bin Mul l r) = fst l * snd r + snd l * fst r + alg1 (Bin Div l r) = (fst l * snd r - snd l * fst r) / snd r ** 2 + alg1 (Bin Power l r) = snd l ** (snd r - 1) * (snd r * fst l + snd l * log (snd l) * fst r) + alg1 (Bin PowerAbs l r) = (abs (snd l) ** (snd r)) * (fst r * log (abs (snd l)) + snd r * fst l / snd l) + alg1 (Bin AQ l r) = ((1 + snd r * snd r) * fst l - snd l * snd r * fst r) / (1 + snd r * snd r) ** 1.5 + + alg2 (Var ix) = var ix + alg2 (Param ix) = param ix + alg2 (Const c) = Fix (Const c) + alg2 (Uni f t) = Fix (Uni f $ snd t) + alg2 (Bin f l r) = Fix (Bin f (snd l) (snd r)) + +-- | Derivative of each supported function +-- For a function h(f) it returns the derivative dh/df +-- +-- >>> derivative Log 2.0 +-- 0.5 +derivative :: Floating a => Function -> a -> a +derivative Id = const 1 +derivative Abs = \x -> x / abs x +derivative Sin = cos +derivative Cos = negate.sin +derivative Tan = recip . (**2.0) . cos +derivative Sinh = cosh +derivative Cosh = sinh +derivative Tanh = (1-) . (**2.0) . tanh +derivative ASin = recip . sqrt . (1-) . (^2) +derivative ACos = negate . recip . sqrt . (1-) . (^2) +derivative ATan = recip . (1+) . (^2) +derivative ASinh = recip . sqrt . (1+) . (^2) +derivative ACosh = \x -> 1 / (sqrt (x-1) * sqrt (x+1)) +derivative ATanh = recip . (1-) . (^2) +derivative Sqrt = recip . (2*) . sqrt +derivative SqrtAbs = \x -> x / (2.0 * abs x ** (3.0/2.0)) +derivative Cbrt = recip . (3*) . (**(1/3)) . (^2) +derivative Square = (2*) +derivative Exp = exp +derivative Log = recip +derivative LogAbs = recip +derivative Recip = negate . recip . (^2) +derivative Cube = (3*) . (^2) +{-# INLINE derivative #-} + +-- | Second-order derivative of supported functions +-- +-- >>> doubleDerivative Log 2.0 +-- -0.25 +doubleDerivative :: Floating a => Function -> a -> a +doubleDerivative Id = const 0 +doubleDerivative Abs = const 0 +doubleDerivative Sin = negate.sin +doubleDerivative Cos = negate.cos +doubleDerivative Tan = \x -> 2 * sin x / (cos x) ^ 3 +doubleDerivative Sinh = sinh +doubleDerivative Cosh = cosh +doubleDerivative Tanh = \x -> -2 * tanh x * (1 / cosh x)^2 +doubleDerivative ASin = \x -> x / (1 - x^2)**(3/2) +doubleDerivative ACos = \x -> x / (1 - x^2)**(3/2) +doubleDerivative ATan = \x -> (-2*x) / (x^2 + 1)^2 +doubleDerivative ASinh = \x -> x / (x^2 + 1)**(3/2) -- check +doubleDerivative ACosh = \x -> 1 / (sqrt (x-1) * sqrt (x+1)) -- check +doubleDerivative ATanh = recip . (1-) . (^2) -- check +doubleDerivative Sqrt = \x -> -1 / (4 * sqrt x^3) +doubleDerivative SqrtAbs = \x -> (-x)*x/(4 * abs x ** (3.5)) +doubleDerivative Cbrt = \x -> -2 / (9 * x * (x^2)**(1/3)) +doubleDerivative Square = const 2 +doubleDerivative Exp = exp +doubleDerivative Log = negate . recip . (^2) +doubleDerivative LogAbs = negate . recip . (^2) +doubleDerivative Recip = (*2) . recip . (^3) +doubleDerivative Cube = (6*) +{-# INLINE doubleDerivative #-} + +-- | Symbolic derivative by a variable +deriveByVar :: Int -> Fix SRTree -> Fix SRTree +deriveByVar = deriveBy False +{-# INLINE deriveByVar #-} + +-- | Symbolic derivative by a parameter +deriveByParam :: Int -> Fix SRTree -> Fix SRTree +deriveByParam = deriveBy True +{-# INLINE deriveByParam #-} diff --git a/src/Data/SRTree/Eval.hs b/src/Data/SRTree/Eval.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Eval.hs @@ -0,0 +1,210 @@ +{-# LANGUAGE LambdaCase #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.Eval +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Evaluation of SRTree expressions +-- +----------------------------------------------------------------------------- +{-# LANGUAGE FlexibleInstances #-} +module Data.SRTree.Eval + ( evalTree + , evalOp + , evalFun + , cbrt + , inverseFunc + , invertibles + , evalInverse + , invright + , invleft + , replicateAs + , SRVector, PVector, SRMatrix + , compMode + ) + where + +import Data.Massiv.Array +import qualified Data.Massiv.Array as M +import Data.SRTree.Internal +import Data.SRTree.Recursion (Fix (..), cata) + +-- | Vector of target values +type SRVector = M.Array D Ix1 Double +-- | Vector of parameter values. Needs to be strict to be readily accesible. +type PVector = M.Array S Ix1 Double +-- | Matrix of features values +type SRMatrix = M.Array S Ix2 Double + +compMode :: M.Comp +compMode = M.Par' + +-- 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 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 + +-- 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 + +-- | 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 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 +{-# INLINE evalTree #-} + +-- evaluates an operator +evalOp :: Floating a => Op -> a -> a -> a +evalOp Add = (+) +evalOp Sub = (-) +evalOp Mul = (*) +evalOp Div = (/) +evalOp Power = (**) +evalOp PowerAbs = \l r -> abs l ** r +evalOp AQ = \l r -> l / sqrt(1 + r*r) +{-# INLINE evalOp #-} + +-- evaluates a function +evalFun :: Floating a => Function -> a -> a +evalFun Id = id +evalFun Abs = abs +evalFun Sin = sin +evalFun Cos = cos +evalFun Tan = tan +evalFun Sinh = sinh +evalFun Cosh = cosh +evalFun Tanh = tanh +evalFun ASin = asin +evalFun ACos = acos +evalFun ATan = atan +evalFun ASinh = asinh +evalFun ACosh = acosh +evalFun ATanh = atanh +evalFun Sqrt = sqrt +evalFun SqrtAbs = sqrt . abs +evalFun Cbrt = cbrt +evalFun Square = (^2) +evalFun Log = log +evalFun LogAbs = log . abs +evalFun Exp = exp +evalFun Recip = recip +evalFun Cube = (^3) +{-# INLINE evalFun #-} + +-- Cubic root +cbrt :: Floating a => a -> a +cbrt x = signum x * abs x ** (1/3) +{-# INLINE cbrt #-} + +-- | Returns the inverse of a function. This is a partial function. +inverseFunc :: Function -> Function +inverseFunc Id = Id +inverseFunc Sin = ASin +inverseFunc Cos = ACos +inverseFunc Tan = ATan +inverseFunc Sinh = ASinh +inverseFunc Cosh = ACosh +inverseFunc Tanh = ATanh +inverseFunc ASin = Sin +inverseFunc ACos = Cos +inverseFunc ATan = Tan +inverseFunc ASinh = Sinh +inverseFunc ACosh = Cosh +inverseFunc ATanh = Tanh +inverseFunc Sqrt = Square +inverseFunc Square = Sqrt +-- inverseFunc Cbrt = (^3) +inverseFunc Log = Exp +inverseFunc Exp = Log +inverseFunc Recip = Recip +-- inverseFunc Abs = Abs -- we assume abs(x) = sqrt(x^2) so y = sqrt(x^2) => x^2 = y^2 => x = sqrt(y^2) = x = abs(y) +inverseFunc x = error $ show x ++ " has no support for inverse function" +{-# INLINE inverseFunc #-} + +-- | evals the inverse of a function +evalInverse :: Floating a => Function -> a -> a +evalInverse Id = id +evalInverse Sin = asin +evalInverse Cos = acos +evalInverse Tan = atan +evalInverse Sinh = asinh +evalInverse Cosh = acosh +evalInverse Tanh = atanh +evalInverse ASin = sin +evalInverse ACos = cos +evalInverse ATan = tan +evalInverse ASinh = sinh +evalInverse ACosh = cosh +evalInverse ATanh = tanh +evalInverse Sqrt = (^2) +evalInverse SqrtAbs = (^2) +evalInverse Square = sqrt +evalInverse Cbrt = (^3) +evalInverse Log = exp +evalInverse LogAbs = exp +evalInverse Exp = log +evalInverse Abs = abs -- we assume abs(x) = sqrt(x^2) so y = sqrt(x^2) => x^2 = y^2 => x = sqrt(y^2) = x = abs(y) +evalInverse Recip = recip +evalInverse Cube = cbrt + +-- | evals the right inverse of an operator +invright :: Floating a => Op -> a -> (a -> a) +invright Add v = subtract v +invright Sub v = (+v) +invright Mul v = (/v) +invright Div v = (*v) +invright Power v = (**(1/v)) +invright PowerAbs v = (**(1/v)) +invright AQ v = (* sqrt (1 + v*v)) + +-- | evals the left inverse of an operator +invleft :: Floating a => Op -> a -> (a -> a) +invleft Add v = subtract v +invleft Sub v = (+v) . negate -- y = v - r => r = v - y +invleft Mul v = (/v) +invleft Div v = (v/) -- y = v / r => r = v/y +invleft Power v = logBase v -- (/(log v)) . log -- y = v ^ r log y = r log v r = log y / log v +invleft PowerAbs v = logBase v . abs +invleft AQ v = (v/) + +-- | List of invertible functions +invertibles :: [Function] +invertibles = [Id, Sin, Cos, Tan, Tanh, ASin, ACos, ATan, ATanh, Sqrt, Square, Log, Exp, Recip] 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 @@ -1,11 +1,12 @@ {-# language FlexibleInstances, DeriveFunctor #-} {-# language ScopedTypeVariables #-} {-# language RankNTypes #-} -{-# language ViewPatterns #-} +{-# language OverloadedStrings #-} +{-# language LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Internal --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -21,25 +22,22 @@ , Op(..) , param , var + , constv , arity , getChildren + , childrenOf + , replaceChildren + , getOperator , countNodes , countVarNodes , countConsts , countParams , countOccurrences - , deriveBy - , deriveByVar - , deriveByParam - , derivative - , forwardMode - , gradParamsFwd - , gradParamsRev - , evalFun - , evalOp - , inverseFunc - , evalTree + , countUniqueTokens + , numberOfVars + , getIntConsts , relabelParams + , relabelVars , constsToParam , floatConstsToParam , paramsToConst @@ -47,15 +45,11 @@ ) where -import Data.SRTree.Recursion ( Fix (..), cata, mutu, accu, cataM ) - -import qualified Data.Vector as V -import Data.Vector ((!)) -import Control.Monad.State -import qualified Data.DList as DL -import Data.Bifunctor (second) - -import Debug.Trace (trace) +import Control.Monad.State (MonadState (get), State, evalState, modify) +import Data.SRTree.Recursion (Fix (..), cata, cataM) +import qualified Data.Set as S +import Data.String (IsString (..)) +import Text.Read (readMaybe) -- | Tree structure to be used with Symbolic Regression algorithms. -- This structure is a fixed point of a n-ary tree. @@ -63,12 +57,14 @@ Var Int -- ^ index of the variables | Param Int -- ^ index of the parameter | Const Double -- ^ constant value, can be converted to a parameter + -- | 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) -- | Supported operators -data Op = Add | Sub | Mul | Div | Power +data Op = Add | Sub | Mul | Div | Power | PowerAbs | AQ deriving (Show, Read, Eq, Ord, Enum) -- | Supported functions @@ -88,10 +84,14 @@ | ACosh | ATanh | Sqrt + | SqrtAbs | Cbrt | Square | Log + | LogAbs | Exp + | Recip + | Cube deriving (Show, Read, Eq, Ord, Enum) -- | create a tree with a single node representing a variable @@ -102,6 +102,26 @@ param :: Int -> Fix SRTree param ix = Fix (Param ix) +-- | create a tree with a single node representing a constant value +constv :: Double -> Fix SRTree +constv x = Fix (Const x) + +-- | the instance of `IsString` allows us to +-- create a tree using a more practical notation: +-- +-- >>> :t "x0" + "t0" * sin("x1" * "t1") +-- Fix SRTree +-- +instance IsString (Fix SRTree) where + fromString [] = error "empty string for SRTree" + fromString ('x':ix) = case readMaybe ix of + Just iy -> Fix (Var iy) + Nothing -> error "wrong format for variable. It should be xi where i is an index. Ex.: \"x0\", \"x1\"." + fromString ('t':ix) = case readMaybe ix of + Just iy -> Fix (Param iy) + Nothing -> error "wrong format for parameter. It should be ti where i is an index. Ex.: \"t0\", \"t1\"." + fromString _ = error "A string can represent a variable or a parameter following the format xi or ti, respectivelly, where i is the index. Ex.: \"x0\", \"t0\"." + instance Num (Fix SRTree) where Fix (Const 0) + r = r l + Fix (Const 0) = l @@ -142,6 +162,9 @@ l / r = Fix $ Bin Div l r {-# INLINE (/) #-} + recip = Fix . Uni Recip + {-# INLINE recip #-} + fromRational = Fix . Const . fromRational {-# INLINE fromRational #-} @@ -188,6 +211,29 @@ logBase l r = log l / log r {-# INLINE logBase #-} +instance Foldable SRTree where + foldMap f = + \case + Bin op l r -> f l <> f r + Uni g t -> f t + _ -> mempty + +instance Traversable SRTree where + traverse f = + \case + Bin op l r -> Bin op <$> f l <*> f r + Uni g t -> Uni g <$> f t + Var ix -> pure (Var ix) + Param ix -> pure (Param ix) + Const x -> pure (Const x) + sequence = + \case + Bin op l r -> Bin op <$> l <*> r + Uni g t -> Uni g <$> t + Var ix -> pure (Var ix) + Param ix -> pure (Param ix) + Const x -> pure (Const x) + -- | Arity of the current node arity :: Fix SRTree -> Int arity = cata alg @@ -200,6 +246,10 @@ {-# INLINE arity #-} -- | Get the children of a node. Returns an empty list in case of a leaf node. +-- +-- >>> map showExpr . getChildren $ "x0" + 2 +-- ["x0", 2] +-- getChildren :: Fix SRTree -> [Fix SRTree] getChildren (Fix (Var {})) = [] getChildren (Fix (Param {})) = [] @@ -208,19 +258,53 @@ getChildren (Fix (Bin _ l r)) = [l, r] {-# INLINE getChildren #-} +-- | Get the children of an unfixed node +-- +childrenOf :: SRTree a -> [a] +childrenOf = + \case + Uni _ t -> [t] + Bin _ l r -> [l, r] + _ -> [] + +-- | replaces the children with elements from a list +replaceChildren :: [a] -> SRTree b -> SRTree a +replaceChildren [l, r] (Bin op _ _) = Bin op l r +replaceChildren [t] (Uni f _) = Uni f t +replaceChildren _ (Var ix) = Var ix +replaceChildren _ (Param ix) = Param ix +replaceChildren _ (Const x) = Const x +replaceChildren xs n = error "ERROR: trying to replace children with not enough elements." +{-# INLINE replaceChildren #-} + +-- | returns a node containing the operator and () as children +getOperator :: SRTree a -> SRTree () +getOperator (Bin op _ _) = Bin op () () +getOperator (Uni f _) = Uni f () +getOperator (Var ix) = Var ix +getOperator (Param ix) = Param ix +getOperator (Const x) = Const x +{-# INLINE getOperator #-} + -- | Count the number of nodes in a tree. -countNodes :: Fix SRTree -> Int +-- +-- >>> countNodes $ "x0" + 2 +-- 3 +countNodes :: Num a => Fix SRTree -> a countNodes = cata alg where - alg Var {} = 1 - alg Param {} = 1 - alg Const {} = 1 - alg (Uni _ t) = 1 + t + alg Var {} = 1 + alg Param {} = 1 + alg Const {} = 1 + alg (Uni _ t) = 1 + t alg (Bin _ l r) = 1 + l + r {-# INLINE countNodes #-} -- | Count the number of `Var` nodes -countVarNodes :: Fix SRTree -> Int +-- +-- >>> countVarNodes $ "x0" + 2 * ("x0" - sin "x1") +-- 3 +countVarNodes :: Num a => Fix SRTree -> a countVarNodes = cata alg where alg Var {} = 1 @@ -231,7 +315,10 @@ {-# INLINE countVarNodes #-} -- | Count the number of `Param` nodes -countParams :: Fix SRTree -> Int +-- +-- >>> countParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0" +-- 3 +countParams :: Num a => Fix SRTree -> a countParams = cata alg where alg Var {} = 0 @@ -242,7 +329,10 @@ {-# INLINE countParams #-} -- | Count the number of const nodes -countConsts :: Fix SRTree -> Int +-- +-- >>> countConsts $ "x0"* 2 + 3 * sin "x0" +-- 2 +countConsts :: Num a => Fix SRTree -> a countConsts = cata alg where alg Var {} = 0 @@ -253,7 +343,10 @@ {-# INLINE countConsts #-} -- | Count the occurrences of variable indexed as `ix` -countOccurrences :: Int -> Fix SRTree -> Int +-- +-- >>> countOccurrences 0 $ "x0"* 2 + 3 * sin "x0" + "x1" +-- 2 +countOccurrences :: Num a => Int -> Fix SRTree -> a countOccurrences ix = cata alg where alg (Var iy) = if ix == iy then 1 else 0 @@ -263,302 +356,142 @@ alg (Bin _ l r) = l + r {-# INLINE countOccurrences #-} --- | 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 :: (Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> a -evalTree xss params f = cata alg - where - alg (Var ix) = xss ! ix - alg (Param ix) = f $ params ! ix - alg (Const c) = f c - alg (Uni g t) = evalFun g t - alg (Bin op l r) = evalOp op l r -{-# INLINE evalTree #-} - -evalOp :: Floating a => Op -> a -> a -> a -evalOp Add = (+) -evalOp Sub = (-) -evalOp Mul = (*) -evalOp Div = (/) -evalOp Power = (**) -{-# INLINE evalOp #-} - -evalFun :: Floating a => Function -> a -> a -evalFun Id = id -evalFun Abs = abs -evalFun Sin = sin -evalFun Cos = cos -evalFun Tan = tan -evalFun Sinh = sinh -evalFun Cosh = cosh -evalFun Tanh = tanh -evalFun ASin = asin -evalFun ACos = acos -evalFun ATan = atan -evalFun ASinh = asinh -evalFun ACosh = acosh -evalFun ATanh = atanh -evalFun Sqrt = sqrt -evalFun Cbrt = cbrt -evalFun Square = (^2) -evalFun Log = log -evalFun Exp = exp -{-# INLINE evalFun #-} - --- | Cubic root -cbrt :: Floating val => val -> val -cbrt x = signum x * abs x ** (1/3) -{-# INLINE cbrt #-} - --- | Returns the inverse of a function. This is a partial function. -inverseFunc :: Function -> Function -inverseFunc Id = Id -inverseFunc Sin = ASin -inverseFunc Cos = ACos -inverseFunc Tan = ATan -inverseFunc Tanh = ATanh -inverseFunc ASin = Sin -inverseFunc ACos = Cos -inverseFunc ATan = Tan -inverseFunc ATanh = Tanh -inverseFunc Sqrt = Square -inverseFunc Square = Sqrt -inverseFunc Log = Exp -inverseFunc Exp = Log -inverseFunc x = error $ show x ++ " has no support for inverse function" -{-# INLINE inverseFunc #-} - --- | Creates the symbolic partial derivative of a tree by variable `dx` (if `p` is `False`) --- or parameter `dx` (if `p` is `True`). -deriveBy :: Bool -> Int -> Fix SRTree -> Fix SRTree -deriveBy p dx = fst (mutu alg1 alg2) +-- | counts the number of unique tokens +-- +-- >>> countUniqueTokens $ "x0" + ("x1" * "x0" - sin ("x0" ** 2)) +-- 8 +countUniqueTokens :: Num a => Fix SRTree -> a +countUniqueTokens = len . cata alg where - alg1 (Var ix) = if not p && ix == dx then 1 else 0 - alg1 (Param ix) = if p && ix == dx then 1 else 0 - alg1 (Const _) = 0 - alg1 (Uni f t) = derivative f (snd t) * fst t - alg1 (Bin Add l r) = fst l + fst r - alg1 (Bin Sub l r) = fst l - fst r - alg1 (Bin Mul l r) = fst l * snd r + snd l * fst r - alg1 (Bin Div l r) = (fst l * snd r - snd l * fst r) / snd r ** 2 - alg1 (Bin Power l r) = snd l ** (snd r - 1) * (snd r * fst l + snd l * log (snd l) * fst r) - - alg2 (Var ix) = var ix - alg2 (Param ix) = param ix - alg2 (Const c) = Fix (Const c) - alg2 (Uni f t) = Fix (Uni f $ snd t) - alg2 (Bin f l r) = Fix (Bin f (snd l) (snd r)) - -newtype Tape a = Tape { untape :: [a] } deriving (Show, Functor) - -instance Num a => Num (Tape a) where - (Tape x) + (Tape y) = Tape $ zipWith (+) x y - (Tape x) - (Tape y) = Tape $ zipWith (-) x y - (Tape x) * (Tape y) = Tape $ zipWith (*) x y - abs (Tape x) = Tape (map abs x) - signum (Tape x) = Tape (map signum x) - fromInteger x = Tape [fromInteger x] - negate (Tape x) = Tape $ map (*(-1)) x -instance Floating a => Floating (Tape a) where - pi = Tape [pi] - exp (Tape x) = Tape (map exp x) - log (Tape x) = Tape (map log x) - sqrt (Tape x) = Tape (map sqrt x) - sin (Tape x) = Tape (map sin x) - cos (Tape x) = Tape (map cos x) - tan (Tape x) = Tape (map tan x) - asin (Tape x) = Tape (map asin x) - acos (Tape x) = Tape (map acos x) - atan (Tape x) = Tape (map atan x) - sinh (Tape x) = Tape (map sinh x) - cosh (Tape x) = Tape (map cosh x) - tanh (Tape x) = Tape (map tanh x) - asinh (Tape x) = Tape (map asinh x) - acosh (Tape x) = Tape (map acosh x) - atanh (Tape x) = Tape (map atanh x) - (Tape x) ** (Tape y) = Tape $ zipWith (**) x y -instance Fractional a => Fractional (Tape a) where - fromRational x = Tape [fromRational x] - (Tape x) / (Tape y) = Tape $ zipWith (/) x y - recip (Tape x) = Tape $ map recip x + len (a, b, c, d, e) = fromIntegral $ length a + length b + length c + length d + length e + alg (Var ix) = (mempty, mempty, S.singleton ix, mempty, mempty) + alg (Param _) = (mempty, mempty, mempty, S.singleton 1, mempty) + alg (Const _) = (mempty, mempty, mempty, mempty, S.singleton 1) + alg (Uni f t) = (mempty, S.singleton f, mempty, mempty, mempty) <> t + alg (Bin op l r) = (S.singleton op, mempty, mempty, mempty, mempty) <> l <> r +{-# INLINE countUniqueTokens #-} --- | Calculates the numerical derivative of a tree using forward mode --- provided a vector of variable values `xss`, a vector of parameter values `theta` and --- a function that changes a Double value to the type of the variable values. -forwardMode :: (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> [a] -forwardMode xss theta f = untape . fst (mutu alg1 alg2) +-- | return the number of unique variables +-- +-- >>> numberOfVars $ "x0" + 2 * ("x0" - sin "x1") +-- 2 +numberOfVars :: Num a => Fix SRTree -> a +numberOfVars = fromIntegral . S.size . cata alg where - n = V.length theta - repMat v = Tape $ replicate n v - zeroes = repMat $ f 0 - twos = repMat $ f 2 - tapeXs = [repMat $ xss ! ix | ix <- [0 .. V.length xss - 1]] - tapeTheta = [repMat $ f (theta ! ix) | ix <- [0 .. n - 1]] - paramVec = [ Tape [if ix==iy then f 1 else f 0 | iy <- [0 .. n-1]] | ix <- [0 .. n-1] ] - - alg1 (Var ix) = zeroes - alg1 (Param ix) = paramVec !! ix - alg1 (Const _) = zeroes - alg1 (Uni f t) = derivative f (snd t) * fst t - alg1 (Bin Add l r) = fst l + fst r - alg1 (Bin Sub l r) = fst l - fst r - alg1 (Bin Mul l r) = (fst l * snd r) + (snd l * fst r) - alg1 (Bin Div l r) = ((fst l * snd r) - (snd l * fst r)) / snd r ** twos - alg1 (Bin Power l r) = snd l ** (snd r - 1) * ((snd r * fst l) + (snd l * log (snd l) * fst r)) - - alg2 (Var ix) = tapeXs !! ix - alg2 (Param ix) = tapeTheta !! ix - alg2 (Const c) = repMat $ f c - alg2 (Uni g t) = fmap (evalFun g) (snd t) - alg2 (Bin op l r) = evalOp op (snd l) (snd r) + alg (Uni f t) = t + alg (Bin op l r) = l <> r + alg (Var ix) = S.singleton ix + alg _ = mempty +{-# INLINE numberOfVars #-} --- | The function `gradParams` 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`. -gradParamsFwd :: (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> (a, [a]) -gradParamsFwd xss theta f = second DL.toList . cata alg +-- | returns the integer constants. We assume an integer constant +-- as those values in which `floor x == ceiling x`. +-- +-- >>> getIntConsts $ "x0" + 2 * "x1" ** 3 - 3.14 +-- [2.0,3.0] +getIntConsts :: Fix SRTree -> [Double] +getIntConsts = cata alg where - n = V.length theta - - alg (Var ix) = (xss ! ix, DL.empty) - alg (Param ix) = (f $ theta ! ix, DL.singleton 1) - alg (Const c) = (f c, DL.empty) - alg (Uni f (v, gs)) = let v' = evalFun f v - dv = derivative f v - in (v', DL.map (*dv) gs) - alg (Bin Add (v1, l) (v2, r)) = (v1+v2, DL.append l r) - alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, DL.append l (DL.map negate r)) - alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, DL.append (DL.map (*v2) l) (DL.map (*v1) r)) - alg (Bin Div (v1, l) (v2, r)) = let dv = (-v1/v2^2) - in (v1/v2, DL.append (DL.map (/v2) l) (DL.map (*dv) r)) - alg (Bin Power (v1, l) (v2, r)) = let dv1 = v1 ** (v2 - 1) - dv2 = v1 * log v1 - in (v1 ** v2, DL.map (*dv1) (DL.append (DL.map (*v2) l) (DL.map (*dv2) r))) - -data TupleF a b = S a | T a b | B a b b deriving Functor -- hi, I'm a tree -type Tuple a = Fix (TupleF a) + alg (Uni f t) = t + alg (Bin op l r) = l <> r + alg (Var ix) = [] + alg (Param _) = [] + alg (Const x) = [x | floor x == ceiling x] +{-# INLINE getIntConsts #-} -gradParamsRev :: forall a . (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> (a, [a]) -gradParamsRev xss theta f t = (getTop fwdMode, DL.toList g) +-- | Relabel the parameters indices incrementaly starting from 0 +-- +-- >>> showExpr . relabelParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0" +-- "x0" + "t0" * sin ("t1" + "x1") - "t2" +relabelParams :: Fix SRTree -> Fix SRTree +relabelParams t = cataM leftToRight alg t `evalState` 0 where - fwdMode = cata forward t - g = accu reverse combine t (1, fwdMode) - - oneTpl x = Fix $ S x - tuple x y = Fix $ T x y - branch x y z = Fix $ B x y z - getTop (Fix (S x)) = x - getTop (Fix (T x y)) = x - getTop (Fix (B x y z)) = x - unCons (Fix (T x y)) = y - getBranches (Fix (B x y z)) = (y,z) - - forward (Var ix) = oneTpl (xss ! ix) - forward (Param ix) = oneTpl (f $ theta ! ix) - forward (Const c) = oneTpl (f c) - forward (Uni f t) = let v = getTop t - in tuple (evalFun f v) t - forward (Bin op l r) = let vl = getTop l - vr = getTop r - in branch (evalOp op vl vr) l r - - reverse (Var ix) (dx, _) = Var ix - reverse (Param ix) (dx, _) = Param ix - reverse (Const v) (dx, _) = Const v - reverse (Uni f t) (dx, unCons -> v) = Uni f (t, (dx * (derivative f $ getTop v), v)) - reverse (Bin op l r) (dx, getBranches -> (vl, vr)) = let (dxl, dxr) = diff op dx (getTop vl) (getTop vr) - in Bin op (l, (dxl, vl)) (r, (dxr, vr)) - - diff Add dx vl vr = (dx, dx) - diff Sub dx vl vr = (dx, negate dx) - diff Mul dx vl vr = (dx * vr, dx * vl) - diff Div dx vl vr = (dx / vr, dx * (-vl/vr^2)) - diff Power dx vl vr = let dxl = dx * vl ** (vr - 1) - dv2 = vl * log vl - in (dxl * vr, dxl * dv2) - - combine (Var ix) s = DL.empty - combine (Param ix) s = DL.singleton $ fst s - combine (Const c) s = DL.empty - combine (Uni _ gs) s = gs - combine (Bin op l r) s = DL.append l r - -derivative :: Floating a => Function -> a -> a -derivative Id = const 1 -derivative Abs = \x -> x / abs x -derivative Sin = cos -derivative Cos = negate.sin -derivative Tan = recip . (**2.0) . cos -derivative Sinh = cosh -derivative Cosh = sinh -derivative Tanh = (1-) . (**2.0) . tanh -derivative ASin = recip . sqrt . (1-) . (^2) -derivative ACos = negate . recip . sqrt . (1-) . (^2) -derivative ATan = recip . (1+) . (^2) -derivative ASinh = recip . sqrt . (1+) . (^2) -derivative ACosh = \x -> 1 / (sqrt (x-1) * sqrt (x+1)) -derivative ATanh = recip . (1-) . (^2) -derivative Sqrt = recip . (2*) . sqrt -derivative Cbrt = recip . (3*) . cbrt . (^2) -derivative Square = (2*) -derivative Exp = exp -derivative Log = recip -{-# INLINE derivative #-} - --- | Symbolic derivative by a variable -deriveByVar :: Int -> Fix SRTree -> Fix SRTree -deriveByVar = deriveBy False + -- | leftToRight (left to right) defines the sequence of processing + 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) --- | Symbolic derivative by a parameter -deriveByParam :: Int -> Fix SRTree -> Fix SRTree -deriveByParam = deriveBy True + -- | any time we reach a Param ix, it replaces ix with current state + -- and increments one to the state. + alg :: SRTree (Fix SRTree) -> State Int (Fix SRTree) + alg (Var ix) = pure $ var ix + alg (Param ix) = do iy <- get; modify (+1); pure (param iy) + alg (Const c) = pure $ Fix $ Const c + alg (Uni f t) = pure $ Fix (Uni f t) + alg (Bin f l r) = pure $ Fix (Bin f l r) --- | Relabel the parameters incrementaly starting from 0 -relabelParams :: Fix SRTree -> Fix SRTree -relabelParams t = cataM lTor alg t `evalState` 0 +-- | Relabel the parameters indices incrementaly starting from 0 +-- +-- >>> showExpr . relabelParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0" +-- "x0" + "t0" * sin ("t1" + "x1") - "t2" +relabelVars :: Fix SRTree -> Fix SRTree +relabelVars t = cataM leftToRight alg t `evalState` 0 where - lTor (Uni f mt) = Uni f <$> mt; - lTor (Bin f ml mr) = Bin f <$> ml <*> mr - lTor (Var ix) = pure (Var ix) - lTor (Param ix) = pure (Param ix) - lTor (Const c) = pure (Const c) + -- | leftToRight (left to right) defines the sequence of processing + 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) + -- | any time we reach a Param ix, it replaces ix with current state + -- and increments one to the state. alg :: SRTree (Fix SRTree) -> State Int (Fix SRTree) - alg (Var ix) = pure $ var ix - alg (Param ix) = do iy <- get; modify (+1); pure (param iy) - alg (Const c) = pure $ Fix $ Const c - alg (Uni f t) = pure $ Fix (Uni f t) + alg (Var ix) = do iy <- get; modify (+1); pure (var iy) + alg (Param ix) = pure $ param ix + alg (Const c) = pure $ Fix $ Const c + alg (Uni f t) = pure $ Fix (Uni f t) alg (Bin f l r) = pure $ Fix (Bin f l r) -- | Change constant values to a parameter, returning the changed tree and a list -- of parameter values +-- +-- >>> snd . constsToParam $ "x0" * 2 + 3.14 * sin (5 * "x1") +-- [2.0,3.14,5.0] constsToParam :: Fix SRTree -> (Fix SRTree, [Double]) constsToParam = first relabelParams . cata alg where first f (x, y) = (f x, y) - alg (Var ix) = (Fix $ Var ix, []) - alg (Param ix) = (Fix $ Param ix, [1.0]) - alg (Const c) = (Fix $ Param 0, [c]) - alg (Uni f t) = (Fix $ Uni f (fst t), snd t) + -- | If the tree already contains a parameter + -- it will return a default value of 1.0 + -- whenever it finds a constant, it changes that + -- to a parameter and adds its content to the singleton list + alg (Var ix) = (Fix $ Var ix, []) + alg (Param ix) = (Fix $ Param ix, [1.0]) + alg (Const c) = (Fix $ Param 0, [c]) + alg (Uni f t) = (Fix $ Uni f (fst t), snd t) alg (Bin f l r) = (Fix (Bin f (fst l) (fst r)), snd l <> snd r) -- | Same as `constsToParam` but does not change constant values that -- can be converted to integer without loss of precision +-- +-- >>> snd . floatConstsToParam $ "x0" * 2 + 3.14 * sin (5 * "x1") +-- [3.14] floatConstsToParam :: Fix SRTree -> (Fix SRTree, [Double]) floatConstsToParam = first relabelParams . cata alg where - first f (x, y) = (f x, y) + first f (x, y) = (f x, y) + combine f (x, y) (z, w) = (f x z, y <> w) + isInt x = floor x == ceiling x - alg (Var ix) = (Fix $ Var ix, []) - alg (Param ix) = (Fix $ Param ix, [1.0]) - alg (Const c) = if floor c == ceiling c then (Fix $ Const c, []) else (Fix $ Param 0, [c]) - alg (Uni f t) = (Fix $ Uni f (fst t), snd t) - alg (Bin f l r) = (Fix (Bin f (fst l) (fst r)), snd l <> snd r) + alg (Var ix) = (var ix, []) + alg (Param ix) = (param ix, [1.0]) + alg (Const c) = if isInt c then (constv c, []) else (param 0, [c]) + alg (Uni f t) = first (Fix . Uni f) t -- (Fix $ Uni f (fst t), snd t) + alg (Bin f l r) = combine ((Fix .) . Bin f) l r -- (Fix (Bin f (fst l) (fst r)), snd l <> snd r) -- | Convert the parameters into constants in the tree +-- +-- >>> showExpr . paramsToConst [1.1, 2.2, 3.3] $ "x0" + "t0" * sin ("t1" * "x0" - "t2") +-- x0 + 1.1 * sin(2.2 * x0 - 3.3) paramsToConst :: [Double] -> Fix SRTree -> Fix SRTree paramsToConst theta = cata alg where - alg (Var ix) = Fix $ Var ix - alg (Param ix) = Fix $ Const (theta !! ix) - alg (Const c) = Fix $ Const c - alg (Uni f t) = Fix $ Uni f t + alg (Var ix) = Fix $ Var ix + alg (Param ix) = Fix $ Const (theta !! ix) + alg (Const c) = Fix $ Const c + alg (Uni f t) = Fix $ Uni f t alg (Bin f l r) = Fix $ Bin f l r diff --git a/src/Data/SRTree/Print.hs b/src/Data/SRTree/Print.hs --- a/src/Data/SRTree/Print.hs +++ b/src/Data/SRTree/Print.hs @@ -1,7 +1,9 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Print --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -12,7 +14,9 @@ ----------------------------------------------------------------------------- module Data.SRTree.Print ( showExpr + , showExprWithVars , printExpr + , printExprWithVars , showTikz , printTikz , showPython @@ -22,42 +26,89 @@ ) where -import Control.Monad.Reader ( asks, runReader, Reader ) -import Data.Char ( toLower ) - +import Control.Monad.Reader (Reader, asks, runReader) +import Data.Char (toLower) import Data.SRTree.Internal -import Data.SRTree.Recursion +import Data.SRTree.Recursion (cata) +-- | converts a tree with protected operators to +-- a conventional math tree +removeProtection :: Fix SRTree -> Fix SRTree +removeProtection = cata $ + \case + Var ix -> Fix (Var ix) + Param ix -> Fix (Param ix) + Const x -> Fix (Const x) + Uni SqrtAbs t -> sqrt (abs t) + Uni LogAbs t -> log (abs t) + Uni Cube t -> t ** 3 + Uni f t -> Fix (Uni f t) + Bin AQ l r -> l / sqrt (1 + r*r) + Bin PowerAbs l r -> abs l ** r + Bin op l r -> Fix (Bin op l r) + +-- | convert a tree into a string in math notation +-- +-- >>> showExpr $ "x0" + sin ( tanh ("t0" + 2) ) +-- "(x0 + Sin(Tanh((t0 + 2.0))))" showExpr :: Fix SRTree -> String -showExpr = cata alg - where - alg (Var ix) = 'x' : show ix - alg (Param ix) = 't' : show ix - alg (Const c) = show c - alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"] - alg (Uni f t) = concat [show f, "(", t, ")"] +showExpr = cata alg . removeProtection + where alg = \case + Var ix -> 'x' : show ix + Param ix -> 't' : show ix + Const c -> show c + Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"] + Uni f t -> concat [show f, "(", t, ")"] +-- | convert a tree into a string in math notation +-- given named vars. +-- +-- >>> showExprWithVar ["mu", "eps"] $ "x0" + sin ( "x1" * tanh ("t0" + 2) ) +-- "(mu + Sin(Tanh(eps * (t0 + 2.0))))" +showExprWithVars :: [String] -> Fix SRTree -> String +showExprWithVars varnames = cata alg . removeProtection + where alg = \case + Var ix -> varnames !! ix + Param ix -> 't' : show ix + Const c -> show c + Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"] + Uni f t -> concat [show f, "(", t, ")"] + +-- | prints the expression printExpr :: Fix SRTree -> IO () printExpr = putStrLn . showExpr +-- | prints the expression +printExprWithVars :: [String] -> Fix SRTree -> IO () +printExprWithVars varnames = putStrLn . showExprWithVars varnames + +-- how to display an operator +showOp :: Op -> String showOp Add = "+" showOp Sub = "-" showOp Mul = "*" showOp Div = "/" showOp Power = "^" +showOp AQ = "_aq_" +showOp PowerAbs = "||^" {-# INLINE showOp #-} -- | Displays a tree as a numpy compatible expression. +-- +-- >>> showPython $ "x0" + sin ( tanh ("t0" + 2) ) +-- "(x[:, 0] + np.sin(np.tanh((t[:, 0] + 2.0))))" showPython :: Fix SRTree -> String -showPython = cata alg +showPython = cata alg . removeProtection where - alg (Var ix) = concat ["x[:, ", show ix, "]"] - alg (Param ix) = concat ["t[:, ", show ix, "]"] - alg (Const c) = show c - alg (Bin Power l r) = concat [l, " ** ", r] - alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"] - alg (Uni f t) = concat [pyFun f, "(", t, ")"] + alg = \case + Var ix -> concat ["x[:, ", show ix, "]"] + Param ix -> concat ["t[:, ", show ix, "]"] + Const c -> show c + Bin Power l r -> concat [l, " ** ", r] + Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"] + Uni f t -> concat [pyFun f, "(", t, ")"] + pyFun Id = "" pyFun Abs = "np.abs" pyFun Sin = "np.sin" @@ -76,39 +127,49 @@ pyFun Square = "np.square" pyFun Log = "np.log" pyFun Exp = "np.exp" + pyFun Cbrt = "np.cbrt" + pyFun Recip = "np.reciprocal" +-- | print the expression in numpy notation printPython :: Fix SRTree -> IO () printPython = putStrLn . showPython --- | Displays a tree as a sympy compatible expression. +-- | Displays a tree as a LaTeX compatible expression. +-- +-- >>> showLatex $ "x0" + sin ( tanh ("t0" + 2) ) +-- "\\left(x_{, 0} + \\operatorname{sin}(\\operatorname{tanh}(\\left(\\theta_{, 0} + 2.0\\right)))\\right)" showLatex :: Fix SRTree -> String -showLatex = cata alg +showLatex = cata alg . removeProtection where - alg (Var ix) = concat ["x_{, ", show ix, "}"] - alg (Param ix) = concat ["\\theta_{, ", show ix, "}"] - alg (Const c) = show c - alg (Bin Power l r) = concat [l, "^{", r, "}"] - alg (Bin op l r) = concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"] - alg (Uni Abs t) = concat ["\\left |", t, "\\right |"] - alg (Uni f t) = concat [showLatexFun f, "(", t, ")"] + alg = \case + Var ix -> concat ["x_{, ", show ix, "}"] + Param ix -> concat ["\\theta_{, ", show ix, "}"] + Const c -> show c + Bin Power l r -> concat [l, "^{", r, "}"] + Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"] + Uni Abs t -> concat ["\\left |", t, "\\right |"] + Uni f t -> concat [showLatexFun f, "(", t, ")"] showLatexFun :: Function -> String showLatexFun f = mconcat ["\\operatorname{", map toLower $ show f, "}"] {-# INLINE showLatexFun #-} +-- | prints expression in LaTeX notation. printLatex :: Fix SRTree -> IO () printLatex = putStrLn . showLatex -- | Displays a tree in Tikz format showTikz :: Fix SRTree -> String -showTikz = cata alg +showTikz = cata alg . removeProtection where + alg = \case + Var ix -> concat ["[$x_{, ", show ix, "}$]\n"] + Param ix -> concat ["[$\\theta_{, ", show ix, "}$]\n"] + Const c -> concat ["[$", show (roundN 2 c), "$]\n"] + Bin op l r -> concat ["[", showOpTikz op, l, r, "]\n"] + Uni f t -> concat ["[", map toLower $ show f, t, "]\n"] + roundN n x = let ten = 10^n in (/ ten) . fromIntegral . round $ x*ten - alg (Var ix) = concat ["[$x_{, ", show ix, "}$]\n"] - alg (Param ix) = concat ["[$\\theta_{, ", show ix, "}$]\n"] - alg (Const c) = concat ["[$", show (roundN 2 c), "$]\n"] - alg (Bin op l r) = concat ["[", showOpTikz op, l, r, "]\n"] - alg (Uni f t) = concat ["[", map toLower $ show f, t, "]\n"] showOpTikz Add = "+\n" showOpTikz Sub = "-\n" @@ -116,4 +177,6 @@ showOpTikz Div = "รท\n" showOpTikz Power = "\\^{}\n" +-- | prints the tree in TikZ format +printTikz :: Fix SRTree -> IO () printTikz = putStrLn . showTikz 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 @@ -2,7 +2,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Random --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -29,13 +29,11 @@ ) where -import System.Random -import Control.Monad.State -import Control.Monad.Reader +import Control.Monad.Reader (ReaderT, asks, runReaderT) +import Control.Monad.State.Strict ( MonadState(state), MonadTrans(lift), StateT ) import Data.Maybe (fromJust) - import Data.SRTree.Internal -import Data.SRTree.Recursion +import System.Random (Random (random, randomR), StdGen, mkStdGen) -- * Class definition of properties that a certain parameter type has. -- @@ -90,10 +88,10 @@ {-# INLINE replaceChild #-} -- Replace the children of a binary tree. -replaceChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree) -replaceChildren (Fix (Bin f _ _)) l r = Just $ Fix (Bin f l r) -replaceChildren _ _ _ = Nothing -{-# INLINE replaceChildren #-} +replaceFixChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree) +replaceFixChildren (Fix (Bin f _ _)) l r = Just $ Fix (Bin f l r) +replaceFixChildren _ _ _ = Nothing +{-# INLINE replaceFixChildren #-} -- | RndTree is a Monad Transformer to generate random trees of type `SRTree ix val` -- given the parameters `p ix val` using the random number generator `StdGen`. @@ -149,6 +147,11 @@ 6 -> pure . Fix $ Bin Power 0 0 -- | Returns a random tree with a limited budget, the parameter `p` must have every property. +-- +-- >>> let treeGen = runReaderT (randomTree 12) (P [0,1] (-10, 10) (2, 3) [Log, Exp]) +-- >>> tree <- evalStateT treeGen (mkStdGen 52) +-- >>> showExpr tree +-- "(-2.7631152121655838 / Exp((x0 / ((x0 * -7.681722660704317) - Log(3.378309080134594)))))" randomTree :: HasEverything p => Int -> RndTree p randomTree 0 = do coin <- lift toss @@ -160,9 +163,14 @@ fromJust <$> case arity node of 0 -> pure $ Just node 1 -> replaceChild node <$> randomTree (budget - 1) - 2 -> replaceChildren node <$> randomTree (budget `div` 2) <*> randomTree (budget `div` 2) + 2 -> replaceFixChildren node <$> randomTree (budget `div` 2) <*> randomTree (budget `div` 2) -- | Returns a random tree with a approximately a number `n` of nodes, the parameter `p` must have every property. +-- +-- >>> let treeGen = runReaderT (randomTreeBalanced 10) (P [0,1] (-10, 10) (2, 3) [Log, Exp]) +-- >>> tree <- evalStateT treeGen (mkStdGen 42) +-- >>> showExpr tree +-- "Exp(Log((((7.784360517385774 * x0) - (3.6412224491658223 ^ x1)) ^ ((x0 ^ -4.09764995657091) + Log(-7.710216839988497)))))" randomTreeBalanced :: HasEverything p => Int -> RndTree p randomTreeBalanced n | n <= 1 = do coin <- lift toss @@ -173,4 +181,4 @@ node <- randomNonTerminal fromJust <$> case arity node of 1 -> replaceChild node <$> randomTreeBalanced (n - 1) - 2 -> replaceChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) + 2 -> replaceFixChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) diff --git a/src/Data/SRTree/Recursion.hs b/src/Data/SRTree/Recursion.hs --- a/src/Data/SRTree/Recursion.hs +++ b/src/Data/SRTree/Recursion.hs @@ -1,5 +1,17 @@ {-# language RankNTypes #-} {-# language DeriveFunctor #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.Recursion +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Recursion schemes +-- +----------------------------------------------------------------------------- module Data.SRTree.Recursion where import Control.Monad ( (>=>) ) diff --git a/src/Text/ParseSR.hs b/src/Text/ParseSR.hs new file mode 100644 --- /dev/null +++ b/src/Text/ParseSR.hs @@ -0,0 +1,302 @@ +{-# language OverloadedStrings #-} +----------------------------------------------------------------------------- +-- | +-- Module : Text.ParseSR +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to parse a string representing an expression +-- +----------------------------------------------------------------------------- +module Text.ParseSR ( parseSR, showOutput, SRAlgs(..), Output(..) ) + where + +import Control.Applicative ((<|>)) +import Data.Attoparsec.ByteString.Char8 +import Data.Attoparsec.Expr +import qualified Data.ByteString.Char8 as B +import Data.Char (toLower) +import Data.List (sortOn) +import Data.SRTree +import qualified Data.SRTree.Print as P +import Debug.Trace (trace) + +-- * Data types + +-- | Parser of a symbolic regression tree with `Int` variable index and +-- numerical values represented as `Double`. The numerical values type +-- can be changed with `fmap`. +type ParseTree = Parser (Fix SRTree) + +-- * Data types and caller functions + +-- | Supported algorithms. +data SRAlgs = TIR | HL | OPERON | BINGO | GOMEA | PYSR | SBP | EPLEX deriving (Show, Read, Enum, Bounded) + +-- | Supported outputs. +data Output = PYTHON | MATH | TIKZ | LATEX deriving (Show, Read, Enum, Bounded) + +-- | Returns the corresponding function from Data.SRTree.Print for a given `Output`. +showOutput :: Output -> Fix SRTree -> String +showOutput PYTHON = P.showPython +showOutput MATH = P.showExpr +showOutput TIKZ = P.showTikz +showOutput LATEX = P.showLatex + +-- | Calls the corresponding parser for a given `SRAlgs` +-- +-- >>> fmap (showOutput MATH) $ parseSR OPERON "lambda,theta" False "lambda ^ 2 - sin(theta*3*lambda)" +-- Right "((x0 ^ 2.0) - Sin(((x1 * 3.0) * x0)))" +parseSR :: SRAlgs -> B.ByteString -> Bool -> B.ByteString -> Either String (Fix SRTree) +parseSR HL header reparam = eitherResult . (`feed` "") . parse (parseHL reparam $ splitHeader header) . putEOL . B.strip +parseSR BINGO header reparam = eitherResult . (`feed` "") . parse (parseBingo reparam $ splitHeader header) . putEOL . B.strip +parseSR TIR header reparam = eitherResult . (`feed` "") . parse (parseTIR reparam $ splitHeader header) . putEOL . B.strip +parseSR OPERON header reparam = eitherResult . (`feed` "") . parse (parseOperon reparam $ splitHeader header) . putEOL . B.strip +parseSR GOMEA header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip +parseSR SBP header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip +parseSR EPLEX header reparam = eitherResult . (`feed` "") . parse (parseGOMEA reparam $ splitHeader header) . putEOL . B.strip +parseSR PYSR header reparam = eitherResult . (`feed` "") . parse (parsePySR reparam $ splitHeader header) . putEOL . B.strip + +eitherResult' :: Show r => Result r -> Either String r +eitherResult' res = trace (show res) $ eitherResult res + +-- * Parsers + +-- | Creates a parser for a binary operator +binary :: B.ByteString -> (a -> a -> a) -> Assoc -> Operator B.ByteString a +binary name fun = Infix (do{ string (B.cons ' ' (B.snoc name ' ')) <|> string name; pure fun }) + +-- | Creates a parser for a unary function +prefix :: B.ByteString -> (a -> a) -> Operator B.ByteString a +prefix name fun = Prefix (do{ string name; pure fun }) + +-- | Envelopes the parser in parens +parens :: Parser a -> Parser a +parens e = do{ string "("; e' <- e; string ")"; pure e' } "parens" + +-- | Parse an expression using a user-defined parser given by the `Operator` lists containing +-- the name of the functions and operators of that SR algorithm, a list of parsers `binFuns` for binary functions +-- a parser `var` for variables, a boolean indicating whether to change floating point values to free +-- parameters variables, and a list of variable names with their corresponding indexes. +parseExpr :: [[Operator B.ByteString (Fix SRTree)]] -> [ParseTree -> ParseTree] -> ParseTree -> Bool -> [(B.ByteString, Int)] -> ParseTree +parseExpr table binFuns var reparam header = do e <- relabelParams <$> expr + many1' space + pure e + where + term = parens expr <|> enclosedAbs expr <|> choice (map ($ expr) binFuns) <|> coef <|> varC "term" + expr = buildExpressionParser table term + coef = if reparam + then do eNumber <- intOrDouble + case eNumber of + Left x -> pure $ fromIntegral x + Right _ -> pure $ param 0 + else Fix . Const <$> signed double "const" + varC = if null header + then var + else var <|> varHeader + + varHeader = choice $ map (uncurry getParserVar) $ sortOn (negate . B.length . fst) header + getParserVar k v = (string k <|> enveloped k) >> pure (Fix $ Var v) + enveloped s = (char ' ' <|> char '(') >> string s >> (char ' ' <|> char ')') >> pure "" + +enumerate :: [a] -> [(a, Int)] +enumerate = (`zip` [0..]) + +splitHeader :: B.ByteString -> [(B.ByteString, Int)] +splitHeader = enumerate . B.split ',' + +-- | Tries to parse as an `Int`, if it fails, +-- parse as a Double. +intOrDouble :: Parser (Either Int Double) +intOrDouble = eitherP parseInt (signed double) + where + parseInt :: Parser Int + parseInt = do x <- signed decimal + c <- peekChar + case c of + Just '.' -> digit >> pure 0 + Just 'e' -> digit >> pure 0 + Just 'E' -> digit >> pure 0 + _ -> pure x + +putEOL :: B.ByteString -> B.ByteString +putEOL bs | B.last bs == '\n' = bs + | otherwise = B.snoc bs '\n' + +-- * Special case functions + +-- | analytic quotient +aq :: Fix SRTree -> Fix SRTree -> Fix SRTree +aq x y = x / sqrt (1 + y ** 2) + +log1p :: Fix SRTree -> Fix SRTree +log1p x = log (1 + x) + +log10 :: Fix SRTree -> Fix SRTree +log10 x = log x / log 10 + +log2 :: Fix SRTree -> Fix SRTree +log2 x = log x / log 2 + +cbrt :: Fix SRTree -> Fix SRTree +cbrt x = x ** (1/3) + +-- Parse `abs` functions as | x | +enclosedAbs :: Num a => Parser a -> Parser a +enclosedAbs expr = do char '|' + e <- expr + char '|' + pure $ abs e + +-- | Parser for binary functions +binFun :: B.ByteString -> (a -> a -> a) -> Parser a -> Parser a +binFun name f expr = do string name + many' space >> char '(' >> many' space + e1 <- expr + many' space >> char ',' >> many' space -- many' space >> char ',' >> many' space + e2 <- expr + many' space >> char ')' + pure $ f e1 e2 + +-- * Custom parsers for SR algorithms + +-- | parser for Transformation-Interaction-Rational. +parseTIR :: Bool -> [(B.ByteString, Int)] -> ParseTree +parseTIR = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [ ] + prefixOps = map (uncurry prefix) + [ ("id", id), ("abs", abs) + , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh) + , ("sin", sin), ("cos", cos), ("tan", tan) + , ("asinh", asinh), ("acosh", acosh), ("atanh", atanh) + , ("asin", asin), ("acos", acos), ("atan", atan) + , ("sqrt", sqrt), ("cbrt", cbrt), ("square", (**2)) + , ("log", log), ("exp", exp) + , ("Id", id), ("Abs", abs) + , ("Sinh", sinh), ("Cosh", cosh), ("Tanh", tanh) + , ("Sin", sin), ("Cos", cos), ("Tan", tan) + , ("ASinh", asinh), ("ACosh", acosh), ("ATanh", atanh) + , ("ASin", asin), ("ACos", acos), ("ATan", atan) + , ("Sqrt", sqrt), ("Cbrt", cbrt), ("Square", (**2)) + , ("Log", log), ("Exp", exp) + ] + binOps = [[binary "^" (**) AssocLeft], [binary "**" (**) AssocLeft] + , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do char 'x' + ix <- decimal + pure $ Fix $ Var ix + "var" + +-- | parser for Operon. +parseOperon :: Bool -> [(B.ByteString, Int)] -> ParseTree +parseOperon = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [ binFun "pow" (**) ] + prefixOps = map (uncurry prefix) + [ ("abs", abs), ("cbrt", cbrt) + , ("acos", acos), ("cosh", cosh), ("cos", cos) + , ("asin", asin), ("sinh", sinh), ("sin", sin) + , ("exp", exp), ("log", log) + , ("sqrt", sqrt), ("square", (**2)) + , ("atan", atan), ("tanh", tanh), ("tan", tan) + ] + binOps = [[binary "^" (**) AssocLeft] + , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do char 'X' + ix <- decimal + pure $ Fix $ Var (ix - 1) -- Operon is not 0-based + "var" + +-- | parser for HeuristicLab. +parseHL :: Bool -> [(B.ByteString, Int)] -> ParseTree +parseHL = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [ binFun "aq" aq ] + prefixOps = map (uncurry prefix) + [ ("logabs", log.abs), ("sqrtabs", sqrt.abs) -- the longer versions should come first + , ("abs", abs), ("exp", exp), ("log", log) + , ("sqrt", sqrt), ("sqr", (**2)), ("cube", (**3)) + , ("cbrt", cbrt), ("sin", sin), ("cos", cos) + , ("tan", tan), ("tanh", tanh) + ] + binOps = [[binary "^" (**) AssocLeft] + , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do char 'x' + ix <- decimal + pure $ Fix $ Var ix + "var" + +-- | parser for Bingo +parseBingo :: Bool -> [(B.ByteString, Int)] -> ParseTree +parseBingo = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [] + prefixOps = map (uncurry prefix) + [ ("abs", abs), ("exp", exp), ("log", log.abs) + , ("sqrt", sqrt.abs) + , ("sinh", sinh), ("cosh", cosh) + , ("sin", sin), ("cos", cos) + ] + binOps = [[binary "^" (**) AssocLeft] + , [binary "/" (/) AssocLeft, binary "" (*) AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do string "X_" + ix <- decimal + pure $ Fix $ Var ix + "var" + +-- | parser for GOMEA +parseGOMEA :: Bool -> [(B.ByteString, Int)] -> ParseTree +parseGOMEA = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [] + prefixOps = map (uncurry prefix) + [ ("exp", exp), ("plog", log.abs) + , ("sqrt", sqrt.abs) + , ("sin", sin), ("cos", cos) + ] + binOps = [[binary "^" (**) AssocLeft] + , [binary "/" (/) AssocLeft, binary "*" (*) AssocLeft, binary "aq" aq AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do string "x" + ix <- decimal + pure $ Fix $ Var ix + "var" + +-- | parser for PySR +parsePySR :: Bool -> [(B.ByteString, Int)] -> ParseTree +parsePySR = parseExpr (prefixOps : binOps) binFuns var + where + binFuns = [ binFun "pow" (**) ] + prefixOps = map (uncurry prefix) + [ ("abs", abs), ("exp", exp) + , ("square", (**2)), ("cube", (**3)), ("neg", negate) + , ("acosh_abs", acosh . (+1) . abs), ("acosh", acosh), ("asinh", asinh) + , ("acos", acos), ("asin", asin), ("atan", atan) + , ("sqrt_abs", sqrt.abs), ("sqrt", sqrt) + , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh) + , ("sin", sin), ("cos", cos), ("tan", tan) + , ("log10", log10), ("log2", log2), ("log1p", log1p) + , ("log_abs", log.abs), ("log10_abs", log10 . abs) + , ("log", log) + ] + binOps = [[binary "^" (**) AssocLeft] + , [binary "/" (/) AssocLeft, binary "*" (*) AssocLeft] + , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft] + ] + var = do string "x" + ix <- decimal + pure $ Fix $ Var ix + "var" diff --git a/src/Text/ParseSR/IO.hs b/src/Text/ParseSR/IO.hs new file mode 100644 --- /dev/null +++ b/src/Text/ParseSR/IO.hs @@ -0,0 +1,73 @@ +{-# language LambdaCase #-} +----------------------------------------------------------------------------- +-- | +-- Module : Text.ParseSR.IO +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to parse multiple expressions from stdin or a text file. +-- +----------------------------------------------------------------------------- +module Text.ParseSR.IO ( withInput, withOutput, withOutputDebug ) + where + +-- import Data.SRTree.EqSat1 +import Algorithm.EqSat.Simplify ( simplifyEqSatDefault ) +import Control.Monad (forM_, unless) +import qualified Data.ByteString.Char8 as B +import Data.SRTree +import Data.SRTree.Recursion (Fix (..)) +import System.IO +import Text.ParseSR (Output, SRAlgs, parseSR, showOutput) + +-- | given a filename, the symbolic regression algorithm, a string of variables name, +-- and two booleans indicating whether to convert float values to parameters and +-- whether to simplify the expression or not, it will read the file and parse everything +-- returning a list of either an error message or a tree. +-- +-- empty filename defaults to stdin +withInput :: String -> SRAlgs -> String -> Bool -> Bool -> IO [Either String (Fix SRTree)] +withInput fname sr hd param simpl = do + h <- if null fname then pure stdin else openFile fname ReadMode + 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 + es = map myParser $ filter (not . null) contents + unless (null fname) $ hClose h + pure es + +-- | outputs a list of either error or trees to a file using the Output format. +-- +-- empty filename defaults to stdout +withOutput :: String -> Output -> [Either String (Fix SRTree)] -> IO () +withOutput fname output exprs = do + h <- if null fname then pure stdout else openFile fname WriteMode + forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err + Right ex -> hPutStrLn h (showOutput output ex) + unless (null fname) $ hClose h + +-- | debug version of output function to check the invalid parsers +withOutputDebug :: String -> Output -> [Either String (Fix SRTree, Fix SRTree)] -> IO () +withOutputDebug fname output exprs = do + h <- if null fname then pure stdout else openFile fname WriteMode + forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err + Right (t1, t2) -> do + hPutStrLn h ("First: " <> showOutput output t1) + hPutStrLn h ("Second: " <> showOutput output t2) + hFlush h + unless (null fname) $ hClose h + +hGetLines :: Handle -> IO [String] +hGetLines h = do + done <- hIsEOF h + if done + then return [] + else do + line <- hGetLine h + (line :) <$> hGetLines h diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -1,12 +1,12 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.35.2. +-- This file has been generated from package.yaml by hpack version 0.37.0. -- -- see: https://github.com/sol/hpack name: srtree -version: 1.0.0.5 -synopsis: A general framework to work with Symbolic Regression expression trees. +version: 2.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; category: Math, Data, Data Structures homepage: https://github.com/folivetti/srtree#readme @@ -27,24 +27,270 @@ library exposed-modules: + Algorithm.EqSat + Algorithm.EqSat.Build + Algorithm.EqSat.DB + Algorithm.EqSat.Egraph + Algorithm.EqSat.Info + Algorithm.EqSat.Queries + 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 + Text.ParseSR + Text.ParseSR.IO other-modules: Paths_srtree hs-source-dirs: src + ghc-options: -fwarn-incomplete-patterns build-depends: - base >=4.16 && <5 - , containers ==0.6.* + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 , random ==1.2.* + , split >=0.2.5 && <0.3 + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 default-language: Haskell2010 +executable egraphGP + main-is: Main.hs + other-modules: + Random + Paths_srtree + hs-source-dirs: + apps/egraphGP + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , optparse-applicative >=0.17 && <0.19 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + +executable eqsatrepr + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/eqsatrepr + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + +executable ieeexplore + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/ieeexplore + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , optparse-applicative >=0.17 && <0.19 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + 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 -O2 -optc-O3 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , optparse-applicative >=0.17 && <0.19 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , 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 -O2 -optc-O3 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , normaldistribution >=1.1.0.3 && <1.2 + , optparse-applicative >=0.17 && <0.19 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + +executable tinygp + main-is: Main.hs + other-modules: + GP + Initialization + Paths_srtree + hs-source-dirs: + apps/tinygp + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -optc-O3 + build-depends: + attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.16 && <5 + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 + , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 + , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 + , optparse-applicative >=0.17 && <0.19 + , random ==1.2.* + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* + , 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 @@ -56,11 +302,27 @@ build-depends: HUnit , ad + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 , base >=4.16 && <5 - , containers ==0.6.* + , bytestring ==0.11.* + , containers >=0.6.7 && <0.8 , dlist ==1.0.* + , exceptions >=0.10.7 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4.4.0 && <1.6 + , ieee754 >=0.8.0 && <0.9 + , lens >=5.2.3 && <5.4 + , list-shuffle >=1.0.0.1 && <1.1 + , massiv >=1.0.4.0 && <1.1 , mtl >=2.2 && <2.4 + , nlopt-haskell >=0.1.3.0 && <0.2 , random ==1.2.* + , split >=0.2.5 && <0.3 , srtree + , statistics >=0.16.2.1 && <0.17 + , transformers >=0.6.1.0 && <0.7 + , unordered-containers ==0.2.* , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 default-language: Haskell2010 diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,8 +1,15 @@ import Data.SRTree +import Data.SRTree.Eval +import Data.SRTree.Derivative +import Data.SRTree.Datasets +import Algorithm.SRTree.AD import qualified Data.Vector as V import Numeric.AD.Double ( grad ) import Test.HUnit +import qualified Data.Massiv.Array as M +import Data.Massiv.Array (D, S, Ix1, Ix2, Comp(..), Sz(..)) +import qualified Foreign as M -- test expressions exprs = [ @@ -13,6 +20,8 @@ , 1 / param 0 * param 1 , param 0 + param 1 + param 0 * param 1 + sin (param 0) + sin (param 1) + cos (param 0) + cos (param 1) + sin (param 0 * param 1) + cos (param 0 * param 1) , sin (exp (param 0) + param 1) + , param 0 / param 1 + , param 0 ** param 1 ] -- autodiff with multiple occurrences of vars @@ -24,6 +33,8 @@ , grad (\[x,y] -> 1 / x * y) [2,3] , grad (\[x,y] -> x + y + x * y + sin x + sin y + cos x + cos y + sin (x * y) + cos (x * y)) [2,3] , grad (\[x,y] -> sin (exp x + y)) [2,3] + , grad (\[x,y] -> x/y) [2,3] + , grad (\[x,y] -> x ** y) [2,3] ] -- autodiff with single occurrences of vars @@ -35,30 +46,39 @@ , grad (\[x,y] -> 1 / x * y) [2,3] , grad (\[a,b,c,d,e,f,g,h,i,j,k,l] -> a + b + c * d + sin e + sin f + cos g + cos h + sin (i * j) + cos (k * l)) [2,3,2,3,2,3,2,3,2,3,2,3] , grad (\[x,y] -> sin (exp x + y)) [2,3] + , grad (\[x,y] -> x/y) [2,3] + , grad (\[x,y] -> x ** y) [2,3] ] -- xs is empty since we are interested in theta -xs :: V.Vector a -xs = V.empty +xs :: M.Array S Ix2 Double +xs = M.singleton 0 + +xs' :: M.Array S Ix2 Double +xs' = M.singleton 0 + +err = M.singleton 1 + -- theta values -thetaMulti, thetaSingle :: V.Vector Double -thetaMulti = V.fromList [2.0, 3.0] -thetaSingle = V.fromList [2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0] +thetaMulti, thetaSingle :: M.Array S Ix1 Double +thetaMulti = M.fromList Seq [2.0, 3.0] +thetaSingle = M.fromList Seq [2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0] -- values from forward mode -forwardVals :: [[Double]] -forwardVals = map (forwardMode xs thetaMulti id) exprs +-- forwardVals :: [[Double]] +forwardVals = map (M.toList . snd . forwardMode xs' thetaMulti err) exprs -- values from grad -- we must relabel the parameters of the expression to sequence values -gradVals :: [(Double, [Double])] -gradVals = map (gradParamsFwd xs thetaSingle id . relabelParams) exprs +--gradVals :: [(Double, [Double])] +gradVals = map (M.toList . snd . forwardModeUnique xs' thetaSingle err . relabelParams) exprs +gradVals' = map (M.toList . snd . reverseModeUnique xs' thetaSingle err . relabelParams) exprs -- values of the evaluated expressions -exprVals :: [Double] -exprVals = map (evalTree xs thetaSingle id . relabelParams) exprs +--exprVals :: [Double] +exprVals = map (evalTree xs' thetaSingle . relabelParams) exprs -refGrad :: [(Double, [Double])] +--refGrad :: [(Double, [Double])] refGrad = zip exprVals autoDiffSingle testDiff :: (Eq a, Show a) => String -> String -> a -> a -> Test @@ -67,10 +87,12 @@ tests :: Test tests = TestList $ zipWith (testDiff "forward mode" "autodiff x forward mode") autoDiffMult forwardVals - <> zipWith (testDiff "opt. grad. parameters" "(evalTree, autodiff) x gradVals") refGrad gradVals - <> zipWith (testDiff "deriveByParam" "deriveByParam x autodiff") (map head autoDiffSingle) (map (head.snd) gradVals) + <> zipWith (testDiff "forward mode" "autodiff x forward mode unique") autoDiffSingle gradVals + <> zipWith (testDiff "reverse mode" "autodiff x reverse mode unique") autoDiffSingle gradVals' main :: IO () main = do result <- runTestTT tests putStrLn $ showCounts result + --ds <- loadDataset "test/wine.csv:3:10:alcohol:liver,deaths,heart" True + --print ds