srtree 1.0.0.5 → 2.0.0.0

raw patch · 40 files changed

+8337/−398 lines, 40 filesdep +attoparsecdep +attoparsec-exprdep +bytestringdep ~containersnew-component:exe:egraphGPnew-component:exe:eqsatreprnew-component:exe:ieeexplorenew-component:exe:srsimplifynew-component:exe:srtoolsnew-component:exe:tinygpPVP ok

version bump matches the API change (PVP)

Dependencies added: attoparsec, attoparsec-expr, bytestring, exceptions, filepath, hashable, ieee754, lens, list-shuffle, massiv, nlopt-haskell, normaldistribution, optparse-applicative, split, statistics, transformers, unordered-containers, zlib

Dependency ranges changed: containers

API changes (from Hackage documentation)

Files

• ChangeLog.md +14/−0
• README.md +262/−13
• apps/egraphGP/Main.hs +598/−0
• apps/egraphGP/Random.hs +54/−0
• apps/eqsatrepr/Main.hs +302/−0
• apps/ieeexplore/Main.hs +101/−0
• apps/srsimplify/Main.hs +103/−0
• apps/srtools/Args.hs +178/−0
• apps/srtools/IO.hs +181/−0
• apps/srtools/Main.hs +31/−0
• apps/srtools/Report.hs +271/−0
• apps/tinygp/GP.hs +222/−0
• apps/tinygp/Initialization.hs +50/−0
• apps/tinygp/Main.hs +72/−0
• src/Algorithm/EqSat.hs +173/−0
• src/Algorithm/EqSat/Build.hs +460/−0
• src/Algorithm/EqSat/DB.hs +351/−0
• src/Algorithm/EqSat/Egraph.hs +270/−0
• src/Algorithm/EqSat/Info.hs +168/−0
• src/Algorithm/EqSat/Queries.hs +92/−0
• src/Algorithm/EqSat/Simplify.hs +214/−0
• src/Algorithm/Massiv/Utils.hs +278/−0
• src/Algorithm/SRTree/AD.hs +323/−0
• src/Algorithm/SRTree/ConfidenceIntervals.hs +447/−0
• src/Algorithm/SRTree/Likelihoods.hs +260/−0
• src/Algorithm/SRTree/ModelSelection.hs +176/−0
• src/Algorithm/SRTree/NonlinearOpt.hs +974/−0
• src/Algorithm/SRTree/Opt.hs +109/−0
• src/Data/SRTree.hs +17/−23
• src/Data/SRTree/Datasets.hs +214/−0
• src/Data/SRTree/Derivative.hs +125/−0
• src/Data/SRTree/Eval.hs +210/−0
• src/Data/SRTree/Internal.hs +229/−296
• src/Data/SRTree/Print.hs +97/−34
• src/Data/SRTree/Random.hs +20/−12
• src/Data/SRTree/Recursion.hs +12/−0
• src/Text/ParseSR.hs +302/−0
• src/Text/ParseSR/IO.hs +73/−0
• srtree.cabal +268/−6
• test/Spec.hs +36/−14

@@ -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

@@ -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 <error message>".++### `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: 

@@ -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<nEvs)++    updateIfNothing ec = do+      mf <- getFitness ec+      case mf of+        Nothing -> do+          t <- getBest ec+          (f, p) <- fitnessFunRep x y x_val y_val t+          insertFitness ec f p+          pure True+        Just _ -> pure False++    getBestFitness = do+      bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+      gets (_fitness . _info . (IM.! bec) . _eClass)++    evalRndSubTree :: RndEGraph (Maybe EClassId)+    evalRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+                        if not (null ecIds)+                          then do rndId <- rnd $ randomFrom ecIds+                                  Just <$> canonical rndId+                          else pure Nothing+++    combineFrom ecs = do+        nt  <- rnd rndNonTerm+        p1  <- rnd (randomFrom ecs)+        p2  <- rnd (randomFrom ecs)+        l1  <- rnd (randomFrom [1..maxSize-2])+        l2  <- rnd (randomFrom [1..(maxSize - l1 - 1)])+        e1  <- randomChildFrom p1 l1+        ml  <- gets (_size . _info . (IM.! e1) . _eClass)+        e2  <- randomChildFrom p2 l2+        case nt of+          Uni Id ()    -> canonical e1+          Uni f ()     -> add myCost (Uni f e1) >>= canonical+          Bin op () () -> do b <- rnd toss+                             if b+                              then add myCost (Bin op e1 e2) >>= canonical+                              else add myCost (Bin op e2 e1) >>= canonical++    getParetoEcsUpTo n = concat <$> (forM [1..maxSize] $ \i -> getBestEcsOfSize  i n)++    getBestEcsOfSize i n = 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." )

@@ -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 #-}

@@ -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]

@@ -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." )

@@ -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." )

@@ -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

@@ -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 "============================================================="

@@ -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"+            )

@@ -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

@@ -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'

@@ -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+    -}

@@ -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." )

@@ -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

@@ -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 ()

@@ -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 <> " | <cond>"++-- 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)

@@ -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 #-}

@@ -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))

@@ -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)

@@ -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

@@ -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: <http://en.wikipedia.org/wiki/Spline_interpolation#Interpolation_using_natural_cubic_spline>.+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

@@ -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 ><! jacob)+  where +    (Sz p)               = M.size theta+    (Sz (m :. n))        = M.size xss+    cmp                  = getComp xss+    -- | if the tree does not use a variable +    -- it will return a single scalar, fromEither fixes this+    fromEither (Left y)  = y+    fromEither (Right y) = M.replicate cmp (Sz m) y++    -- if it is a variable, returns the value of that variable and an array of zeros (Jacobian)+    alg (Var ix) = do tape  <- M.newMArray (Sz2 m p) 0 +                                 >>= UMA.unsafeFreeze cmp+                      pure (Left (xss <! ix), tape)++    -- if it is a constant, returns the value of the constant and array of zeros +    alg (Const c) = do tape <- M.newMArray (Sz2 m p) 0+                                 >>= UMA.unsafeFreeze cmp+                       pure (Right c, tape)++    -- 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 <! ix, DL.empty)+      alg (Param ix)      = (replicateAs xss $ theta ! ix, DL.singleton one)+      alg (Const c)       = (replicateAs xss c, DL.empty)+      alg (Uni f (v, gs)) = let v' = evalFun f v+                                dv = derivative f v+                             in (v', DL.map (*dv) gs)+      alg (Bin Add (v1, l) (v2, r)) = (v1+v2, DL.append l r)+      alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, DL.append l (DL.map negate r))+      alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, DL.append (DL.map (*v2) l) (DL.map (*v1) r))+      alg (Bin Div (v1, l) (v2, r)) = let dv = ((-v1)/(v2*v2)) +                                       in (v1/v2, DL.append (DL.map (/v2) l) (DL.map (*dv) r))+      alg (Bin Power (v1, l) (v2, r)) = let dv1 = v1 ** (v2 - one)+                                            dv2 = v1 * log v1+                                         in (v1 ** v2, DL.map (*dv1) (DL.append (DL.map (*v2) l) (DL.map (*dv2) r)))+      alg (Bin PowerAbs (v1, l) (v2, r)) = let dv1 = abs v1 ** v2+                                               dv2 = DL.map (* (log (abs v1))) r+                                               dv3 = DL.map (*(v2 / v1)) l+                                           in (abs v1 ** v2, DL.map (*dv1) (DL.append dv2 dv3))+      alg (Bin AQ (v1, l) (v2, r)) = let dv1 = DL.map (*(1 + v2*v2)) l+                                         dv2 = DL.map (*(-v1*v2)) r+                                     in (v1/sqrt(1 + v2*v2), DL.map (/(1 + v2*v2)**1.5) $ DL.append dv1 dv2)++data TupleF a b = Single a | T a b | Branch a b b deriving Functor -- hi, I'm a tree+type Tuple a = Fix (TupleF a)++-- | Same as above, but using reverse mode, that is even faster.+reverseModeUnique :: SRMatrix+                  -> 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 <! ix)+      forward (Param ix)   = oneTpl (Right $ theta ! ix)+      forward (Const c)    = oneTpl (Right c)+      forward (Uni g t)    = let v = getTop t+                             in tuple (applyUni g v) t+      forward (Bin op l r) = let vl = getTop l+                                 vr = getTop r+                              in branch (applyBin op vl vr) l r++++      -- reverse walks from the root to the leaf calculating the+      -- partial derivative with respect to an arbitrary variable+      -- up to that point+      reverse (Var ix)     (dx,   _)         = Var ix+      reverse (Param ix)   (dx,   _)         = Param ix+      reverse (Const v)    (dx,   _)         = Const v+      reverse (Uni f t)    (dx, unCons -> v) =+          let g' = applyDer f (getTop v)+          in Uni f (t, ( applyBin Mul dx g', v ))+      reverse (Bin op l r) (dx, getBranches -> (vl, vr)) =+          let (dxl, dxr) = diff op dx (getTop vl) (getTop vr)+           in Bin op (l, (dxl, vl)) (r, (dxr, vr))++      -- dx is the current derivative so far+      -- fx is the evaluation of the left branch+      -- gx is the evaluation of the right branch+      --+      -- this should return a tuple, where the left element is+      -- dx * d op(f(x), g(x)) / d f(x) and+      -- the right branch dx * d op (f(x), g(x)) / d g(x)+      diff Add dx fx gy = (dx, dx)+      diff Sub dx fx gy = (dx, negate' dx)+      diff Mul dx fx gy = (applyBin Mul dx gy, applyBin Mul dx fx)+      diff Div dx fx gy = (applyBin Div dx gy, applyBin Mul dx (applyBin Div (negate' fx) (applyBin Mul gy gy)))+      diff Power dx fx gy = let dxl = applyBin Mul dx (applyBin Power fx (applyBin Sub gy (Right 1)))+                                dv2 = applyBin Mul fx (applyUni Log fx)+                            in (applyBin Mul dxl gy, applyBin Mul dxl dv2)+      diff PowerAbs dx fx gy = let dxl = applyBin Mul (applyBin Mul gy fx) (applyBin PowerAbs fx (applyBin Sub gy (Right 2)))+                                   dxr = applyBin Mul (applyUni LogAbs fx) (applyBin PowerAbs fx gy)+                               in (applyBin Mul dxl dx, applyBin Mul dxr dx)+      diff AQ dx fx gy = let dxl = applyUni Recip (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1)))+                             dxy = applyBin Div (applyBin Mul fx gy) (applyUni Cube (applyUni Sqrt (applyBin Add (applyUni Square gy) (Right 1))))+                         in (applyBin Mul dxl dx, applyBin Mul dxy dx)+++      -- once we reach a leaf with a parameter, we return a singleton+      -- with that derivative upwards until the root+      --combine :: (forall s . MArray (PrimState (ST s)) S Int Double) -> SRTree () -> (Either SRVector Double, a) -> ()+      combine j (Var ix) s = 0+      combine j (Const _) s = 0+      combine j (Param ix) s = unsafePerformIO $ do+                                 case fst s of+                                   Left v  -> do v' <- dotM v err+                                                 UMA.unsafeWrite j ix v'+                                   Right v -> UMA.unsafeWrite j ix $ M.foldrS (\x acc -> x*v + acc) 0 err+                                 UMA.unsafeRead j ix+      combine j (Uni f gs) s = gs+      combine j (Bin op l r) s = l+r+++-- | 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 <! ix, DL.empty)+      alg (Param ix)      = (replicateAs xss $ theta ! ix, DL.singleton one)+      alg (Const c)       = (replicateAs xss c, DL.empty)+      alg (Uni f (v, gs)) = let v' = evalFun f v+                                dv = derivative f v+                             in (v', DL.map (*dv) gs)+      alg (Bin Add (v1, l) (v2, r)) = (v1+v2, DL.append l r)+      alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, DL.append l (DL.map negate r))+      alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, DL.append (DL.map (*v2) l) (DL.map (*v1) r))+      alg (Bin Div (v1, l) (v2, r)) = let dv = ((-v1)/(v2*v2))+                                       in (v1/v2, DL.append (DL.map (/v2) l) (DL.map (*dv) r))+      alg (Bin Power (v1, l) (v2, r)) = let dv1 = v1 ** (v2 - one)+                                            dv2 = v1 * log v1+                                         in (v1 ** v2, DL.map (*dv1) (DL.append (DL.map (*v2) l) (DL.map (*dv2) r)))+      alg (Bin PowerAbs (v1, l) (v2, r)) = let dv1 = abs v1 ** v2+                                               dv2 = DL.map (* (log (abs v1))) r+                                               dv3 = DL.map (*(v2 / v1)) l+                                           in (abs v1 ** v2, DL.map (*dv1) (DL.append dv2 dv3))+      alg (Bin AQ (v1, l) (v2, r)) = let dv1 = DL.map (*(1 + v2*v2)) l+                                         dv2 = DL.map (*(-v1*v2)) r+                                     in (v1/sqrt(1 + v2*v2), DL.map (/(1 + v2*v2)**1.5) $ DL.append dv1 dv2)

@@ -0,0 +1,447 @@+{-# language ViewPatterns, ScopedTypeVariables, MultiWayIf, FlexibleContexts #-}+-----------------------------------------------------------------------------+-- |+-- Module      :  Algorithm.SRTree.ConfidenceIntervals +-- 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.ConfidenceIntervals where++import qualified Data.Massiv.Array as A+import Data.Massiv.Array (Ix2(..), (*.), (!+!), (!*!))+import Data.Massiv.Array.Numeric ( identityMatrix )+import Statistics.Distribution ( ContDistr(quantile) )+import Statistics.Distribution.StudentT ( studentT )+import Statistics.Distribution.FDistribution ( fDistribution )+import qualified Data.Vector.Storable as VS+import Data.SRTree+import Data.SRTree.Eval+import Data.SRTree.Recursion ( cata )+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.Opt+    ( minimizeNLLNonUnique, minimizeNLLWithFixedParam )+import Data.List ( sortOn, nubBy )+import Data.Maybe ( fromMaybe )+import Algorithm.SRTree.NonlinearOpt+import Algorithm.Massiv.Utils+import System.IO.Unsafe ( unsafePerformIO )+import Control.Monad.Catch ( catch )++import Debug.Trace ( trace, traceShow )++-- | profile likelihood algorithms: Bates (classical), ODE (faster), Constrained (fastest)+-- The Constrained approach returns only the endpoints.+data PType = Bates | ODE | Constrained deriving (Show, Read)++-- | Confidence Interval using Laplace approximation or profile likelihood.+data CIType = Laplace BasicStats | Profile BasicStats [ProfileT]++-- | Basic stats of the data: covariance of parameters, correlation, standard errors +data BasicStats = MkStats { _cov    :: SRMatrix+                          , _corr   :: SRMatrix+                          , _stdErr :: PVector+                          } deriving (Eq, Show)++-- | a confience interval is composed of the point estimate (`est_`), lower bound (`_lower_`)+-- and upper bound (`upper_`)+data CI = CI { est_   :: Double+             , lower_ :: Double+             , upper_ :: Double+             } deriving (Eq, Show, Read)++--  | A profile likelihood is composed of a vector of tau values that traces the likelihood, +--  the matrix of thetas for each profile, the local optima, and two splines that converts +--  taus to theta and vice-versa. +data ProfileT = ProfileT { _taus      :: PVector+                         , _thetas    :: SRMatrix+                         , _opt       :: Double+                         , _tau2theta :: Double -> Double+                         , _theta2tau :: Double -> Double+                         }++-- 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<p && j<p -> 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

@@ -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)+

@@ -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 #-}

@@ -0,0 +1,974 @@+{-# OPTIONS_GHC -Wall #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TypeApplications #-}++{- |+Module      :  Numeric.NLOPT+Copyright   :  (c) Matthew Peddie 2017+License     :  BSD3+Maintainer  :  Matthew Peddie <mpeddie@gmail.com>+Stability   :  provisional+Portability :  GHC++This module provides a high-level, @hmatrix@-compatible interface to+the <http://ab-initio.mit.edu/wiki/index.php/NLopt NLOPT> 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 <http://ab-initio.mit.edu/wiki/index.php/NLopt_Introduction NLOPT introduction>++  * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Reference NLOPT reference manual>++  * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms NLOPT algorithm manual>++= 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+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- 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+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- 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+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- 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

@@ -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

@@ -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

@@ -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),+    (<!),+  )+import Data.Massiv.Array qualified as M+import Data.Maybe (fromJust)+import Data.SRTree.Eval (PVector, SRMatrix, compMode)+import Data.Vector qualified as V+import System.FilePath (takeExtension)+import Text.Read (readMaybe)++-- | Loads a list of list of bytestrings to a matrix of double+loadMtx :: [[B.ByteString]] -> 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 <!) ixs+    y       = datum <! iy+    x_train = M.computeAs S $ M.extractFromTo' (st :. 0) (end+1 :. p) x+    y_train = M.computeAs S $ M.extractFromTo' st (end+1) y +    x_val   = M.computeAs S $ M.throwEither $ M.deleteRowsM st (Sz1 $ end - st + 1) x+    y_val   = M.computeAs S $ M.throwEither $ M.deleteColumnsM st (Sz1 $ end - st + 1) y+{-# inline processData #-}+

@@ -0,0 +1,125 @@+{-# LANGUAGE OverloadedStrings #-}+-----------------------------------------------------------------------------+-- |+-- Module      :  Data.SRTree.Derivative +-- Copyright   :  (c) Fabricio Olivetti 2021 - 2024+-- License     :  BSD3+-- Maintainer  :  fabricio.olivetti@gmail.com+-- Stability   :  experimental+-- Portability :  FlexibleInstances, DeriveFunctor, ScopedTypeVariables+--+-- Symbolic derivative of SRTree expressions+--+-----------------------------------------------------------------------------+module Data.SRTree.Derivative+        ( derivative+        , doubleDerivative+        , deriveByVar+        , deriveByParam+        )+        where++import Data.SRTree.Internal+import Data.SRTree.Recursion (Fix (..), mutu)+import Data.Attoparsec.ByteString.Char8 (double)++-- | Creates the symbolic partial derivative of a tree by variable `dx` (if `p` is `False`)+-- or parameter `dx` (if `p` is `True`).+-- This uses mutual recursion where the first recursion (alg1) holds the derivative w.r.t. +-- the current node and the second (alg2) holds the original tree.+--+-- >>> 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 #-}