srtree 1.0.0.5 → 2.0.1.8
raw patch · 40 files changed
Files
- ChangeLog.md +69/−0
- README.md +269/−14
- apps/srsimplify/Main.hs +103/−0
- apps/srtools/Args.hs +184/−0
- apps/srtools/IO.hs +230/−0
- apps/srtools/Main.hs +34/−0
- apps/srtools/Report.hs +280/−0
- apps/tinygp/GP.hs +252/−0
- apps/tinygp/Initialization.hs +50/−0
- apps/tinygp/Main.hs +116/−0
- apps/tinygp/Util.hs +63/−0
- src/Algorithm/EqSat.hs +186/−0
- src/Algorithm/EqSat/Build.hs +640/−0
- src/Algorithm/EqSat/DB.hs +373/−0
- src/Algorithm/EqSat/Egraph.hs +358/−0
- src/Algorithm/EqSat/Info.hs +208/−0
- src/Algorithm/EqSat/Queries.hs +221/−0
- src/Algorithm/EqSat/SearchSR.hs +223/−0
- src/Algorithm/EqSat/SearchSRCache.hs +244/−0
- src/Algorithm/EqSat/Simplify.hs +280/−0
- src/Algorithm/Massiv/Utils.hs +278/−0
- src/Algorithm/SRTree/AD.hs +559/−0
- src/Algorithm/SRTree/ConfidenceIntervals.hs +448/−0
- src/Algorithm/SRTree/Likelihoods.hs +635/−0
- src/Algorithm/SRTree/ModelSelection.hs +190/−0
- src/Algorithm/SRTree/NonlinearOpt.hs +976/−0
- src/Algorithm/SRTree/Opt.hs +135/−0
- src/Data/SRTree.hs +25/−23
- src/Data/SRTree/Datasets.hs +281/−0
- src/Data/SRTree/Derivative.hs +127/−0
- src/Data/SRTree/Eval.hs +215/−0
- src/Data/SRTree/Internal.hs +300/−293
- src/Data/SRTree/Print.hs +120/−35
- src/Data/SRTree/Random.hs +81/−27
- src/Data/SRTree/Recursion.hs +12/−0
- src/Numeric/Optimization/NLOPT/Bindings.hs +1075/−0
- src/Text/ParseSR.hs +402/−0
- src/Text/ParseSR/IO.hs +73/−0
- srtree.cabal +189/−9
- test/Spec.hs +1/−73
ChangeLog.md view
@@ -1,5 +1,74 @@ # Changelog for srtree +## 2.0.1.7 ++- Added log10 MSE fitness function ++## 2.0.1.6++- Added Fractional Bayes model selection++## 2.0.1.5++- Fix `refit` to only replace the fitness if it improves the fitness +- Fix `paretoFront` ++## 2.0.1.4++- Added `loadTrainingOnly`, `splitData`, and `loadX` to `Data.SRTree.Datasets`+- Added `getFitness`, `getTheta`, `getSize`, `isSizeOf`, `getBestFitness` to `Algorithm.EqSat.Egraph`+- Added `parseNonTerms` to `Text.ParseSR` +- Added module `Algorithm.EqSat.SearchSR` with support functions for SR algorithms with EqSat++## 2.0.1.3++- Fix compatibility with stackage nightly ++## 2.0.1.2++- Fix bug where the parameters were printed as `t[:,ix]` instead of `t[ix]` in `showPython`++## 2.0.1.1++- MSE loss is now the default+- Renamed `--distribution` argument to `--loss` in eggp and easter.+- Fixed bug with Gaussian distribution and fixed number of parameters.+- Fixed bug in which `--number-params 0` would create parameters.+- Fixed bug in `rEGGression` that pattern matched equivalent expressions.+- Support to `--numpy` flag that prints the output as a numpy expression (experimental, eggp only).+- Support to `--simplify` flag that simplifies the expressions before displaying (experimental, eggp only).++## 2.0.1.0++- Support to Multiview Symbolic Regression in eggp and symregg.+- Support to `--number-params` argument that limits the maximum number of parameters and allow repated parameters in an expression.++## 2.0.0.4++- Cleaned up test cases (they were deprecated), will include new ones later ++## 2.0.0.3++- Fixed compatibility with random-1.3.0 and GHC-9.12.1 +- Fixed bug in Bernoulli distribution +- Removed `log(sqrt(x))` rule in parametric rules due to generating longer expressions +- Fixed DL calculation without the correct number of parameters +- Fixed memory issue when querying pattern distribution ++## 2.0.0.0 ++- Complete refactoring of the library+- 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
README.md view
@@ -1,28 +1,283 @@-# 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** (SR). -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`.+This repository is also the home for different algorithm implementations for SR and software tools to support the post-processing of SR models (please refer to their corresponding README files): -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.+- [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.+- [rEGGression](https://github.com/folivetti/reggression/blob/main/README.md): nonlinear regression models exploration and query system with e-graphs (egg).+- [symregg](https://github.com/folivetti/symregg/blob/main/README.md): Equality graph Assisted Search Technique for Equation Recovery.+- [eggp](https://github.com/folivetti/eggp/blob/main/README.md): E-graph Genetic Programming. -The `SRTree` structure has instances for `Num, Fractional, Floating` which allows to create an expression as a valid Haskell expression such as:+## SRTree +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:++- `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, 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. +++## 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: - support more advanced functions - support conditional branching (`IF-THEN-ELSE`)-+- document egraph-search and ieeexplore
+ apps/srsimplify/Main.hs view
@@ -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." )
+ apps/srtools/Args.hs view
@@ -0,0 +1,184 @@+module Args where++import Data.Char ( toLower, toUpper )+import Data.List ( intercalate )+import Algorithm.SRTree.Likelihoods ( Distribution (..) )+import Algorithm.SRTree.ConfidenceIntervals ( PType (..) )+import Options.Applicative+import Text.ParseSR ( SRAlgs (..) )+import Text.Read ( readMaybe )++-- Data type to store command line arguments+data Args = Args+ { from :: SRAlgs+ , infile :: String+ , outfile :: String+ , dataset :: String+ , test :: String+ , niter :: Int+ , hasHeader :: Bool+ , simpl :: Bool+ , dist :: Distribution+ , restart :: Bool+ , rseed :: Int+ , toScreen :: Bool+ , useProfile :: Bool+ , simple :: Bool+ , sigma :: Double+ , alpha :: Double+ , ptype :: PType+ } deriving Show++-- parser of command line arguments+opt :: Parser Args+opt = Args+ <$> option sralgsReader+ ( long "from"+ <> short 'f'+ <> metavar ("[" <> intercalate "|" sralgsHelp <> "]")+ <> help "Input expression format" )+ <*> strOption+ ( long "input"+ <> short 'i'+ <> metavar "INPUT-FILE"+ <> showDefault+ <> value ""+ <> help "Input file containing expressions. \+ \ Empty string gets expression from stdin." )+ <*> strOption+ ( long "output"+ <> short 'o'+ <> metavar "OUTPUT-FILE"+ <> showDefault+ <> value ""+ <> help "Output file to store the stats in CSV format. \+ \ Empty string prints expressions to stdout." )+ <*> strOption+ ( long "dataset"+ <> short 'd'+ <> metavar "DATASET-FILENAME"+ <> help "Filename of the dataset used for optimizing the parameters. \+ \ Empty string omits stats that make use of the training data. \+ \ It will auto-detect and handle gzipped file based on gz extension. \+ \ It will also auto-detect the delimiter.\n\+ \ The filename can include extra information: \+ \ filename.csv:start:end:target:cols:yerr:xerr where start and end \+ \ corresponds to the range of rows that should be used for fitting,\+ \ target is the column index (or name) of the target variable and cols\+ \ is a comma separated list of column indices or names of the variables\+ \ in the same order as used by the symbolic model.\+ \ The yerr field corresponds to the column with the error of the target,\+ \ while xerr a comma separated indices of the columns with the error of the\+ \ variables. If nothing passed, it will ignore measurement errors." )+ <*> strOption+ ( long "test"+ <> metavar "TEST"+ <> showDefault+ <> value ""+ <> help "Filename of the test dataset.\+ \ Empty string omits stats that make use of the training data.\+ \ It can have additional information as in the training set,\+ \ but the validation range will be discarded." )+ <*> option auto+ ( long "niter"+ <> metavar "NITER"+ <> showDefault+ <> value 10+ <> help "Number of iterations for the optimization algorithm.")+ <*> switch+ ( long "hasheader"+ <> help "Uses the first row of the csv file as header.")+ <*> switch+ ( long "simplify"+ <> help "Apply basic simplification." )+ <*> option distRead+ ( long "distribution"+ <> metavar ("[" <> intercalate "|" distHelp <> "]")+ <> showDefault+ <> value Gaussian+ <> help "Minimize negative log-likelihood following one of\+ \ the avaliable distributions.\+ \ The default is Gaussian."+ )+ <*> switch+ ( long "restart"+ <> help "If set, it samples the initial values of\+ \ the parameters using a Gaussian distribution N(0, 1),\+ \ otherwise it uses the original values of the expression." )+ <*> option auto+ ( long "seed"+ <> metavar "SEED"+ <> showDefault+ <> value (-1)+ <> help "Random seed to initialize the parameters values.\+ \ Used only if restart is enabled.")+ <*> switch+ ( long "report"+ <> help "If set, reports the analysis in a user-friendly\+ \ format instead of csv. It will also include\+ \ confidence interval for the parameters and predictions" )+ <*> switch+ ( long "profile"+ <> help "If set, it will use profile likelihood to calculate the CIs." )+ <*> switch+ ( long "simple"+ <> help "If set, calculates only SSE.")+ <*> option auto+ ( long "sigma"+ <> metavar "SIGMA"+ <> showDefault+ <> value 0.001+ <> help "Estimation of error for Guassian distribution.")+ <*> option auto+ ( long "alpha"+ <> metavar "ALPHA"+ <> showDefault+ <> value 0.05+ <> help "Significance level for confidence intervals.")+ <*> option auto+ ( long "ptype"+ <> metavar "[Bates | ODE | Constrained]"+ <> showDefault+ <> value Constrained+ <> help "Profile Likelihood method. Default: Constrained. NOTE: Constrained method only calculates the endpoint."+ )++-- helper functions to show the possible options+mkDescription :: Show a => [a] -> [String]+mkDescription = map (envelope '\'' . map toLower . show) + where+ envelope :: a -> [a] -> [a]+ envelope c xs = c : xs <> [c]+{-# INLINE mkDescription #-}++sralgsHelp :: [String]+sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..]+{-# INLINE sralgsHelp #-}++distHelp :: [String]+distHelp = mkDescription [toEnum 0 :: Distribution ..]+{-# INLINE distHelp #-}++-- helper functions to parse the options+mkReader :: Read a => String -> (a -> b) -> String -> ReadM b+mkReader err val sr = eitherReader + $ case readMaybe sr of+ Nothing -> pure (Left err)+ Just x -> pure (Right (val x))++sralgsReader :: ReadM SRAlgs+sralgsReader =+ str >>= (mkReader errMsg id . map toUpper)+ where+ errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp++--s2Reader :: ReadM (Maybe Double)+--s2Reader =+-- str >>= \s -> mkReader ("wrong format " <> s) Just s++distRead :: ReadM Distribution+distRead =+ str >>= \s -> mkReader ("unsupported distribution " <> s) id (capitalize s)+ where+ capitalize "" = ""+ capitalize (c:cs) = toUpper c : if length cs == 2 then map toUpper cs else map toLower cs
+ apps/srtools/IO.hs view
@@ -0,0 +1,230 @@+{-# language BlockArguments #-}+{-# language LambdaCase #-}+module IO where++import System.IO ( hClose, hPutStrLn, openFile, stderr, stdout, IOMode(WriteMode), Handle )+import qualified Data.Massiv.Array as A+import Data.List ( intercalate )+import Control.Monad ( unless, forM_ )+import System.Random ( StdGen )++import Data.SRTree ( SRTree (..), Fix (..), var, floatConstsToParam, relabelVars )+import Algorithm.SRTree.Likelihoods ( Distribution (..) )+import Algorithm.SRTree.ConfidenceIntervals ( printCI, BasicStats(_stdErr, _corr), CI )+import qualified Data.SRTree.Print as P+import Data.SRTree.Eval ( compMode )++import Args ( Args(outfile, alpha,dist,niter,sigma) )+import Report+import Data.SRTree.Recursion ( cata )++import Debug.Trace ( trace, traceShow )++-- Header of CSV file+csvHeader :: String+csvHeader = intercalate "," (basicFields <> optFields <> modelFields)+{-# inline csvHeader #-}++csvHeaderSimple :: String+csvHeaderSimple = intercalate "," (basicFields <> optFields)+{-# inline csvHeaderSimple #-}++-- Open file if filename is not empty+openWriteWithDefault :: Handle -> String -> IO Handle+openWriteWithDefault dflt "" = pure dflt+openWriteWithDefault _ fname = openFile fname WriteMode+{-# INLINE openWriteWithDefault #-}++-- procecss a single tree and return all the available stats+processTree :: Args -- command line arguments+ -> StdGen -- random number generator+ -> Datasets -- datasets+ -> Fix SRTree -- expression in tree format+ -> Int -- index of the parsed expression + -> (BasicInfo, SSE, SSE, Info, (BasicStats, [CI], [CI], [CI], [CI]))+processTree args seed dset t ix = (basic, sseOrig, sseOpt, info, cis)+ where+ (tree, theta0') = floatConstsToParam t+ theta0 = if dist args == Gaussian+ then theta0' <> [sigma args]+ else theta0'++ basic = getBasicStats args seed dset tree theta0 ix+ treeVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> _expr basic+ (_, Nothing) -> _expr basic+ (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix+ sseOrig = getSSE dset t+ sseOpt = getSSE dset (_expr basic)+ info = getInfo args dset (_expr basic) treeVal+ cis = getCI args dset basic (alpha args)++processTreeSimple :: Args -- command line arguments+ -> StdGen -- random number generator+ -> Datasets -- datasets+ -> Fix SRTree -- expression in tree format+ -> Int -- index of the parsed expression+ -> (BasicInfo, SSE, SSE)+processTreeSimple args seed dset t ix = (basic, sseOrig, sseOpt)+ where+ (tree, theta0') = floatConstsToParam t+ theta0 = if dist args == Gaussian+ then theta0' <> [sigma args]+ else theta0'++ basic = getBasicStats args seed dset tree theta0 ix+ treeVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> _expr basic+ (_, Nothing) -> _expr basic+ (Just xV, Just yV) -> _expr $ getBasicStats args seed dset{_xTr = xV, _yTr = yV} tree theta0 ix+ sseOrig = getSSE dset t+ sseOpt = getSSE dset (_expr basic)++-- print the results to a csv format (except CI)+printResults :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO ()+printResults args seed dset varnames exprs = do+ hStat <- openWriteWithDefault stdout (outfile args)+ hPutStrLn hStat csvHeader + forM_ (zip [0..] exprs) + \(ix, tree) -> + case tree of+ Left err -> hPutStrLn stderr ("invalid expression: " <> err)+ Right t -> let treeData = processTree args seed dset t ix+ in hPutStrLn hStat (toCsv treeData varnames)+ unless (null (outfile args)) (hClose hStat)++printResultsSimple :: Args -> StdGen -> Datasets -> [String] -> [Either String (Fix SRTree)] -> IO ()+printResultsSimple args seed dset varnames exprs = do+ hStat <- openWriteWithDefault stdout (outfile args)+ hPutStrLn hStat csvHeaderSimple+ forM_ (zip [0..] exprs)+ \(ix, tree) ->+ case tree of+ Left err -> hPutStrLn stderr ("invalid expression: " <> err)+ Right t -> let treeData = processTreeSimple args seed dset t ix+ in hPutStrLn hStat (toCsvSimple treeData varnames)+ unless (null (outfile args)) (hClose hStat)++-- change the stats into a string+toCsv :: (BasicInfo, SSE, SSE, Info, e) -> [String] -> String+toCsv (basic, sseOrig, sseOpt, info, _) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt <> sInfo)+ where+ sBasic = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic)+ , show (_nNodes basic), show (_nParams basic)+ , intercalate ";" (map show (_params basic))+ , show (_nEvals basic)+ ]+ sSSEOrig = map (showF sseOrig) [_sseTr, _sseVal, _sseTe]+ sSSEOpt = map (showF sseOpt) [_sseTr, _sseVal, _sseTe]+ sInfo = map (showF info) [_bic, _bicVal, _aic, _aicVal, _evidence, _evidenceVal, _mdl, _mdlFreq, _mdlLatt, _mdlVal, _mdlFreqVal, _mdlLattVal, _nllTr, _nllVal, _nllTe, _cc, _cp]+ <> [intercalate ";" (map show (_fisher info))]+ showF p f = show (f p)++toCsvSimple :: (BasicInfo, SSE, SSE) -> [String] -> String+toCsvSimple (basic, sseOrig, sseOpt) varnames = intercalate "," (sBasic <> sSSEOrig <> sSSEOpt)+ where+ sBasic = [ show (_index basic), show (_fname basic), P.showExprWithVars varnames (_expr basic)+ , show (_nNodes basic), show (_nParams basic)+ , intercalate ";" (map show (_params basic))+ , show (_nEvals basic)+ ]+ sSSEOrig = map (showF sseOrig) [_sseTr, _sseVal, _sseTe]+ sSSEOpt = map (showF sseOpt) [_sseTr, _sseVal, _sseTe]+ showF p f = show (f p)++-- get trees of transformed features+getTransformedFeatures :: Fix SRTree -> (Fix SRTree, [Fix SRTree])+getTransformedFeatures = cata $+ \case+ Var ix -> (Fix $ Var ix, [])+ Param ix -> (Fix $ Param ix, [])+ Const x -> (Fix $ Const x, [])+ Uni f (t, vars) -> (Fix $ Uni f t, vars)+ Bin op (l, vs1) (r, vs2) -> case (hasNoParam l, hasNoParam r) of+ (False, True) -> let vs = vs1 <> vs2+ in (Fix $ Bin op l (var $ length vs), vs <> [r])+ (True, False) -> let vs = vs1 <> vs2+ in (Fix $ Bin op (var $ length vs) r, vs <> [l])+ ( _, _) -> (Fix $ Bin op l r, vs1 <> vs2) -- vs1 == vs2 == []++ where+ hasNoParam = cata $+ \case+ Var ix -> True+ Param ix -> False+ Const x -> if floor x == ceiling x then True else False+ Uni f t -> t+ Bin op l r -> l && r++allAreVars :: [Fix SRTree] -> Bool+allAreVars = all isOnlyVar+ where+ isOnlyVar (Fix (Var _)) = True+ isOnlyVar _ = False++-- print the information on screen (including CIs)+printResultsScreen :: Args -> StdGen -> Datasets -> [String] -> String -> [Either String (Fix SRTree)] -> IO ()+printResultsScreen args seed dset varnames targt exprs = do+ forM_ (zip [0..] exprs) + \(ix, tree) -> + case tree of+ Left err -> do putStrLn ("invalid expression: " <> err)+ Right t -> let treeData = processTree args seed dset t ix+ in printToScreen ix treeData+ where+ decim :: Int -> Double -> Double+ decim n x = (fromIntegral . (round :: Double -> Integer)) (x * 10^n) / 10^n+ sdecim n = show . decim n+ nplaces = 4+++ printToScreen ix (basic, _, sseOpt, info, (sts, cis, pis_tr, pis_val, pis_te)) =+ do let (transformedT, newvars) = getTransformedFeatures (_expr basic)+ varnames' = ['z': show ix | ix <- [0 .. length newvars - 1]]+ putStrLn $ "=================== EXPR " <> show ix <> " =================="+ putStr $ targt <> " ~ f(" <> intercalate ", " varnames <> ") = "+ putStrLn $ P.showExprWithVars varnames (_expr basic)++ unless (allAreVars newvars) do+ putStrLn "\nExpression and transformed features: "+ putStr $ targt <> " ~ f(" <> intercalate ", " varnames' <> ") = "+ putStrLn $ P.showExprWithVars varnames' (relabelVars transformedT)+ forM_ (zip varnames' newvars) \(vn, tv) -> do+ putStrLn $ vn <> " = " <> P.showExprWithVars varnames tv++ putStrLn "\n---------General stats:---------\n"+ putStrLn $ "Number of nodes: " <> show (_nNodes basic)+ putStrLn $ "Number of params: " <> show (_nParams basic)+ putStrLn $ "theta = " <> show (_params basic)++ putStrLn "\n----------Performance:--------\n"+ putStrLn $ "SSE (train.): " <> sdecim nplaces (_sseTr sseOpt)+ putStrLn $ "SSE (val.): " <> sdecim nplaces (_sseVal sseOpt)+ putStrLn $ "SSE (test): " <> sdecim nplaces (_sseTe sseOpt)+ putStrLn $ "NegLogLiklihood (train.): " <> sdecim nplaces (_nllTr info)+ putStrLn $ "NegLogLiklihood (val.): " <> sdecim nplaces (_nllVal info)+ putStrLn $ "NegLogLiklihood (test): " <> sdecim nplaces (_nllTe info)++ putStrLn "\n------Selection criteria:-----\n"+ putStrLn $ "BIC: " <> sdecim nplaces (_bic info)+ putStrLn $ "AIC: " <> sdecim nplaces (_aic info)+ putStrLn $ "MDL: " <> sdecim nplaces (_mdl info)+ putStrLn $ "MDL (freq.): " <> sdecim nplaces (_mdlFreq info)+ putStrLn $ "Functional complexity: " <> sdecim nplaces (_cc info)+ putStrLn $ "Parameter complexity: " <> sdecim nplaces (_cp info)++ putStrLn "\n---------Uncertainties:----------\n"+ putStrLn "Correlation of parameters: " + putStrLn $ show $ A.map (decim 2) (_corr sts)+ putStrLn $ "Std. Err.: " <> show (A.map (decim nplaces) (_stdErr sts))+ putStrLn "\nConfidence intervals:\n\nlower <= val <= upper"+ mapM_ (printCI nplaces) cis+ putStrLn "\nConfidence intervals (predictions training):\n\nlower <= val <= upper"+ mapM_ (printCI nplaces) pis_tr+ unless (null pis_val) do+ putStrLn "\nConfidence intervals (predictions validation):\n\nlower <= val <= upper"+ mapM_ (printCI nplaces) pis_val+ unless (null pis_te) do+ putStrLn "\nConfidence intervals (predictions test):\n\nlower <= val <= upper"+ mapM_ (printCI nplaces) pis_te+ putStrLn "============================================================="
+ apps/srtools/Main.hs view
@@ -0,0 +1,34 @@+module Main (main) where++import Data.ByteString.Char8 ( pack, unpack, split )+import Options.Applicative+import System.Random ( getStdGen, mkStdGen )+import Text.ParseSR.IO ( withInput )++import Args+import IO+import Report++main :: IO ()+main = do+ args <- execParser opts+ g <- getStdGen+ (dset, varnames, tgname) <- getDataset args++ let seed = if rseed args < 0 + then g + else mkStdGen (rseed args)+ varnames' = map unpack $ split ',' $ pack varnames+ withInput (infile args) (from args) varnames False (simpl args)+ >>= if toScreen args+ then printResultsScreen args seed dset varnames' tgname -- full report on screen+ else if simple args+ then printResultsSimple args seed dset varnames' -- csv file+ else printResults args seed dset varnames' -- csv file+ where + opts = info (opt <**> helper)+ ( fullDesc <> progDesc "Optimize the parameters of\+ \ Symbolic Regression expressions."+ <> header "srtools - a CLI tool to (re)optimize the numeric\+ \ parameters of symbolic regression expressions"+ )
+ apps/srtools/Report.hs view
@@ -0,0 +1,280 @@+module Report where++import qualified Data.Vector.Storable as VS+import qualified Data.Massiv.Array as A+import Data.Massiv.Array ( Sz(..) )+import Data.Maybe ( fromMaybe )+import Statistics.Distribution.FDistribution ( fDistribution )+import Statistics.Distribution.ChiSquared ( chiSquared )+import Statistics.Distribution ( quantile )+import System.Random ( StdGen, split, randomRs )++import Data.SRTree ( SRTree, Fix (..), floatConstsToParam, paramsToConst, countNodes )+import Data.SRTree.Eval+import Algorithm.SRTree.AD ( forwardModeUniqueJac )+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.ModelSelection ( aic, bic, evidence, logFunctional, logParameters, mdl, mdlFreq, mdlLatt )+import Algorithm.SRTree.ConfidenceIntervals+import Algorithm.SRTree.Opt (minimizeNLLWithFixedParam, minimizeNLL)+import Data.SRTree.Datasets ( loadDataset )+import Data.SRTree.Print ( showExpr )+import Debug.Trace ( trace, traceShow )++import Args++-- store the datasets split into training, validation and test+data Datasets = DS { _xTr :: SRMatrix+ , _yTr :: PVector+ , _xVal :: Maybe SRMatrix+ , _yVal :: Maybe PVector+ , _xTe :: Maybe SRMatrix+ , _yTe :: Maybe PVector+ , _yErrTr :: Maybe PVector+ , _yErrVal :: Maybe PVector+ , _yErrTe :: Maybe PVector+ }++-- basic fields name+basicFields :: [String]+basicFields = [ "Index"+ , "Filename"+ , "Expression"+ , "Number_of_nodes"+ , "Number_of_parameters"+ , "Parameters"+ , "Number_of_evaluations"+ ]++-- basic information about the tree+data BasicInfo = Basic { _index :: Int+ , _fname :: String+ , _expr :: Fix SRTree+ , _nNodes :: Int+ , _nParams :: Int+ , _params :: [Double]+ , _nEvals :: Int+ }++-- optimization fields+optFields :: [String]+optFields = [ "SSE_train_orig"+ , "SSE_val_orig"+ , "SSE_test_orig"+ , "SSE_train_opt"+ , "SSE_val_opt"+ , "SSE_test_opt"+ ]++-- optimization information+data SSE = SSE { _sseTr :: Double+ , _sseVal :: Double+ , _sseTe :: Double+ }++-- model selection fields+modelFields :: [String]+modelFields = [ "BIC"+ , "BIC_val"+ , "AIC"+ , "AIC_val"+ , "Evidence"+ , "EvidenceVal"+ , "MDL"+ , "MDL_Freq"+ , "MDL_Lattice"+ , "MDL_val"+ , "MDL_Freq_val"+ , "MDL_Lattice_val"+ , "NegLogLikelihood_train"+ , "NegLogLikelihood_val"+ , "NegLogLikelihood_test"+ , "LogFunctional"+ , "LogParameters"+ , "Fisher"+ ]++-- model selection information+data Info = Info { _bic :: Double+ , _bicVal :: Double+ , _aic :: Double+ , _aicVal :: Double+ , _evidence :: Double+ , _evidenceVal :: Double+ , _mdl :: Double+ , _mdlFreq :: Double+ , _mdlLatt :: Double+ , _mdlVal :: Double+ , _mdlFreqVal :: Double+ , _mdlLattVal :: Double+ , _nllTr :: Double+ , _nllVal :: Double+ , _nllTe :: Double+ , _cc :: Double+ , _cp :: Double+ , _fisher :: [Double]+ }++-- load the datasets+getDataset :: Args -> IO (Datasets, String, String)+getDataset args = do+ ((xTr, yTr, xVal, yVal), (yErrTr, yErrVal), varnames, tgname) <- loadDataset (dataset args) (hasHeader args)+ let (A.Sz m) = A.size yVal+ let (mXVal, mYVal) = if m == 0+ then (Nothing, Nothing)+ else (Just xVal, Just yVal)+ (mXTe, mYTe, mYErrTe) <- if null (test args)+ then pure (Nothing, Nothing, Nothing)+ else do ((xTe, yTe, _, _), (yErrTe, _), _, _) <- loadDataset (test args) (hasHeader args)+ pure (Just xTe, Just yTe, yErrTe)+ pure (DS xTr yTr mXVal mYVal mXTe mYTe yErrTr yErrVal mYErrTe, varnames, tgname)++getBasicStats :: Args -> StdGen -> Datasets -> Fix SRTree -> [Double] -> Int -> BasicInfo+getBasicStats args seed dset tree theta0 ix+ | anyNaN = getBasicStats args (snd $ split seed) dset tree theta0 ix+ | otherwise = Basic ix (infile args) tOpt nNodes nParams params nEvs+ where+ -- (tree', theta0) = floatConstsToParam tree+ thetas = if restart args+ then A.fromList compMode $ take nParams (randomRs (-1.0, 1.0) seed)+ else A.fromList compMode theta0+ (t,_,nEvs) = minimizeNLL (dist args) (_yErrTr dset) (niter args) (_xTr dset) (_yTr dset) tree thetas+ tOpt = paramsToConst (A.toList t) tree+ nNodes = countNodes tOpt :: Int+ nParams = length theta0+ params = A.toList t+ anyNaN = A.any isNaN t++getSSE :: Datasets -> Fix SRTree -> SSE+getSSE dset tree = SSE tr val te+ where+ (t, th) = floatConstsToParam tree+ tr = sse (_xTr dset) (_yTr dset) t (A.fromList compMode th)+ val = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ (Just xVal, Just yVal) -> sse xVal yVal t (A.fromList compMode th)+ te = case (_xTe dset, _yTe dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ (Just xTe, Just yTe) -> sse xTe yTe t (A.fromList compMode th)++getInfo :: Args -> Datasets -> Fix SRTree -> Fix SRTree -> Info+getInfo args dset tree treeVal =+ Info { _bic = bic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _bicVal = bicVal+ , _aic = aic dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _aicVal = aicVal+ , _evidence = evidence dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _evidenceVal = evidenceVal+ , _mdl = mdl dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _mdlFreq = mdlFreq dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _mdlLatt = mdlLatt dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _mdlVal = mdlVal+ , _mdlFreqVal = mdlFreqVal+ , _mdlLattVal = mdlLattVal+ , _nllTr = nllTr+ , _nllVal = nllVal+ , _nllTe = nllTe+ , _cc = logFunctional tOpt+ , _cp = logParameters dist' (_yErrTr dset) xTr yTr thetaOpt' tOpt+ , _fisher = A.toList $ fisherNLL dist' (_yErrTr dset) xTr yTr tOpt thetaOpt'+ }+ where+ (xTr, yTr) = (_xTr dset, _yTr dset)+ (xVal, yVal) = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> (xTr, yTr)+ (_, Nothing) -> (xTr, yTr)+ (Just a, Just b) -> (a, b)+ (tOpt, thetaOpt_nosig) = floatConstsToParam tree+ thetaOpt = if dist args == Gaussian+ then thetaOpt_nosig <> [sigma args]+ else thetaOpt_nosig+ thetaOpt' = A.fromList compMode thetaOpt++ (tOptVal, thetaOptVal_nosig) = floatConstsToParam treeVal+ thetaOptVal = if dist args == Gaussian+ then thetaOptVal_nosig <> [sigma args]+ else thetaOptVal_nosig+ thetaOptVal' = A.fromList compMode thetaOptVal++ dist' = dist args++ nllTr = nll dist' (_yErrTr dset) (_xTr dset) (_yTr dset) tOpt (A.fromList compMode thetaOpt)+ bicVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> bic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ aicVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> aic dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ evidenceVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> evidence dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ nllVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> nll dist' (_yErrVal dset) xVal yVal tOptVal (A.fromList compMode thetaOptVal)+ mdlVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> mdl dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ mdlFreqVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> mdlFreq dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ mdlLattVal = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ _ -> mdlLatt dist' (_yErrVal dset) xVal yVal thetaOptVal' tOptVal+ nllTe = case (_xTe dset, _yTe dset) of+ (Nothing, _) -> 0.0+ (_, Nothing) -> 0.0+ (Just xTe, Just yTe) -> nll dist' (_yErrTe dset) xTe yTe tOpt (A.fromList compMode thetaOpt)++getCI :: Args -> Datasets -> BasicInfo -> Double -> (BasicStats, [CI], [CI], [CI], [CI])+getCI args dset basic alpha' = (stats', cis, pis_tr, pis_val, pis_te)+ where+ (Sz n) = A.size yTr+ (tree, _) = floatConstsToParam (_expr basic)+ theta = _params basic+ tau_max = (quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - 0.01))+ tau_max' = sqrt $ quantile (fDistribution (_nParams basic) (n - _nParams basic)) (1 - alpha')+ (xTr, yTr) = (_xTr dset, _yTr dset)+ dist' = dist args+ stats' = getStatsFromModel dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta)+ profiles = getAllProfiles (ptype args) dist' (_yErrTr dset) xTr yTr tree (A.fromList compMode theta) (_stdErr stats') estCIs alpha'+ method = if useProfile args+ then Profile stats' profiles+ else Laplace stats'+ predFun = A.computeAs A.S . predict dist' tree (A.fromList compMode theta)++ prof estPi th t =+ let (thOpt, _, _) = minimizeNLL dist' (_yErrTr dset) 100 xTr yTr t th+ ssr = sse xTr yTr t thOpt+ est = sqrt $ ssr / fromIntegral (n - _nParams basic)+ stdErr = _stdErr stats' A.! 0+ fun = case ptype args of+ Bates -> getProfile dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max 0+ ODE -> getProfileODE dist' (_yErrTr dset) xTr yTr t thOpt stdErr estPi tau_max 0+ Constrained -> getProfileCnstr dist' (_yErrTr dset) xTr yTr t thOpt stdErr tau_max' 0+ in case fun of+ Left th' -> trace "found better optima" $ prof estPi th' t+ Right p -> (_tau2theta p, _opt p)+ jac xss = forwardModeUniqueJac xss (A.fromList compMode theta) tree -- FIX++ estCIs = paramCI (Laplace stats') n (A.fromList compMode theta) 0.001+ cis = paramCI method n (A.fromList compMode theta) alpha'++ estPIS_tr = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []+ estPIS_val = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []+ estPIS_te = predictionCI (Laplace stats') dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' []++ pis_tr = predictionCI method dist' predFun jac prof xTr tree (A.fromList compMode theta) alpha' estPIS_tr+ pis_val = case (_xVal dset, _yVal dset) of+ (Nothing, _) -> []+ (Just xVal, _) -> predictionCI method dist' predFun jac prof xVal tree (A.fromList compMode theta) alpha' estPIS_val+ pis_te = case (_xTe dset, _yTe dset) of+ (Nothing, _) -> []+ (Just xTe, _) -> predictionCI method dist' predFun jac prof xTe tree (A.fromList compMode theta) alpha' estPIS_te
+ apps/tinygp/GP.hs view
@@ -0,0 +1,252 @@+{-# LANGUAGE ImportQualifiedPost #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE BangPatterns #-}+module GP where++import Data.SRTree+import Algorithm.SRTree.Opt+import Algorithm.SRTree.Likelihoods+import Data.SRTree.Print+import Data.SRTree.Eval+import Data.SRTree.Recursion ( cata )+import System.Random+import Control.Monad.State.Strict+import Control.Monad+import Data.Vector qualified as V+import Control.Monad (when)+import Data.Massiv.Array qualified as M+import Debug.Trace ( traceShow, trace )+import Util+import Data.List ( intercalate, maximumBy )+import qualified Data.Vector.Mutable as MV++data Method = Grow | Full | BTC+type Rng a = StateT StdGen IO a++type GenUni = Fix SRTree -> Fix SRTree +type GenBin = Fix SRTree -> Fix SRTree -> Fix SRTree+type FitFun = Individual -> Rng Individual++data Individual = Individual { _tree :: Fix SRTree, _fit :: Double, _params :: [PVector] }++instance Show Individual where + show (Individual t f p) = showExpr t <> "," <> show f <> "," <> show p ++toss :: Rng Bool+toss = state random+{-# INLINE toss #-}++randomRange :: (Ord val, Random val) => (val, val) -> Rng val+randomRange rng = state (randomR rng)+{-# INLINE randomRange #-}++randomFrom :: [a] -> Rng a+randomFrom funs = do n <- randomRange (0, length funs - 1)+ pure $ funs !! n+{-# INLINE randomFrom #-}++randomFromV :: V.Vector a -> Rng a+randomFromV funs = do n <- randomRange (0, length funs - 1)+ pure $ funs V.! n+{-# INLINE randomFromV #-}++countNodes' :: Fix SRTree -> Int+countNodes' = cata alg + where + alg (Var _) = 1+ alg (Param _) = 1+ alg (Const _) = 0+ alg (Bin _ l r) = 1 + l + r+ alg (Uni Abs t) = t+ alg (Uni _ t) = 1 + t+{-# INLINE countNodes' #-}+++randomTree :: HyperParams -> Bool -> Rng (Fix SRTree)+randomTree hp grow + | depth <= 1 || size <= 2 = randomFrom term + | (min_depth >= 0 || (depth > 2 && not grow)) && size > 2 = genNonTerm + | otherwise = genTermOrNon+ where + min_depth = _minDepth hp+ depth = _maxDepth hp+ size = _maxSize hp+ term = _term hp+ nonterm = _nonterm hp++ genNonTerm =+ do et <- randomFrom nonterm+ case et of + Left uniT -> uniT <$> randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow+ Right binT -> do l <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1} grow+ r <- randomTree hp{_minDepth = min_depth-1, _maxDepth = depth - 1, _maxSize = size - 1 - countNodes' l} grow+ pure (binT l r)+ genTermOrNon = do r <- toss+ if r+ then randomFrom term + else genNonTerm ++data HyperParams = + HP { _minDepth :: Int + , _maxDepth :: Int+ , _maxSize :: Int + , _popSize :: Int+ , _tournSize :: Int+ , _pc :: Double + , _pm :: Double + , _term :: [Fix SRTree]+ , _nonterm :: [Either GenUni GenBin] + }++tournament :: HyperParams -> V.Vector Individual -> Rng Individual+tournament hp pop = do+ selection <- replicateM (_tournSize hp) (randomFromV $ V.filter (not.isNaN._fit) pop)+ let maxFitness = maximum (fmap _fit selection)+ champions = V.filter ((== maxFitness) . _fit) pop+ if null selection+ then randomFromV pop+ else randomFromV champions++randomIndividual :: HyperParams -> FitFun -> Bool -> Rng Individual+randomIndividual hyperparams fitFun grow = do+ t <- randomTree hyperparams grow + let p = countParams t+ --theta' <- replicateM p (randomRange (-1,1))+ fitFun $ Individual t 0.0 [] -- (M.fromList compMode theta' :: PVector)+ --pure ind+ --if isInfinite (_fit ind)+ -- then randomIndividual hyperparams fitFun grow + -- else pure ind++initialPop :: HyperParams -> FitFun -> Rng (V.Vector Individual)+initialPop hyperparams fitFun = do + let depths = [3 .. _maxDepth hyperparams]+ pop <- forM depths $ \md -> + do let m = _popSize hyperparams `div` (_maxDepth hyperparams - 3 + 1)+ g = V.fromList . take m $ cycle [True, False]+ mapM (randomIndividual hyperparams{ _maxDepth = md} fitFun) g+ pure (V.concat pop)++fitnessMV :: Distribution -> [(SRMatrix, PVector, Maybe PVector)] -> Individual -> Rng Individual+fitnessMV dist datas ind = do+ fs <- forM datas (fitness dist ind)+ let fitOpt = minimum $ map fst fs+ pure ind{_fit = fitOpt, _params = map snd fs}++fitness :: Distribution -> Individual -> (SRMatrix, PVector, Maybe PVector) -> Rng (Double, PVector)+fitness dist ind (x, y, e) = do+ let tree = relabelParams $ _tree ind+ p = countParams tree+ theta1' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))+ theta2' <- M.fromList M.Seq <$> replicateM p (randomRange (-1,1))+ let (theta1, f1, _) = minimizeNLL dist e 50 x y tree theta1'+ (theta2, f2, _) = minimizeNLL dist e 50 x y tree theta2'+ fit1 = if isNaN f1 then (-1.0/0.0) else negate f1+ fit2 = if isNaN f1 then (-1.0/0.0) else negate f2+ thetaOpt = if fit1 > fit2 then theta1 else theta2+ fitOpt = max fit1 fit2+ pure (fitOpt, thetaOpt)+++mutate :: HyperParams -> Individual -> Rng (Maybe Individual)+mutate hp ind = do+ let sz = countNodes' (_tree ind)+ p <- state $ randomR (0, sz-1)+ b <- state random+ t <- go p (_maxSize hp) (_tree ind)+ --(t, b) <- go sz (_pm hp) (_tree ind)+ if b <= _pm hp && countNodes t <= _maxSize hp+ then pure . Just $ Individual t 0.0 []+ else pure Nothing+ where+ go 0 msz t = randomTree hp{_maxSize = msz-1} True+ go n msz (Fix (Uni f t)) = Fix . Uni f <$> go (n-1) (msz-1) t+ go n msz (Fix (Bin op l r)) = do+ let nl = countNodes l+ nr = countNodes r+ if nl <= n - 1+ then Fix . Bin op l <$> go (n-nl-1) (msz-nl-1) r+ else do l' <- go (n-1) (msz-nr-1) l+ pure $ Fix $ Bin op l' r++crossover :: HyperParams -> Individual -> Individual -> Rng (Maybe Individual)+crossover hp ind1 ind2 = do+ b <- state random+ if b < (_pc hp)+ then do let n1 = countNodes $ _tree ind1+ n2 = countNodes $ _tree ind2+ p1 <- state $ randomR (0, n1-1)+ p2 <- state $ randomR (0, n2-1)+ let part1 = pickLeft p1 $ _tree ind1+ part2 = pickRight p2 $ _tree ind2+ t = part1 part2+ n = countNodes t+ if n <= _maxSize hp+ then pure . Just $ ind1{_tree = t}+ else pure Nothing+ else pure Nothing+ where+ pickRight :: Int -> Fix SRTree -> Fix SRTree+ pickRight 0 node = node+ pickRight n (Fix (Uni f t)) = pickRight (n-1) t+ pickRight n (Fix (Bin op l r)) = let nl = countNodes l+ in if nl <= n-1+ then pickRight (n-nl-1) r+ else pickRight (n-1) l+ pickLeft :: Int -> Fix SRTree -> (Fix SRTree -> Fix SRTree)+ pickLeft 0 node = \t -> t+ pickLeft n (Fix (Uni f t)) = let g = pickLeft (n-1) t in \t' -> Fix $ Uni f (g t')+ pickLeft n (Fix (Bin op l r)) = let nl = countNodes l+ in if nl <= n-1+ then let g = pickLeft (n-nl-1) r in \t -> Fix $ Bin op l (g t)+ else let g = pickLeft (n-1) l in \t -> Fix $ Bin op (g t) r+++evolve :: HyperParams -> FitFun -> V.Vector Individual -> Rng Individual+evolve hp fitFun pop = do + parent1 <- tournament hp pop+ parent2 <- tournament hp pop + mChild <- crossover hp parent1 parent2+ child' <- case mChild of+ Nothing -> mutate hp parent1+ Just child -> mutate hp child+ --let p = countParams (_tree child')+ --theta' <- M.fromList compMode <$> replicateM p (randomRange (-1,1))+ case child' of+ Nothing -> pure parent1+ Just c -> fitFun c++printFinal dist ind dataTrains dataTests = do+ let tree = relabelParams $ _tree ind+ thetas = _params ind+ mseTrain = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTrains+ mseTest = maximum $ map (\(theta, (x,y,e)) -> nll dist e x y tree theta) $ zip thetas dataTests+ r2Train = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTrains+ r2Test = minimum $ map (\(theta, (x,y,e)) -> r2 x y tree theta) $ zip thetas dataTests+ thetaStr = intercalate "_" $ map (intercalate ";" . map show . M.toList) thetas+ putStrLn "id,Expression,theta,size,MSE_train,MSE_test,R2_train,R2_test"+ putStr $ "0," <> showExpr tree <> "," <> thetaStr <> "," <> show (countNodes tree) <> "," <> show mseTrain <> "," <> show mseTest <> "," <> show r2Train <> "," <> show r2Test++report :: Int -> V.Vector Individual -> IO ()+report gen = mapM_ reportOne+ where reportOne ind = do putStr (show gen)+ putStr ": "+ putStr (showExpr (_tree ind))+ putStr " - " + putStr (show (_fit ind))+ putStr " "+ print (map M.toList $ _params ind)+{-# INLINE report #-}++evolution :: Int -> HyperParams -> FitFun -> Rng (Individual)+evolution gen hp fitFun = do + pop <- initialPop hp fitFun+ --liftIO $ report (-1) pop+ go gen pop+ where + go 0 !pop = pure $ pop V.! 0+ go n !pop = do+ let best = V.maximumOn _fit $ V.filter (not.isNaN._fit) pop+ pop' <- V.modify (\v -> MV.write v 0 best) <$> V.replicateM (_popSize hp) (evolve hp fitFun pop)+ --liftIO $ report (gen-n) pop'+ go (n-1) pop'
+ apps/tinygp/Initialization.hs view
@@ -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+ -}
+ apps/tinygp/Main.hs view
@@ -0,0 +1,116 @@+module Main (main) where++import GP ( HyperParams(HP), fitnessMV, evolution, printFinal )+import Data.SRTree+import System.Random ( getStdGen )+import Control.Monad.State.Strict ( evalStateT )+import Data.SRTree.Datasets ( loadDataset ) +import Options.Applicative+import Data.Massiv.Array +import Util+import Algorithm.SRTree.Likelihoods+import Data.SRTree.Datasets++-- Data type to store command line arguments+data Args = Args+ { dataset :: String,+ _testData :: String,+ popSize :: Int,+ gens :: Int,+ _maxSize :: Int,+ pc :: Double,+ pm :: Double,+ _nonterminals :: String,+ _nTournament :: Int,+ _distribution :: Distribution+ }+ deriving (Show)+++-- parser of command line arguments+opt :: Parser Args+opt = Args+ <$> strOption+ ( long "dataset"+ <> short 'd'+ <> metavar "INPUT-FILE"+ <> help "CSV dataset." )+ <*> strOption+ ( long "test"+ <> value ""+ <> metavar "INPUT-FILE"+ <> help "CSV dataset." )+ <*> option auto+ ( long "population"+ <> short 'p'+ <> metavar "POP-SIZE"+ <> showDefault+ <> value 100+ <> help "Population size." )+ <*> option auto+ ( long "generations"+ <> short 'g'+ <> metavar "GENS"+ <> showDefault+ <> value 100+ <> help "Number of generations." )+ <*> option auto+ ( long "max-size"+ <> metavar "SIZE"+ <> showDefault+ <> value 20+ <> help "maximum expression size." )+ <*> option auto+ ( long "probCx"+ <> metavar "PC"+ <> showDefault+ <> value 0.9+ <> help "Crossover probability." )+ <*> option auto+ ( long "probMut"+ <> metavar "PM"+ <> showDefault+ <> value 0.3+ <> help "Mutation probability." )+ <*> strOption+ ( long "non-terminals"+ <> value "Add,Sub,Mul,Div,PowerAbs,Recip"+ <> showDefault+ <> help "set of non-terminals to use in the search."+ )+ <*> option auto+ ( long "tournament-size"+ <> value 2+ <> showDefault+ <> help "tournament size."+ )+ <*> option auto+ ( long "distribution"+ <> value MSE+ <> showDefault+ <> help "distribution of the data.")++nonterms = [Right (+), Right (-), Right (*), Right (/), Right (\l r -> Fix $ Bin PowerAbs l r), Left recip, Left log, Left exp, Left (\t -> Fix $ Uni SqrtAbs t)]++main :: IO ()+main = do+ args <- execParser opts+ g <- getStdGen+ --(x, y, _) <- loadTrainingOnly (dataset args) True+ --(x_test, y_test, _) <- loadTrainingOnly (_testData args) True++ let datasets = words (dataset args)+ dataTrains <- Prelude.mapM (flip loadTrainingOnly True) datasets -- load all datasets+ dataTests <- if null (_testData args)+ then pure dataTrains+ else Prelude.mapM (flip loadTrainingOnly True) $ words (_testData args)++ let hp = HP 3 10 (_maxSize args) (popSize args) (_nTournament args) (pc args) (pm args) terms (parseNonTerms $ _nonterminals args)+ (Sz2 _ nFeats) = size . getX $ head dataTrains+ terms = [var ix | ix <- [0 .. nFeats-1]] <> [param ix | ix <- [0 .. 5]]+ best <- evalStateT (evolution (gens args) hp (fitnessMV (_distribution args) dataTrains)) g+ printFinal (_distribution args) best dataTrains dataTests+ where+ opts = info (opt <**> helper)+ ( fullDesc <> progDesc "Very simple example of GP using SRTree."+ <> header "tinyGP - a very simple example of GP using SRTRee." )
+ apps/tinygp/Util.hs view
@@ -0,0 +1,63 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE TupleSections #-}++module Util where++import qualified Data.Map.Strict as Map+import Data.Massiv.Array as MA hiding (forM_, forM)+import Data.SRTree+import Data.SRTree.Eval+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Build+import Algorithm.EqSat.Info++import Algorithm.SRTree.NonlinearOpt+import System.Random+import Algorithm.SRTree.Likelihoods+--import Algorithm.SRTree.ModelSelection+--import Algorithm.SRTree.Opt+import qualified Data.IntMap.Strict as IM+import Control.Monad.State.Strict+import Control.Monad ( when, replicateM, forM, forM_ )+import Data.Maybe ( fromJust )+import Data.List ( maximumBy )+import Data.Function ( on )+import List.Shuffle ( shuffle )+import Data.List.Split ( splitOn )+import Data.Char ( toLower )+import qualified Data.IntSet as IntSet+import Data.SRTree.Datasets+import Algorithm.EqSat.Queries+++ {-+type DataSet = (SRMatrix, PVector, Maybe PVector)+++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b+-}++parseNonTerms = Prelude.map toNonTerm . splitOn ","+ where+ binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+ uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+ toNonTerm xs' = let xs = Prelude.map toLower xs'+ in case binTerms Map.!? xs of+ Just op -> Right $ \l r -> Fix $ Bin op l r+ Nothing -> case uniTerms Map.!? xs of+ Just f -> Left $ \t -> Fix $ Uni f t+ Nothing -> error $ "invalid non-terminal " <> show xs+
+ src/Algorithm/EqSat.hs view
@@ -0,0 +1,186 @@+{-# 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 <- getBestExpr root+ --info <- gets ((IntMap.! root) . _eClass)+ --info2 <- gets ((IntMap.! 9) . _eClass)+ --traceShow (info, info2) $+ if not end -- if had an early stop+ then do modify' (const emptyGraph) >> eqSat best rules costFun it -- reapplies eqsat on the best so far+ else pure best++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 -- TODO: partial db is still incomplete + --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' > 1500 -- 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..]+ (rls, sch') = runState (matchAll rules') IntMap.empty+ --matches <- mapM (\rule -> map (rule,) <$> match (source rule)) $ concat rls+ --mapM_ (uncurry (applyMergeOnlyMatch costFun)) $ take 500 $ concat matches+ matches <- getNMatches 500 rls+ rebuild costFun+ -- recalculate heights+ --calculateHeights+ where+ rules' = concatMap replaceEqRules rules++ -- replaces the equality rules with two one-way rules+ replaceEqRules :: Rule -> [Rule]+ replaceEqRules (p1 :=> p2) = [p1 :=> p2]+ replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1]+ replaceEqRules (r :| cond) = map (:| cond) $ replaceEqRules r++ getNMatches n [] = pure []+ getNMatches 0 _ = pure []+ getNMatches n ([]:rss) = getNMatches n rss+ getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> match (source r)+ let (x, y) = splitAt n matches+ m = length x+ if m == n+ then pure matches+ else do matches' <- getNMatches (n - length x) (rs:rss)+ pure (matches <> matches')+++-- | matches the rules given a scheduler+matchWithScheduler :: Int -> Int -> Rule -> Scheduler [Rule] -- [(Rule, (Map ClassOrVar ClassOrVar, ClassOrVar))]+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
+ src/Algorithm/EqSat/Build.hs view
@@ -0,0 +1,640 @@+{-# 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 Control.Monad.Identity+import Data.SRTree.Recursion (cataM)+import Algorithm.EqSat.Info+import qualified Data.IntSet as IntSet+import Data.Maybe+import Data.Sequence (Seq(..), (><))+import Data.List ( nub )+import Debug.Trace (trace, traceShow)++-- | adds a new or existing e-node (merging if necessary)+add :: Monad m => CostFun -> ENode -> EGraphST m EClassId+add costFun enode =+ do enode'' <- canonize enode -- canonize e-node+ constEnode <- calculateConsts enode''+ enode' <- case constEnode of+ ConstVal x -> pure $ Const x+ ParamIx x -> pure $ Param x+ _ -> case enode'' of+ Bin Sub c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)+ pure $ case constType of+ ParamIx x -> Bin Add c1 c2+ _ -> enode''+ Bin Div c1 c2 -> do constType <- gets (_consts . _info . (IntMap.! c2) . _eClass)+ pure $ case constType of+ ParamIx x -> Bin Mul c1 c2+ _ -> enode''+ _ -> pure $ enode''++ 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))+{-# INLINE rebuild #-}++-- | 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')+{-# INLINE repair #-}++-- | 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')+ . over (eDB . refits) (Set.insert ecId')+ _ <- modifyEClass costFun ecId'+ pure ()+{-# INLINE repairAnalysis #-}++-- | merge to equivalent e-classes+merge :: Monad m => CostFun -> EClassId -> EClassId -> EGraphST m EClassId+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 . IntMap.insert subO led) -- points sub e-class to leader to maintain consistency+ let -- create new e-class with same id as led+ newC = EClass led+ (_eNodes ledC `Set.union` _eNodes subC)+ (_parents ledC <> _parents subC)+ (min (_height ledC) (_height subC))+ (joinData (_info ledC) (_info subC))++ modify' $ over eClass (IntMap.insert led newC . IntMap.delete sub) -- delete sub e-class and replace leader+ . over (eDB . worklist) (_parents subC <>) -- insert parents of sub into worklist+ when (_info newC /= _info ledC) -- if there was change in data,+ $ modify' $ over (eDB . analysis) (_parents ledC <>) -- insert parents into analysis+ . over (eDB . refits) (Set.insert led)+ when (_info newC /= _info subC)+ $ modify' $ over (eDB . analysis) (_parents subC <>)+ updateDBs newC led ledC ledO sub subC subO+ modifyEClass costFun led+ --forM_ (_eNodes newC) $ \en -> addToDB (decodeEnode en) led+ pure led++ getLeaderSub c1 c1O c2 c2O =+ 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})+ when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId)+ case maybeEid of+ Nothing -> pure ecId+ Just eid' -> merge costFun eid' ecId++ ParamIx x -> do+ let en = Param x+ c <- calculateCost costFun en+ ens <- gets (_eNodes . (IntMap.! ecId) . _eClass)+ let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en }+ maybeEid <- gets ((Map.!? en) . _eNodeToEClass)+ modify' $ over eClass (IntMap.insert ecId ec{_eNodes = Set.insert (encodeEnode en) (_eNodes ec), _info = infoEc})+ when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (Set.insert ecId)+ -- TODO: what happen to the orphans?+ case maybeEid of+ Nothing -> pure ecId+ Just eid' -> merge costFun eid' ecId++ _ -> 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)+{-# INLINE createDB #-}++createDBBest :: Monad m => EGraphST m DB+createDBBest = do modify' $ over (eDB . patDB) (const Map.empty)+ ecls <- gets (Prelude.map (\(eId, ec) -> (_best (_info ec), eId)) . IntMap.toList . _eClass)+ mapM_ (uncurry addToDB) ecls+ gets (_patDB . _eDB)++-- | `addToDB` adds an e-node and e-class id to the database+addToDB :: Monad m => ENode -> EClassId -> EGraphST m () -- State DB ()+addToDB enode' eid = do+ eid' <- canonical eid+ isConst <- gets (_consts . _info . (IntMap.! eid') . _eClass)+ let enode = case isConst of+ ConstVal x -> Const x+ ParamIx x -> Param x+ _ -> enode'+ let ids = eid : childrenOf enode -- we will add the e-class id and the children ids+ op = getOperator enode -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB+ trie <- gets ((Map.!? op) . _patDB . _eDB) -- gets the entry for op, if it exists+ case populate trie ids of -- populates the trie+ Nothing -> pure ()+ Just t -> modify' $ over (eDB . patDB) (Map.insert op t) -- if something was created, insert back into the DB+{-# INLINE addToDB #-}++-- | Populates an IntTrie with a sequence of e-class ids+populate :: Maybe IntTrie -> [EClassId] -> Maybe IntTrie+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))+{-# INLINE populate #-}++canonizeMap :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m (Map ClassOrVar ClassOrVar, ClassOrVar)+canonizeMap (subst, cv) = (,cv) <$> traverse g subst -- Map.fromList <$> traverse f (Map.toList subst)+ 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)+{-# INLINE canonizeMap #-}++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 ()+{-# INLINE applyMatch #-}++applyMergeOnlyMatch :: Monad m => CostFun -> Rule -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m ()+applyMergeOnlyMatch costFun rule match' =+ do let conds = getConditions rule+ match <- canonizeMap match'+ validHeight <- isValidHeight match+ validConds <- mapM (`isValidConditions` match) conds+ when (validHeight && and validConds) $+ do maybe_eid <- classOfENode costFun (fst match) (target rule)+ case maybe_eid of+ Nothing -> pure ()+ Just eid -> do merge costFun (getInt (snd match)) eid+ pure ()+{-# INLINE applyMergeOnlyMatch #-}++-- | gets the e-node of the target of the rule+-- TODO: add consts and modify+classOfENode :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m (Maybe EClassId)+classOfENode costFun subst (VarPat c) = do let maybeEid = getInt <$> subst Map.!? Right (fromEnum c)+ 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)+{-# INLINE classOfENode #-}++-- | adds the target of the rule into the e-graph+reprPrat :: Monad m => CostFun -> Map ClassOrVar ClassOrVar -> Pattern -> EGraphST m EClassId+reprPrat costFun subst (VarPat c) = canonical $ getInt $ subst Map.! Right (fromEnum c)+reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target)+ add costFun (replaceChildren newChildren target)+{-# INLINE reprPrat #-}++isValidHeight :: Monad m => (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool+isValidHeight match = do+ h <- case snd match of+ Left ec -> do ec' <- canonical ec+ gets (_height . (IntMap.! ec') . _eClass)+ Right _ -> pure 0+ pure $ h < 15+{-# INLINE isValidHeight #-}++-- | returns `True` if the condition of a rule is valid for that match+isValidConditions :: Monad m => Condition -> (Map ClassOrVar ClassOrVar, ClassOrVar) -> EGraphST m Bool+isValidConditions cond match = gets $ cond (fst match)+{-# INLINE isValidConditions #-}++-- * Tree to e-graph conversion and utility functions++-- | Creates an e-graph from an expression tree+fromTree :: Monad m => CostFun -> Fix SRTree -> EGraphST m EClassId+fromTree costFun = cataM sequence (add costFun)+{-# INLINE fromTree #-}++-- | Builds an e-graph from multiple independent trees+fromTrees :: Monad m => CostFun -> [Fix SRTree] -> EGraphST m [EClassId]+fromTrees costFun = foldM (\rs t -> do eid <- fromTree costFun t; pure (eid:rs)) []+{-# INLINE fromTrees #-}++countParamsEg :: EGraph -> EClassId -> Int+countParamsEg eg rt = countParams . runIdentity $ getBestExpr rt `evalStateT` eg+countParamsUniqEg :: EGraph -> EClassId -> Int+countParamsUniqEg eg rt = countParamsUniq . runIdentity $ getBestExpr rt `evalStateT` eg+++-- | gets the best expression given the default cost function+getBestExpr :: Monad m => EClassId -> EGraphST m (Fix SRTree)+getBestExpr eid = do eid' <- canonical eid+ best <- gets (_best . _info . (IntMap.! eid') . _eClass)+ childs <- mapM getBestExpr $ childrenOf best+ pure . Fix $ replaceChildren childs best+{-# INLINE getBestExpr #-}++getBestENode eid = do eid' <- canonical eid+ gets (_best . _info . (IntMap.! eid') . _eClass)+{-# INLINE getBestENode #-}++-- | 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+{-# INLINE getExpressionFrom #-}++-- | 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)+{-# INLINE getAllExpressionsFrom #-}++getNExpressionsFrom :: Monad m => Int -> EClassId -> EGraphST m [Fix SRTree]+getNExpressionsFrom n eId' = getNExpressionsFrom' n 15 eId' ++getNExpressionsFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree]+getNExpressionsFrom' _ 0 _ = pure []+getNExpressionsFrom' n d eId' = do+ eId <- canonical eId'+ nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)+ (concat <$> go n d nodes)+ where+ isTerm (Var _) = True+ isTerm (Const _) = True+ isTerm (Param _) = True+ isTerm _ = False+ toTree (Var ix) = Fix $ Var ix+ toTree (Const x) = Fix $ Const x+ toTree (Param ix) = Fix $ Param ix+ toTree _ = undefined++ go n' _ [] = pure []+ go n' 0 ts = pure []+ go n' d (node:ns) = do+ tt <- Prelude.map Fix <$> case node of+ Bin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l+ r' <- getNExpressionsFrom' n' (d-1) r+ pure $ Prelude.take n [Bin op li ri | li <- l', ri <- r']+ Uni f t -> Prelude.map (Uni f) <$> getNExpressionsFrom' n' (d-1) t+ Var ix -> pure [Var ix]+ Const x -> pure [Const x]+ Param ix -> pure [Param ix]+ let n'' = n' - length tt+ if n'' <= 0+ then pure [tt]+ else do ts <- go n'' (d-1) ns+ pure (tt:ts)++getNEclassFrom :: Monad m => Int -> EClassId -> EGraphST m [[EClassId]]+getNEclassFrom n eid = getNEclassFrom' n 15 eid++getNEclassFrom' :: Monad m => Int -> Int -> EClassId -> EGraphST m [[EClassId]]+getNEclassFrom' _ 0 _ = pure []+getNEclassFrom' n d eId' = do+ eId <- canonical eId'+ nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! eId) . _eClass)+ (Prelude.map (eId:) <$> go n d nodes)+ where+ --go :: Int -> Int -> [ENode] -> EGraphST m [[EClassId]]+ go n' _ [] = pure []+ go n' 0 ts = pure []+ go n' d (node:ns) = do+ tt <- case node of+ Bin op l r -> do l' <- getNEclassFrom' n' (d-1) l+ r' <- getNEclassFrom' n' (d-1) r+ pure $ Prelude.take n [li <> ri | li <- l', ri <- r']+ Uni f t -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]]+ Var ix -> pure [[]]+ Const x -> pure [[]]+ Param ix -> pure [[]]+ pure tt+ --let n'' = n' - length tt+ --if n'' <= 0+ -- then pure [tt]+ -- else do ts <- go n'' (d-1) ns+ -- pure (tt:ts)++getAllChildEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildEClasses eId' = do+ eId <- canonical eId'+ IntSet.toList <$> go [eId] IntSet.empty++ where+ hasNoTerminal :: [ENode] -> Bool+ hasNoTerminal = all (not . null . childrenOf) + getNodes :: Monad m => EClassId -> EGraphST m [ENode]+ getNodes n = gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)++ go :: Monad m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet+ go [] visited = pure visited+ go queue visited = do + nodes <- concatMap childrenOf . concat . filter hasNoTerminal <$> mapM getNodes queue+ eids <- filter (\e -> e `IntSet.notMember` visited) <$> (mapM canonical nodes)+ go eids (visited `IntSet.union` IntSet.fromList queue)+ {-+ go n = do nodes <- gets (map decodeEnode . Set.toList . _eNodes . (IntMap.! n) . _eClass)+ let hasTerminal = any (null . childrenOf) nodes+ eids <- mapM canonical $ concatMap childrenOf nodes+ if hasTerminal+ then pure [n]+ else do eids' <- mapM go eids+ pure ((n : eids) <> concat eids')+ -}+{-# INLINE getAllChildEClasses #-}++getAllChildBestEClasses :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildBestEClasses eId' = do+ eId <- canonical eId'+ nub <$> go eId++ where+ go :: Monad m => Int -> EGraphST m [Int]+ go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass)+ let hasTerminal = (null . childrenOf) node+ eids <- mapM canonical $ childrenOf node+ if hasTerminal+ then pure [n]+ else do eids' <- mapM go eids+ pure ((n : eids) <> concat eids')++getAllChildBestEClassesRep :: Monad m => EClassId -> EGraphST m [EClassId]+getAllChildBestEClassesRep eId' = do+ eId <- canonical eId'+ go eId++ where+ go :: Monad m => Int -> EGraphST m [Int]+ go n = do node <- gets (_best . _info . (IntMap.! n) . _eClass)+ let hasTerminal = (null . childrenOf) node+ eids <- mapM canonical $ childrenOf node+ if hasTerminal+ then pure [n]+ else do eids' <- mapM go eids+ pure (n : concat eids')++-- | returns a random expression rooted at e-class `eId`+getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree)+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+{-# INLINE getRndExpressionFrom #-}++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+{-# INLINE cleanMaps #-}++forceState :: Monad m => StateT s m ()+forceState = get >>= \ !_ -> return ()+{-# INLINE forceState #-}
+ src/Algorithm/EqSat/DB.hs view
@@ -0,0 +1,373 @@+{-# 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 Data.SRTree.Recursion (cata)++import Debug.Trace++-- A Pattern is either a fixed-point of a tree or an+-- index to a pattern variable. The pattern variable matches anything. +data Pattern = Fixed (SRTree Pattern) | VarPat Char deriving (Show, Eq, Ord) -- Fixed structure of a pattern or a variable that matches anything++-- The instance for `IsString` for a `Pattern` is +-- valid only for a single letter char from a-zA-Z. +-- 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++tree2pat :: Fix SRTree -> Pattern+tree2pat = cata alg+ where+ alg (Param ix) = if ix >= 100 then VarPat (toEnum $ ix - 100 + 65) else Fixed $ Param ix+ alg (Var ix) = Fixed $ Var ix+ alg (Const x) = Fixed $ Const x+ alg (Bin op l r) = Fixed $ Bin op l r+ alg (Uni f t) = Fixed $ Uni f t+-- A rule is either a directional rule where pat1 can be replaced by pat2, a bidirectional rule +-- where pat1 can be replaced or replace pat2, or a pattern with a conditional function +-- describing when to apply the rule +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+{-# INLINE target #-}++source :: Rule -> Pattern+source (r :| _) = source r+source (s :=> _) = s+source (s :==: _) = s+{-# INLINE source #-}++getConditions :: Rule -> [Condition]+getConditions (r :| c) = c : getConditions r+getConditions _ = []+{-# INLINE getConditions #-}++cleanDB :: Monad m => EGraphST m ()+cleanDB = modify' $ over (eDB. patDB) (const Map.empty)+{-# INLINE cleanDB #-}++-- | Returns the substitution rules+-- for every match of the pattern `source` inside the e-graph.+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]+{-# INLINE match #-}++-- | 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)+{-# INLINE compileToQuery #-}++-- get the value from the Either Int Int+getInt :: ClassOrVar -> Int+getInt (Left a) = a+getInt (Right a) = a+{-# INLINE getInt #-}++-- | returns the list of the children values+getElems :: SRTree a -> [a]+getElems (Bin _ l r) = [l,r]+getElems (Uni _ t) = [t]+getElems _ = []+{-# INLINE getElems #-}++-- | Creates the substituion map for+-- the pattern variables for each one of the+-- 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)+{-# INLINE genericJoin #-}++ -- [Map.insert x classId y | classId <- domainX db x atoms+ -- , y <- go (updateVar x classId atoms) vars]+++-- | returns the e-class id for a certain variable that+-- matches the pattern described by the atoms+domainX :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar]+domainX var atoms root = do+ let atoms' = filter (elemOfAtom var) atoms -- :: [ClassOrVar] -- look only in the atoms with this var+ map Left <$> intersectAtoms var atoms' root -- find the intersection of possible keys by each atom+{-# INLINE domainX #-}+ --let ss = (map Left+ -- $ intersectAtoms var db+ -- $+ -- in ss++-- | returns all e-class id that can matches this sequence of atoms+intersectAtoms :: Monad m => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId]+intersectAtoms _ [] 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+{-# INLINE intersectAtoms #-}++-- | 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)+{-# INLINE intersectTries #-}++-- | 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'+{-# INLINE updateVar #-}++-- | 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+{-# INLINE isDiffFrom #-}++-- | 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+{-# INLINE elemOfAtom #-}++-- | 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)+{-# INLINE orderedVars #-}
+ src/Algorithm/EqSat/Egraph.hs view
@@ -0,0 +1,358 @@+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE StrictData #-}+{-# LANGUAGE DeriveGeneric #-}+{-# 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 hiding ( get, put )+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 Data.Binary+import qualified Data.Binary as Bin+import qualified Data.Massiv.Array as MA++import GHC.Generics++import Debug.Trace++type EClassId = Int -- NOTE: DO NOT CHANGE THIS, this will break the use of IntMap and IntSet+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+type ECache = IntMap.IntMap PVector++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+{-# INLINE getSmallest #-}++getGreatest :: Ord a => RangeTree a -> (a, EClassId)+getGreatest rt = case rt of+ Empty -> error "empty finger"+ t :|> x -> x+{-# INLINE getGreatest #-}++data EGraph = EGraph { _canonicalMap :: ClassIdMap EClassId -- maps an e-class id to its canonical form+ , _eNodeToEClass :: Map ENode EClassId -- maps an e-node to its e-class id+ , _eClass :: ClassIdMap EClass -- maps an e-class id to its e-class data+ , _eDB :: EGraphDB+ } deriving (Show, Generic)++data EGraphDB = EDB { _worklist :: HashSet (EClassId, ENode) -- e-nodes and e-class schedule for analysis+ , _analysis :: HashSet (EClassId, ENode) -- e-nodes and e-class that changed data+ , _refits :: HashSet EClassId+ , _patDB :: DB -- database of patterns+ , _fitRangeDB :: RangeTree Double -- database of valid fitness+ , _dlRangeDB :: RangeTree Double+ , _sizeDB :: IntMap IntSet -- database of model sizes+ , _sizeFitDB :: IntMap (RangeTree Double) -- hacky! Size x Fitness DB+ , _sizeDLDB :: IntMap (RangeTree Double)+ , _unevaluated :: IntSet -- set of not-evaluated e-classes+ , _nextId :: Int -- next available id+ } deriving (Show, Generic)++data EClass = EClass { _eClassId :: Int -- e-class id (maybe we don't need that here)+ , _eNodes :: HashSet ENodeEnc -- set of e-nodes inside this e-class+ , _parents :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s+ , _height :: Int -- height+ , _info :: EClassData -- data+ } deriving (Show, Eq, Generic)++data Consts = NotConst | ParamIx Int | ConstVal Double deriving (Show, Eq, Generic)+data Property = Positive | Negative | NonZero | Real deriving (Show, Eq, Generic) -- TODO: incorporate properties++data EClassData = EData { _cost :: Cost+ , _best :: ENode+ , _consts :: Consts+ , _fitness :: Maybe Double -- NOTE: this cannot be NaN+ , _dl :: Maybe Double+ , _theta :: [PVector]+ , _size :: Int+ -- , _properties :: Property+ -- TODO: include evaluation of expression from this e-class+ } deriving (Show, Generic)++-- * Serialization+instance Generic (EClassId, ENode)++instance Binary (SRTree EClassId) where+ put (Var ix) = put (0 :: Word8) >> put ix+ put (Param ix) = put (1 :: Word8) >> put ix+ put (Const x) = put (2 :: Word8) >> put x+ put (Uni f t) = put (3 :: Word8) >> put (fromEnum f) >> put t+ put (Bin op l r) = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r++ get = do t <- get :: Get Word8+ case t of+ 0 -> Var <$> get+ 1 -> Param <$> get+ 2 -> Const <$> get+ 3 -> Uni <$> (toEnum <$> get) <*> get+ 4 -> Bin <$> (toEnum <$> get) <*> get <*> get++instance Binary (SRTree ()) where+ put (Var ix) = put (0 :: Word8) >> put ix+ put (Param ix) = put (1 :: Word8) >> put ix+ put (Const x) = put (2 :: Word8) >> put x+ put (Uni f t) = put (3 :: Word8) >> put (fromEnum f)+ put (Bin op l r) = put (4 :: Word8) >> put (fromEnum op)++ get = do t <- get :: Get Word8+ case t of+ 0 -> Var <$> get+ 1 -> Param <$> get+ 2 -> Const <$> get+ 3 -> Uni <$> (toEnum <$> get) <*> pure ()+ 4 -> Bin <$> (toEnum <$> get) <*> pure () <*> pure ()++instance (Binary a, Hashable a) => Binary (HashSet a) where+ put hs = put (Set.toList hs)+ get = Set.fromList <$> get++instance Binary PVector where+ put xs = put (MA.toList xs)+ get = MA.fromList compMode <$> get++instance Binary IntTrie+instance Binary EClass+instance Binary Consts+instance Binary Property+instance Binary EClassData+instance Binary EGraphDB+instance Binary EGraph++instance Eq EClassData where+ EData c1 b1 cs1 ft1 dl1 _ s1 == EData c2 b2 cs2 ft2 dl2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && dl1==dl2 && s1==s2++-- The database maps a symbol to an IntTrie+-- The IntTrie stores the possible paths from a certain e-class+-- 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 (Generic)++-- Shows the IntTrie as {keys} -> {show IntTries}+instance Show IntTrie where+ show (IntTrie k t) = let keys = intercalate "," (map show $ Set.toList k)+ tries = intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t)+ in "{" <> keys <> "} - {" <> tries <> "}"++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+{-# INLINE emptyGraph #-}++-- | returns an empty e-graph DB+emptyDB :: EGraphDB+emptyDB = EDB Set.empty Set.empty Set.empty Map.empty FingerTree.empty FingerTree.empty IntMap.empty IntMap.empty IntMap.empty IntSet.empty 0+{-# INLINE emptyDB #-}++-- | 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)+{-# INLINE trie #-}++-- | Check whether an e-class is a constant value+isConst :: Monad m => EClassId -> EGraphST m Bool+isConst eid = do ec <- gets ((IntMap.! eid) . _eClass)+ case (_consts . _info) ec of+ ConstVal _ -> pure True+ _ -> pure False+{-# INLINE isConst #-}++getFitness :: Monad m => EClassId -> EGraphST m (Maybe Double)+getFitness c = gets (_fitness . _info . (IntMap.! c) . _eClass)+{-# INLINE getFitness #-}+getTheta :: Monad m => EClassId -> EGraphST m ([PVector])+getTheta c = gets (_theta . _info . (IntMap.! c) . _eClass)+{-# INLINE getTheta #-}+getSize :: Monad m => EClassId -> EGraphST m Int+getSize c = gets (_size . _info . (IntMap.! c) . _eClass)+{-# INLINE getSize #-}+isSizeOf :: (Int -> Bool) -> EClass -> Bool+isSizeOf p = p . _size . _info+{-# INLINE isSizeOf #-}+getBestFitness :: Monad m => EGraphST m (Maybe Double)+getBestFitness = do+ bec <- (gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical)+ gets (_fitness . _info . (IntMap.! bec) . _eClass)+getDL :: Monad m => EClassId -> EGraphST m (Maybe Double)+getDL c = gets (_dl . _info . (IntMap.! c) . _eClass)+{-# INLINE getDL #-}
+ src/Algorithm/EqSat/Info.hs view
@@ -0,0 +1,208 @@+-----------------------------------------------------------------------------+-- |+-- 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 dl1 p1 sz1) (EData c2 b2 cn2 fit2 dl2 p2 sz2) =+ --EData (min c1 c2) b (combineConsts cn1 cn2) (minMaybe fit1 fit2) (bestParam p1 p2 fit1 fit2) (min sz1 sz2)+ EData (min c1 c2) (choose b1 b2) (choose cn1 cn2) (maxMaybe fit1 fit2) (choose dl1 dl2) (choose p1 p2) (choose sz1 sz2)+ where+ isFst = c1 <= c2+ choose x y = if isFst then x else y+ chooseF x y = if maxIsFst then x else y++ maxIsFst = case (fit1, fit2) of+ (Nothing, Nothing) -> True+ (Nothing, Just f) -> False+ (Just f , Nothing) -> True+ (Just f1, Just f2) -> f1 >= f2++ maxMaybe Nothing x = x+ maxMaybe x Nothing = x+ maxMaybe x y = max x y++ bestParam Nothing x _ _ = x+ bestParam x Nothing _ _ = x+ bestParam x y (Just f1) (Just f2) = if f1 >= f2 then x else y++ 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 (ParamIx ix) (ConstVal x) = ConstVal x+ combineConsts (ConstVal x) (ParamIx ix) = ConstVal x -- p - p = 0+ combineConsts x y = error (show x <> " " <> show y)++-- | Calculate e-node data (constant values and cost)+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+ --ParamIx x -> ParamIx 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+ eId <- canonical eId'+ tree <- getBestExpr' eId+ let p = fromIntegral (length params)+ let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0+ ec <- gets ((IntMap.! eId) . _eClass)+ let oldFit = _fitness . _info $ ec+ --when (oldFit < Just fit) $ do+ let newInfo = (_info ec){_fitness = Just fit, _theta = params}+ newEc = ec{_info = newInfo}+ sz = _size newInfo+ modify' $ over eClass (IntMap.insert eId newEc)+ if (isNothing oldFit)+ then modify' $ over (eDB . unevaluated) (IntSet.delete eId)+ . over (eDB . fitRangeDB) (insertRange eId fit)+ . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)+ . over (eDB . dlRangeDB) (insertRange eId f_compl)+ else modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId (fromJust oldFit))++insertDL :: Monad m => EClassId -> Double -> EGraphST m ()+insertDL eId fit' = do+ let fit = negate fit'+ ec <- gets ((IntMap.! eId) . _eClass)+ let sz = _size . _info $ ec+ newInfo = (_info ec){_dl = Just fit'}+ newEc = ec{_info=newInfo}+ modify' $ over eClass (IntMap.insert eId newEc)+ modify' $ over (eDB . dlRangeDB) (insertRange eId fit)+ . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith (><) sz Empty)++-- | TODO: remove from here gets the best expression given the default cost function+getBestExpr' :: Monad m => EClassId -> EGraphST m (Fix SRTree)+getBestExpr' eid = do eid' <- canonical eid+ best <- gets (_best . _info . (IntMap.! eid') . _eClass)+ childs <- mapM getBestExpr' $ childrenOf best+ pure . Fix $ replaceChildren childs best+
+ src/Algorithm/EqSat/Queries.hs view
@@ -0,0 +1,221 @@+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TupleSections #-}+-----------------------------------------------------------------------------+-- |+-- 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 qualified Data.Sequence as FingerTree+import qualified Data.Foldable as Foldable+import Data.SRTree (childrenOf)++import Debug.Trace++-- this is too slow for now, it needs a db of its own+-- basically a db for each query we need+getEClassesThat :: Monad m => (EClass -> Bool) -> EGraphST m [EClassId]+getEClassesThat 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 => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopECLassThat b n p = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB)+ >>= go n []+ where+ go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go 0 bests rt = pure bests+ go m bests rt = case rt of+ Empty -> pure bests+ t :|> y -> do let x = snd y+ ecId <- canonical x+ ec <- gets ((IntMap.! ecId) . _eClass)+ if (isInfinite . fromJust . _fitness . _info $ ec)+ then go m bests t+ else if p ec+ then go (m-1) (ecId:bests) t+ else go m bests t++getTopEClassInRange :: Monad m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId]+getTopEClassInRange b n p range = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB)+ >>= go n [] range+ where+ inRange v (x, y)+ | v >= x && v <= y = 0+ | v < x = -1+ | v > y = 1+ | otherwise = 1 ++ go :: Monad m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId]+ go _ bests [] _ = pure bests + go 0 bests (r:rs) rt = go n bests rs rt+ go m bests (r:rs) rt = case rt of+ Empty -> pure bests+ t :|> y -> do let x = snd y+ ecId <- canonical x+ ec <- gets ((IntMap.! ecId) . _eClass)+ if (isInfinite . fromJust . _fitness . _info $ ec)+ then go m bests (r:rs) t+ else do let v = p ec + case (v `inRange` r) of+ 0 -> go (m-1) (ecId:bests) (r:rs) t -- it is in range, go to the next range + -1 -> go n bests rs (t :|> y) -- it is smaller than the range, get the first n of the next range+ 1 -> go m bests (r:rs) t -- y is still greater than the range, keep looking in the same range++getTopECLassIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopECLassIn b n p ecs' = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB)+ >>= go n []+ where+ ecs = Set.fromList ecs'+ go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go 0 bests rt = pure bests+ go m bests rt = case rt of+ Empty -> pure bests+ t :|> y -> do let x = snd y+ ecId <- canonical x+ ec <- gets ((IntMap.! ecId) . _eClass)+ if (isInfinite . fromJust . _fitness . _info $ ec)+ then go m bests t -- pure bests+ else if ecId `Set.member` ecs && p ec+ then go (m-1) (ecId:bests) t+ else go m bests t++getTopECLassNotIn :: Monad m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopECLassNotIn b n p ecs' = do+ let f = if b then _fitRangeDB else _dlRangeDB+ gets (f . _eDB)+ >>= go n []+ where+ ecs = Set.fromList ecs'++ go :: Monad m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go 0 bests rt = pure bests+ go m bests rt = case rt of+ Empty -> pure bests+ t :|> y -> do let x = snd y+ ecId <- canonical x+ ec <- gets ((IntMap.! ecId) . _eClass)+ if (isInfinite . fromJust . _fitness . _info $ ec)+ then go m bests t+ else if not (ecId `Set.member` ecs) && p ec+ then go (m-1) (ecId:bests) t+ else go m bests t++getAllEvaluatedEClasses :: Monad m => EGraphST m [EClassId]+getAllEvaluatedEClasses = do+ gets (_fitRangeDB . _eDB)+ >>= go []+ where+ go :: Monad m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId]+ go bests rt = case rt of+ Empty -> pure bests+ t :|> y -> do let x = snd y+ ecId <- canonical x+ ec <- gets ((IntMap.! ecId) . _eClass)+ if (isInfinite . fromJust . _fitness . _info $ ec)+ then go bests t+ else go (ecId:bests) t++getTopEClassWithSize :: Monad m => Bool -> Int -> Int -> EGraphST m [EClassId]+getTopEClassWithSize b sz n = do+ let fun = if b then _sizeFitDB else _sizeDLDB+ gets (go n [] . (IntMap.!? sz) . fun . _eDB)+ -- >>= mapM canonical+ where+ -- go :: Monad m => Int -> [EClassId] -> Maybe (RangeTree Double) -> EGraphST m [EClassId]+ go _ bests Nothing = []+ go 0 bests (Just rt) = bests+ go m bests (Just rt) = case rt of+ Empty -> bests+ t :|> (f, x) -> if isInfinite f || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t)++getTopFitEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopFitEClassThat = getTopECLassThat True+getTopDLEClassThat :: Monad m => Int -> (EClass -> Bool) -> EGraphST m [EClassId]+getTopDLEClassThat = getTopECLassThat False+getTopFitEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassIn = getTopECLassIn True+getTopDLEClassIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassIn = getTopECLassIn False+getTopFitEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopFitEClassNotIn = getTopECLassNotIn True+getTopDLEClassNotIn :: Monad m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId]+getTopDLEClassNotIn = getTopECLassNotIn True+getTopFitEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]+getTopFitEClassWithSize = getTopEClassWithSize True+getTopDLEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId]+getTopDLEClassWithSize = getTopEClassWithSize False++rebuildAllRanges :: Monad m => EGraphST m ()+rebuildAllRanges = do szF <- gets (_sizeFitDB._eDB) >>= traverse rebuildRange+ dlF <- gets (_sizeDLDB._eDB) >>= traverse rebuildRange+ fR <- gets (_fitRangeDB._eDB) >>= rebuildRange+ dR <- gets (_dlRangeDB._eDB) >>= rebuildRange++ modify' $ over (eDB.fitRangeDB) (const fR)+ . over (eDB.dlRangeDB) (const dR)+ . over (eDB.sizeFitDB) (const szF)+ . over (eDB.sizeDLDB) (const dlF)++canonizeRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)+canonizeRange = traverse (\(x, eid) -> (x,) <$> canonical eid)++rebuildRange :: Monad m => RangeTree Double -> EGraphST m (RangeTree Double)+rebuildRange rt = go Set.empty Empty <$> canonizeRange rt+ where+ go :: Set.HashSet EClassId -> RangeTree Double -> RangeTree Double -> RangeTree Double+ go seen root Empty = root+ go seen root (xs :|> (x,eid)) = go (Set.insert eid seen)+ (if Set.member eid seen+ then root+ else (x, eid) :<| root)+ xs -- (Prelude.filter ((/= eid) . snd) xs)+
+ src/Algorithm/EqSat/SearchSR.hs view
@@ -0,0 +1,223 @@+-----------------------------------------------------------------------------+-- |+-- Module : Algorithm.EqSat.Search+-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability :+--+-- Support functions for search symbolic expressions with e-graphs+--+-----------------------------------------------------------------------------++module Algorithm.EqSat.SearchSR where++import Data.SRTree+import Data.SRTree.Datasets+import System.Random+import Control.Monad.State.Strict+import Algorithm.EqSat.Egraph+import Algorithm.SRTree.Likelihoods+import qualified Data.IntMap as IM+import qualified Data.IntSet as IntSet+import qualified Data.SRTree.Random as Random+import Data.Function ( on )+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.NonlinearOpt+import Control.Monad ( when, replicateM, forM, forM_ )+import Algorithm.EqSat.Egraph+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Info+import Algorithm.EqSat.Build+import Data.Maybe ( fromJust )+import Data.SRTree.Random+import Algorithm.EqSat.Queries+import Data.List ( maximumBy )+import qualified Data.Map.Strict as Map++-- Environment of an e-graph with support to random generator and IO+type RndEGraph a = EGraphST (StateT StdGen IO) a++io :: IO a -> RndEGraph a+io = lift . lift+{-# INLINE io #-}+rnd :: StateT StdGen IO a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _) = 1+myCost (Const _) = 1+myCost (Param _) = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t) = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f t+while p arg prog = do if (p arg)+ then do arg' <- prog arg+ while p arg' prog+ else pure arg++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> PVector -> (Double, PVector)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig =+ if isNaN val -- || isNaN tr+ then (-(1/0), theta) -- infinity+ else (val, theta)+ where+ --tree = relabelParams _tree+ nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ (theta, _, _) = minimizeNLL' VAR1 distribution mYErr nIter x y tree thetaOrig+ evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+ --tr = evalF x y mYErr+ val = evalF x_val y_val mYErr_val++--{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, PVector)+fitnessFunRep nRep nIter distribution dataTrain dataVal tree = do+ let nParams = countParamsUniq tree + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+ let fits = maximumBy (compare `on` fst) $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal tree) thetaOrigs+ pure fits+--{-# INLINE fitnessFunRep #-}+++fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [PVector])+fitnessMV shouldReparam nRep nIter distribution dataTrainsVals _tree = do+ let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree+ response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep nRep nIter distribution dt dv tree+ pure (minimum (Prelude.map fst response), Prelude.map snd response)++++++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId+insertExpr t fitFun = do+ ecId <- fromTree myCost t >>= canonical+ (f, p) <- fitFun t+ insertFitness ecId f p+ pure ecId+ where powabs l r = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+ mf <- getFitness ec+ case mf of+ Nothing -> do+ t <- getBestExpr ec+ (f, p) <- fitFun t+ insertFitness ec f p+ pure True+ Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+ if not (null ecIds)+ then do rndId' <- rnd $ randomFrom ecIds+ rndId <- canonical rndId'+ constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+ case constType of+ NotConst -> pure $ Just rndId+ _ -> pure Nothing+ else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)+getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n)++getBestExprWithSize n =+ do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+ if (not (null ec))+ then do+ bestFit <- getFitness $ head ec+ bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+ pure [(head ec, bestFit)]+ else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+ do grow <- rnd toss+ n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+ t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+ fromTree myCost t >>= canonical++refit fitFun ec = do+ t <- getBestExpr ec+ (f, p) <- fitFun t+ mf <- getFitness ec+ case mf of+ Nothing -> insertFitness ec f p+ Just f' -> when (f > f') $ insertFitness ec f p++--printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest fitFun printExprFun = do+ bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+ bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)+ --refit fitFun bec+ --io.print $ "should be " <> show bestFit+ printExprFun 0 bec++--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))+ where+ go :: Int -> Int -> Double -> RndEGraph [[String]]+ go n ix f+ | n > maxSize = pure []+ | otherwise = do+ ecList <- getBestExprWithSize n+ if not (null ecList)+ then do let (ec, mf) = head ecList+ f' = fromJust mf+ improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f'+ ec' <- canonical ec+ if improved+ then do refit fitFun ec'+ t <- printExprFun ix ec'+ ts <- go (n+1) (ix + if improved then 1 else 0) (max f f')+ pure (t:ts)+ else go (n+1) (ix + if improved then 1 else 0) (max f f')+ else go (n+1) ix f++evaluateUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ forM_ ec $ \c -> do+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p++evaluateRndUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ c <- rnd . randomFrom $ ec+ t <- getBestExpr c+ (f, p) <- fitFun t+ insertFitness c f p+ pure c++-- | check whether an e-node exists or does not exist in the e-graph+doesExist, doesNotExist :: ENode -> RndEGraph Bool+doesExist en = gets ((Map.member en) . _eNodeToEClass)+doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)++-- | check whether the partial tree defined by a list of ancestors will create+-- a non-existent expression when combined with a certain e-node.+doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool+doesNotExistGens [] en = gets ((Map.notMember en) . _eNodeToEClass)+doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass)+ if b+ then pure True+ else case mGrand of+ Nothing -> pure False+ Just gf -> do ec <- gets ((Map.! en) . _eNodeToEClass)+ en' <- canonize (gf ec)+ doesNotExistGens grands en'++-- | check whether combining a partial tree `parent` with the e-node `en'`+-- will create a new expression+checkToken parent en' = do en <- canonize en'+ mEc <- gets ((Map.!? en) . _eNodeToEClass)+ case mEc of+ Nothing -> pure True+ Just ec -> do ec' <- canonical ec+ ec'' <- canonize (parent ec')+ not <$> doesExist ec''
+ src/Algorithm/EqSat/SearchSRCache.hs view
@@ -0,0 +1,244 @@+-----------------------------------------------------------------------------+-- |+-- Module : Algorithm.EqSat.Search+-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability :+--+-- Support functions for search symbolic expressions with e-graphs+--+-----------------------------------------------------------------------------++module Algorithm.EqSat.SearchSRCache where++import Data.SRTree+import Data.SRTree.Datasets+import System.Random+import Control.Monad.State.Strict+import Algorithm.EqSat.Egraph+import Algorithm.SRTree.Likelihoods+import qualified Data.IntMap as IM+import qualified Data.IntSet as IntSet+import qualified Data.SRTree.Random as Random+import Data.Function ( on )+import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.NonlinearOpt+import Control.Monad ( when, replicateM, forM, forM_ )+import Algorithm.EqSat.Egraph+import Algorithm.SRTree.Opt+import Algorithm.EqSat.Info+import Algorithm.EqSat.Build+import Data.Maybe ( fromJust )+import Data.SRTree.Random+import Algorithm.EqSat.Queries+import Data.List ( maximumBy )+import qualified Data.Map.Strict as Map+import Control.Monad.Identity++import Debug.Trace++-- Environment of an e-graph with support to random generator and IO+type RndEGraph a = EGraphST (StateT StdGen (StateT [ECache] IO)) a++io :: IO a -> RndEGraph a+io = lift . lift . lift+{-# INLINE io #-}+getCache :: StateT [ECache] IO a -> RndEGraph a+getCache = lift . lift+rnd :: StateT StdGen (StateT [ECache] IO) a -> RndEGraph a+rnd = lift+{-# INLINE rnd #-}++myCost :: SRTree Int -> Int+myCost (Var _) = 1+myCost (Const _) = 1+myCost (Param _) = 1+myCost (Bin _ l r) = 2 + l + r+myCost (Uni _ t) = 3 + t++while :: Monad f => (t -> Bool) -> t -> (t -> f t) -> f t+while p arg prog = do if (p arg)+ then do arg' <- prog arg+ while p arg' prog+ else pure arg++fitnessFun :: Int -> Distribution -> DataSet -> DataSet -> EGraph -> EClassId -> ECache -> PVector -> (Double, PVector, ECache)+fitnessFun nIter distribution (x, y, mYErr) (x_val, y_val, mYErr_val) egraph root cache thetaOrig =+ if isNaN val -- || isNaN tr+ then (-(1/0), theta,cache') -- infinity+ else (val, theta, cache')+ where+ tree = runIdentity $ getBestExpr root `evalStateT` egraph+ nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ (theta, val, _, cache') = minimizeNLLEGraph VAR1 distribution mYErr nIter x y egraph root cache thetaOrig+ evalF a b c = negate $ nll distribution c a b tree $ if nParams == 0 then thetaOrig else theta+ -- val = evalF x_val y_val mYErr_val++--{-# INLINE fitnessFun #-}++fitnessFunRep :: Int -> Int -> Distribution -> DataSet -> DataSet -> EClassId -> ECache -> RndEGraph (Double, PVector, ECache)+fitnessFunRep nRep nIter distribution dataTrain dataVal root cache = do+ egraph <- get+ let nParams = countParamsUniqEg egraph root + if distribution == ROXY then 3 else if distribution == Gaussian then 1 else 0+ fst' (a, _, _) = a+ thetaOrigs <- replicateM nRep (rnd $ randomVec nParams)+ let fits = maximumBy (compare `on` fst') $ Prelude.map (fitnessFun nIter distribution dataTrain dataVal egraph root cache) thetaOrigs+ pure fits+--{-# INLINE fitnessFunRep #-}+++fitnessMV :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector])+fitnessMV shouldReparam nRep nIter distribution dataTrainsVals root = do+ -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree+ -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm'+ caches <- getCache get+ response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache+ getCache $ put (Prelude.map trd response)+ pure (minimum (Prelude.map fst' response), Prelude.map snd' response)+ where fst' (a, _, _) = a+ snd' (_, a, _) = a+ trd (_, _, a) = a++fitnessMVNoCache :: Bool -> Int -> Int -> Distribution -> [(DataSet, DataSet)] -> EClassId -> RndEGraph (Double, [PVector])+fitnessMVNoCache shouldReparam nRep nIter distribution dataTrainsVals root = do+ -- let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree+ -- WARNING: this should be done BEFORE inserting into egraph, so it's up to the algorithm'+ caches <- getCache get+ response <- forM (Prelude.zip dataTrainsVals caches) $ \((dt, dv), cache) -> fitnessFunRep nRep nIter distribution dt dv root cache+ pure (minimum (Prelude.map fst' response), Prelude.map snd' response)+ where fst' (a, _, _) = a+ snd' (_, a, _) = a+ trd (_, _, a) = a++++-- RndEGraph utils+-- fitFun fitnessFunRep rep iter distribution x y mYErr x_val y_val mYErr_val+insertExpr :: Fix SRTree -> (Fix SRTree -> RndEGraph (Double, [PVector])) -> RndEGraph EClassId+insertExpr t fitFun = do+ ecId <- fromTree myCost t >>= canonical+ (f, p) <- fitFun t+ insertFitness ecId f p+ pure ecId+ where powabs l r = Fix (Bin PowerAbs l r)++updateIfNothing fitFun ec = do+ mf <- getFitness ec+ case mf of+ Nothing -> do+ --t <- getBestExpr ec+ (f, p) <- fitFun ec+ insertFitness ec f p+ pure True+ Just _ -> pure False++pickRndSubTree :: RndEGraph (Maybe EClassId)+pickRndSubTree = do ecIds <- gets (IntSet.toList . _unevaluated . _eDB)+ if not (null ecIds)+ then do rndId' <- rnd $ randomFrom ecIds+ rndId <- canonical rndId'+ constType <- gets (_consts . _info . (IM.! rndId) . _eClass)+ case constType of+ NotConst -> pure $ Just rndId+ _ -> pure Nothing+ else pure Nothing++getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n)+getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n)++getBestExprWithSize n =+ do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical+ if (not (null ec))+ then do+ bestFit <- getFitness $ head ec+ bestP <- gets (_theta . _info . (IM.! (head ec)) . _eClass)+ pure [(head ec, bestFit)]+ else pure []++insertRndExpr maxSize rndTerm rndNonTerm =+ do grow <- rnd toss+ n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. maxSize])+ t <- rnd $ Random.randomTree 3 8 n rndTerm rndNonTerm grow+ fromTree myCost t >>= canonical++refit fitFun ec = do+ --t <- getBestExpr ec+ (f, p) <- fitFun ec+ mf <- getFitness ec+ case mf of+ Nothing -> insertFitness ec f p+ Just f' -> when (f > f') $ insertFitness ec f p++--printBest :: (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+printBest fitFun printExprFun = do+ bec <- gets (snd . getGreatest . _fitRangeDB . _eDB) >>= canonical+ bestFit <- gets (_fitness. _info . (IM.! bec) . _eClass)+ --refit fitFun bec+ --io.print $ "should be " <> show bestFit+ printExprFun 0 bec++--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph ()+paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0))+ where+ go :: Int -> Int -> Double -> RndEGraph [[String]]+ go n ix f+ | n > maxSize = pure []+ | otherwise = do+ ecList <- getBestExprWithSize n+ if not (null ecList)+ then do let (ec, mf) = head ecList+ f' = fromJust mf+ improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f'+ ec' <- canonical ec+ if improved+ then do refit fitFun ec'+ t <- printExprFun ix ec'+ ts <- go (n+1) (ix + if improved then 1 else 0) (max f f')+ pure (t:ts)+ else go (n+1) (ix + if improved then 1 else 0) (max f f')+ else go (n+1) ix f++evaluateUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ forM_ ec $ \c -> do+ --t <- getBestExpr c+ (f, p) <- fitFun c+ insertFitness c f p++evaluateRndUnevaluated fitFun = do+ ec <- gets (IntSet.toList . _unevaluated . _eDB)+ c <- rnd . randomFrom $ ec+ --t <- getBestExpr c+ (f, p) <- fitFun c+ insertFitness c f p+ pure c++-- | check whether an e-node exists or does not exist in the e-graph+doesExist, doesNotExist :: ENode -> RndEGraph Bool+doesExist en = gets ((Map.member en) . _eNodeToEClass)+doesNotExist en = gets ((Map.notMember en) . _eNodeToEClass)++-- | check whether the partial tree defined by a list of ancestors will create+-- a non-existent expression when combined with a certain e-node.+doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool+doesNotExistGens [] en = gets ((Map.notMember en) . _eNodeToEClass)+doesNotExistGens (mGrand:grands) en = do b <- gets ((Map.notMember en) . _eNodeToEClass)+ if b+ then pure True+ else case mGrand of+ Nothing -> pure False+ Just gf -> do ec <- gets ((Map.! en) . _eNodeToEClass)+ en' <- canonize (gf ec)+ doesNotExistGens grands en'++-- | check whether combining a partial tree `parent` with the e-node `en'`+-- will create a new expression+checkToken parent en' = do en <- canonize en'+ mEc <- gets ((Map.!? en) . _eNodeToEClass)+ case mEc of+ Nothing -> pure True+ Just ec -> do ec' <- canonical ec+ ec'' <- canonize (parent ec')+ not <$> doesExist ec''
+ src/Algorithm/EqSat/Simplify.hs view
@@ -0,0 +1,280 @@+{-# 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, rewritesParams, rewriteBasic, rewritesFun, rewritesSimple, rewritesWithConstant, myCost ) 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" * "y" :=> "y" * "x"+ , "x" + "y" :=> "y" + "x"+ --, ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"+ --, (powabs "x" "y") * (powabs "x" "z") :=> powabs "x" ("y" + "x")+ , ("x" + "y") + "z" :=> "x" + ("y" + "z")+ , ("x" + "y") - "z" :=> "x" + ("y" - "z")+ --, ("x" + "y") - "z" :=> "x" + ("y" - "z") -- TODO: check that I don't need that+ , ("x" * "y") * "z" :=> "x" * ("y" * "z")+ , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")+ , "x" - ("y" + "z") :=> ("x" - "y") - "z" -- TODO: check that I don't this+ , "x" - ("y" - "z") :=> ("x" - "y") + "z" -- TODO+ , ("x" * "y") / "z" :=> ("x" / "z") * "y" :| isNotZero "z" -- TODO: inv(x) <=> x^-1 , x/y <=> x*y^-1+ , "x" * ("y" / "z") :=> ("x" / "z") * "y" :| isNotZero "z" -- ^+ , "x" / ("y" * "z") :=> ("x" / "z") / "y" :| isNotZero "z" -- ^ TODO: 0 ^-1 check+ , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") - ("z" * "x") :=> ("w" - "z") * "x" -- TODO: handle sub :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y") -- TODO handle with power :| isConstPt "w" :| isConstPt "z" :| isNotZero "z"+ -- TODO: a + b*y :=> b * (a/b + y) :| isNotZero b+ , (("x" * "y") + ("z" * "w")) :=> "x" * ("y" + ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"+ -- , "a" * (("x" * "y") + ("z" * "w")) :=> ("a" * "x") * ("y" + ("z" / "x") * "w") :| isConstPt "a" :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"+ , (("x" * "y") - ("z" * "w")) :=> "x" * ("y" - ("z" / "x") * "w") :| isConstPt "x" :| isConstPt "z" :| isNotZero "x"+ , (("x" * "y") * ("z" * "w")) :=> ("x" * "z") * ("y" * "w") :| isConstPt "x" :| isConstPt "z"+ , "x" * "x" :=> "x" ** 2 + , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 + , "x" ** 2 + "x" * "y" :=> "x" * ("x" + "y")+ -- , "x" + "y" :=> "y" * ("x" * "y" ** (-1) + 1) :| isNotZero "y" -- GABRIEL + -- , "x" + "y" * "z" :=> "y" * ("x" * "y" ** (-1) + "z") :| isNotZero "y" -- GABRIEL + ]++-- rules for nonlinear functions +rewritesFun :: [Rule]+rewritesFun =+ [+ log (exp "x") :==: exp (log "x")+ , log (exp "x") :=> "x"+ -- , exp (log "x") :=> "x" -- :| isPositive "x" ??? exp(log(x)), x, log(exp(0))+ , log ("x" * "y") :=> log "x" + log "y" :| isConstPos "x" :| isConstPos "y"+ -- , log ("x" / "y") :=> log "x" - log "y" :| isConstPos "x" :| isConstPos "y"+ , log ("x" ** "y") :=> "y" * log "x"+ , log (powabs "x" "y") :=> "y" * log (abs "x")+ --, sqrt ("x" ** "y") :=> "x" ** ("y" / 2) :| isEven "y"+ -- , sqrt ("y" * "x") :=> sqrt "y" * sqrt "x" --+ --, sqrt ("y" / "x") :=> sqrt "y" / sqrt "x"+ , abs ("x" * "y") :=> abs "x" * abs "y" -- :| isConstPt "x"+ , 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"+ , sqrt ("x" * "x") :=> abs "x"+ ]++-- Rules that reduces redundant parameters+constReduction :: [Rule]+constReduction =+ [+ 0 + "x" :=> "x"+ -- , "x" - 0 :=> "x"+ --, 1 * "x" :=> "x"+ -- , 0 / "x" :=> 0 :| isNotZero "x"+ --, "x" - "x" :=> 0 :| isNotParam "x"+ --, "x" / "x" :=> 1 :| isNotZero "x" :| isNotParam "x"+ , "x" ** 1 :=> "x"+ , powabs "x" 1 :=> abs "x"++ -- , "x" * (1 / "x") :=> 1 :| isNotParam "x" :| isNotZero "x"+ -- , negate ("x" * "y") :=> (negate "x") * "y" :| isConstPt "x"++ , "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isPositive "x"+ , (powabs "x" "y") * (powabs "x" "z") :=> powabs "x" ("y" + "x")+ , ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isPositive "x"+ , powabs (powabs "x" "y") "z" :=> powabs "x" ("y" * "z")+ , ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isPositive "x" :| isPositive "y"++ --, "x" ** "y" * "x" ** "z" :==: "x" ** ("y" + "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"+ --, ("x" ** "y") ** "z" :==: "x" ** ("y" * "z") :| isInteger "y" :| isInteger "z" :| isNotZero "x"+ --, ("x" * "y") ** "z" :==: "x" ** "z" * "y" ** "z" :| isInteger "z" :| isNotZero "x" :| isNotZero "y"++ ]++rewritesWithConstant :: [Rule]+rewritesWithConstant =+ [+ "x" * "x" :=> "x" ** 2+ , "x" - "x" :=> 0+ , "x" / "x" :=> 1 :| isNotZero "x"+ , "x" ** "y" * "x" :=> "x" ** ("y" + 1) :| isPositive "x"+ , 1 ** "x" :=> 1+ , powabs 1 "x" :=> 1+ , log (sqrt "x") :=> 0.5 * log "x" :| isNotParam "x"+ , "x" ** (1/2) :==: sqrt "x" -- <==>+ , powabs "x" (1/2) :=> sqrt (abs "x")+ , "x" ** (1/3) :==: Fixed (Uni Cbrt "x")+ , 0 * "x" :=> 0 :| isValid "x" -- :| isNotParam "x"+ , 0 ** "x" :=> 0 :| isPositive "x"+ , powabs 0 "x" :=> 0+ , 0 - "x" :=> negate "x"+ , "x" + negate "y" :==: "x" - "y"+ ]+rewritesWithParam :: [Rule]+rewritesWithParam =+ [+ -- "x" * "x" :=> "x" ** Fixed (Param 0)+ "x" - "x" :=> Fixed (Param 0)+ , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"+ , 1 ** "x" :=> Fixed (Param 0)+ , powabs 1 "x" :=> Fixed (Param 0)+ -- , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"+ ]++rewritesSimple :: [Rule]+rewritesSimple =+ [+ "x" * "y" :=> "y" * "x"+ , "x" + "y" :=> "y" + "x"+ , ("x" ** "y") * ("x" ** "z") :=> "x" ** ("y" + "z") -- :| isPositive "x"+ , ("x" + "y") + "z" :=> "x" + ("y" + "z")+ , ("x" * "y") * "z" :=> "x" * ("y" * "z")+ , ("x" * "y") + ("x" * "z") :=> "x" * ("y" + "z")+ , "x" - ("y" + "z") :=> ("x" - "y") - "z" -- TODO: check that I don't this+ , "x" - ("y" - "z") :=> ("x" - "y") + "z" -- TODO+ , ("x" * "y") / "z" :=> ("x" / "z") * "y" :| isNotZero "z" -- TODO: inv(x) <=> x^-1 , x/y <=> x*y^-1+ , "x" * ("y" / "z") :=> ("x" / "z") * "y" :| isNotZero "z" -- ^+ , "x" / ("y" * "z") :=> ("x" / "z") / "y" :| isNotZero "z" -- ^ TODO: 0 ^-1 check+ , ("w" * "x") + ("z" * "x") :=> ("w" + "z") * "x" -- :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") - ("z" * "x") :=> ("w" - "z") * "x" -- TODO: handle sub :| isConstPt "w" :| isConstPt "z"+ , ("w" * "x") / ("z" * "y") :=> ("w" / "z") * ("x" / "y")+ , log (exp "x") :=> "x"+ , exp (log "x") :=> "x"+ , log ("x" * "y") :=> log "x" + log "y"+ , log ("x" ** "y") :=> "y" * log "x"+ , abs ("x" * "y") :=> abs "x" * abs "y"+ , abs ("x" ** "y") :=> abs "x" ** "y"+ , abs ("x" - "y") :=> abs ("y" - "x")+ , recip (recip "x") :=> "x" :| isNotZero "x"+ , "x" * "x" :=> "x" ** Fixed (Param 0)+ , "x" - "x" :=> Fixed (Param 0)+ , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x"+ , 1 ** "x" :=> Fixed (Param 0)+ , log (sqrt "x") :=> Fixed (Param 0) * log "x" :| isNotParam "x"+ ]+powabs l r = Fixed (Bin PowerAbs l r)++-- | default cost function for simplification+-- TODO:+-- num_params:+-- 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 _) = 3+myCost (Param _) = 3+myCost (Bin op l r) = 2 + l + r+myCost (Uni _ t) = 3 + t++-- all rewrite rules+rewrites :: [Rule]+rewrites = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithConstant+rewritesParams :: [Rule]+rewritesParams = rewriteBasic <> constReduction <> rewritesFun <> rewritesWithParam++-- | 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
+ src/Algorithm/Massiv/Utils.hs view
@@ -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
+ src/Algorithm/SRTree/AD.hs view
@@ -0,0 +1,559 @@+{-# language FlexibleInstances, DeriveFunctor #-}+{-# language ScopedTypeVariables #-}+{-# language RankNTypes #-}+{-# language ViewPatterns #-}+{-# language FlexibleContexts #-}+{-# language BangPatterns #-}+{-# language TypeApplications #-}+{-# language MultiWayIf #-}++-----------------------------------------------------------------------------+-- |+-- 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+ ( reverseModeArr+ , reverseModeEGraph+ , reverseModeGraph+ , forwardModeUniqueJac+ , evalCache+ ) where++import Control.Monad (forM_, foldM, when)+import Control.Monad.ST ( runST )+import Data.Bifunctor (bimap, first, second)+import qualified Data.DList as DL+import Data.Massiv.Array hiding (forM_, map, replicate, zipWith)+import qualified Data.Massiv.Array as M+import qualified Data.Massiv.Array.Unsafe as UMA+import Data.Massiv.Core.Operations (unsafeLiftArray)+import Data.SRTree.Derivative ( derivative )+import Data.SRTree.Eval+ ( SRVector, evalFun, evalOp, SRMatrix, PVector, replicateAs )+import Data.SRTree.Internal+import Data.SRTree.Print (showExpr)+import Data.SRTree.Recursion ( cataM, cata, accu )+import qualified Data.Vector as V+import Debug.Trace (trace, traceShow)+import GHC.IO (unsafePerformIO)+import qualified Data.IntMap.Strict as IntMap+import Data.List ( foldl' )+import qualified Data.Vector.Storable as VS+import Control.Scheduler +import Data.Maybe ( fromJust, isJust )+import Algorithm.EqSat.Egraph++import Control.Monad.State.Strict+import Control.Monad.Identity++--import UnliftIO.Async++import qualified Data.Map.Strict as Map++evalCache :: SRMatrix -> EGraph -> ECache -> EClassId -> VS.Vector Double -> ECache+evalCache xss egraph cache root' theta = cache'+ where+ (Sz2 _ m') = M.size xss+ m = Sz1 m'+ root = canon root'+ p = VS.length theta+ comp = M.getComp xss+ one :: Array S Ix1 Double+ one = M.replicate comp m 1++ canon rt = case _canonicalMap egraph IntMap.!? rt of+ Nothing -> error "wrong canon"+ Just rt' -> if rt == rt' then rt else canon rt'++ getNode rt' = let rt = canon rt'+ cls = _eClass egraph IntMap.! rt+ in (_best . _info) cls++ getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph+ in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n'++ ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty)+ where+ evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ evalCached rt = insertKey rt++ insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ insertKey key' = do+ let key = canon key'+ isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst)+ isCachedLocal <- gets ((key `IntMap.member`) . snd . fst)+ when (not isCachedLocal && not isCachedGlobal) $ do+ let node = getNode key+ (ev, toLocal) <- evalKey node+ modify' (insKey node ev toLocal)+ getVal key++ evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ evalKey (Var ix) = pure $ (M.computeAs S $ xss <! ix, False)+ evalKey (Const v) = pure $ (M.replicate comp m v, False)+ evalKey (Param ix) = pure $ (M.replicate comp m (theta VS.! ix), True)+ evalKey (Uni f t) = do (v, b) <- getVal t+ pure $ (M.computeAs S . M.map (evalFun f) $ v, b)+ evalKey (Bin op l r) = do (vl, bl) <- getVal l+ (vr, br) <- getVal r+ pure $ (M.computeAs S $ M.zipWith (evalOp op) vl vr, bl || br)++ insKey (Var _) _ _ s = s+ insKey (Const _) _ _ s = s+ insKey (Param _) _ _ s = s+ insKey node v toLocal ((global,local), s) =+ let k = getId node+ in if toLocal+ then ((global, IntMap.insert k v local), s)+ else ((IntMap.insert k v global, local), s)++ insertLocal k v = do (c1, c2) <- get+ put (c1, IntMap.insert k v c2)+ insertGlobal k v = do (c1, c2) <- get+ put (IntMap.insert k v c1, c2)+ getVal rt' = do let rt = canon rt'+ n = getNode rt+ case n of+ Var ix -> evalKey n+ Const v -> evalKey n+ Param ix -> evalKey n+ _ -> getFromCache rt+ getFromCache rt = do+ global <- gets ((IntMap.!? rt) . fst . fst)+ local <- gets ((IntMap.!? rt) . snd . fst)+ if | isJust global -> pure (fromJust global, False)+ | isJust local -> pure (fromJust local, True)+ | otherwise -> insertKey rt++-- reverse mode applied directly on an e-graph. Supports caching.+-- assumes root points to the loss function, so for an expression+-- f(x) and the loss (y - (f(x))^2), root will point to "^"+reverseModeEGraph :: SRMatrix -> PVector -> Maybe PVector -> EGraph -> ECache -> EClassId -> VS.Vector Double -> (Array D Ix1 Double, VS.Vector Double)+reverseModeEGraph xss ys mYErr egraph cache root' theta =+ (delay $ rootVal+ , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]]+ )+ where+ rootVal = extractCache (cache'' IntMap.!? root', localcache' IntMap.!? root')+ root = canon root'+ yErr = fromJust mYErr+ m = M.size ys+ p = VS.length theta+ comp = M.getComp xss+ one :: Array S Ix1 Double+ one = M.replicate comp m 1++ canon rt = case _canonicalMap egraph IntMap.!? rt of+ Nothing -> error "wrong canon"+ Just rt' -> if rt == rt' then rt else canon rt'++ getNode rt' = let rt = canon rt'+ cls = _eClass egraph IntMap.! rt+ in (_best . _info) cls++ getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph+ in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n'++ ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty)+ where+ evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ evalCached rt = insertKey rt++ insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ insertKey key' = do+ let key = canon key'+ isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst)+ isCachedLocal <- gets ((key `IntMap.member`) . snd . fst)+ when (not isCachedLocal && not isCachedGlobal) $ do+ let node = getNode key+ (ev, toLocal) <- evalKey node+ modify' (insKey node ev toLocal)+ getVal key++ evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool)+ evalKey (Var ix) = pure $ if | ix == -1 -> (ys, False)+ | ix == -2 -> (yErr, False)+ | otherwise -> (M.computeAs S $ xss <! ix, False)+ evalKey (Const v) = pure $ (M.replicate comp m v, False)+ evalKey (Param ix) = pure $ (M.replicate comp m (theta VS.! ix), True)+ evalKey (Uni f t) = do (v, b) <- getVal t+ pure $ (M.computeAs S . M.map (evalFun f) $ v, b)+ evalKey (Bin op l r) = do (vl, bl) <- getVal l+ (vr, br) <- getVal r+ pure $ (M.computeAs S $ M.zipWith (evalOp op) vl vr, bl || br)++ insKey (Var _) _ _ s = s+ insKey (Const _) _ _ s = s+ insKey (Param _) _ _ s = s+ insKey node v toLocal ((global,local), s) =+ let k = getId node+ in if toLocal+ then ((global, IntMap.insert k v local), s)+ else ((IntMap.insert k v global, local), s)++ insertLocal k v = do (c1, c2) <- get+ put (c1, IntMap.insert k v c2)+ insertGlobal k v = do (c1, c2) <- get+ put (IntMap.insert k v c1, c2)+ getVal rt' = do let rt = canon rt'+ n = getNode rt+ case n of+ Var ix -> evalKey n+ Const v -> evalKey n+ Param ix -> evalKey n+ _ -> getFromCache rt+ getFromCache rt = do+ global <- gets ((IntMap.!? rt) . fst . fst)+ local <- gets ((IntMap.!? rt) . snd . fst)+ if | isJust global -> pure (fromJust global, False)+ | isJust local -> pure (fromJust local, True)+ | otherwise -> insertKey rt++ extractCache (Nothing, Nothing) = error "no root info"+ extractCache (Just r, _) = r+ extractCache (_, Just r) = r++ ((cache'', localcache'), cachedGrad) = calcGrad root one `execState` ((cache', localcache), Map.empty)++ calcGrad :: Int -> Array S Ix1 Double -> State ((IntMap.IntMap (Array S Ix1 Double), IntMap.IntMap (Array S Ix1 Double)), Map.Map (SRTree Int) (Array S Ix1 Double)) ()+ calcGrad rt v = do let node = getNode rt+ case node of+ Bin op l r -> do xl <- fst <$> getVal l+ xr <- fst <$> getVal r+ (dl, dr) <- diff op v xl xr l r+ calcGrad l dl+ calcGrad r dr+ Uni f t -> do x <- fst <$> getVal t+ calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x))+ Param ix -> modify' (insertGrad v (Param ix))+ _ -> pure ()+ where+ insertGrad v k ((a, b), g) = ((a, b), Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g)++ --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double)+ diff Add dx fx gy l r = pure (dx, dx)+ diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx)+ diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx)+ diff Div dx fx gy l r = do+ let k = getId (Bin Div l r)+ v <- fst <$> getVal k+ pure (M.computeAs S $ M.zipWith (/) dx gy+ , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy))+ diff Power dx fx gy l r = do+ let k = getId (Bin Power l r)+ v <- fst <$> getVal k+ pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v+ , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v)++ diff PowerAbs dx fx gy l r = do+ let k = getId (Bin PowerAbs l r)+ v <- fst <$> getVal k+ let v2 = M.map abs fx+ v3 = M.computeAs S $ M.zipWith (*) fx gy+ pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2+ , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v)++ diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx+ dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl+ in pure (M.computeAs S $ dxl, M.computeAs S $ dxy)++ fixNaN x = if isNaN x then 0 else x+++reverseModeGraph :: SRMatrix -> PVector -> Maybe PVector -> VS.Vector Double -> Fix SRTree -> (Array D Ix1 Double, VS.Vector Double)+reverseModeGraph xss ys mYErr theta tree = (delay $ cachedVal' IntMap.! root+ , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]])+ where+ yErr = fromJust mYErr+ --ys = delay ys'+ m = M.size ys+ p = VS.length theta+ comp = M.getComp xss+ one :: Array S Ix1 Double+ one = M.replicate comp m 1+ (key2int, int2key, cachedVal, (subtract 1) -> root) = cataM leftToRight alg tree `execState` (Map.empty, IntMap.empty, IntMap.empty, 0)+ (key2int', int2key', cachedVal', cachedGrad) = calcGrad root one `execState` (key2int, int2key, cachedVal, Map.empty)++ calcGrad :: Int -> Array S Ix1 Double -> State (Map.Map (SRTree Int) Int, IntMap.IntMap (SRTree Int), IntMap.IntMap (Array S Ix1 Double), Map.Map (SRTree Int) (Array S Ix1 Double)) ()+ calcGrad key v = do node <- gets ((IntMap.! key) . _int2key)+ case node of+ Bin op l r -> do xl <- gets (getVal l)+ xr <- gets (getVal r)+ (dl, dr) <- diff op v xl xr l r+ calcGrad l dl+ calcGrad r dr+ Uni f t -> do x <- gets (getVal t)+ calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x))+ Param ix -> modify' (insertGrad v (Param ix))+ _ -> pure ()+ where+ _int2key (_, b, _, _) = b+ insertGrad v k (a, b, c, g) = (a, b, c, Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g)++ graph (a, _, _, _) = a+ insKey key ev (a, b, c, d) = (Map.insert key d a, IntMap.insert d key b, IntMap.insert d ev c, d+1)+ -- get the values from the cache+ getVal key (a, b, c, d) = c IntMap.! key+ -- maps the the struct to an integer key+ getKey key (a, b, c, d) = a Map.! key++ -- this tells the order in which we traverse the tree+ leftToRight (Uni f mt) = Uni f <$> mt;+ leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+ leftToRight (Var ix) = pure (Var ix)+ leftToRight (Param ix) = pure (Param ix)+ leftToRight (Const c) = pure (Const c)++ evalKey (Var ix) = pure $ if ix == -1+ then ys+ else if ix == -2+ then yErr+ else M.computeAs S $ xss <! ix+ evalKey (Const v) = pure $ M.replicate comp m v+ evalKey (Param ix) = pure $ M.replicate comp m (theta VS.! ix)+ evalKey (Uni f t) = M.computeAs S . M.map (evalFun f) <$> gets (getVal t)+ evalKey (Bin op l r) = M.computeAs S <$> (M.zipWith (evalOp op) <$> gets (getVal l) <*> gets (getVal r))++ alg (Var ix) = insertKey (Var ix)+ alg (Param ix) = insertKey (Param ix)+ alg (Const v) = insertKey (Const v)+ alg (Uni f t) = insertKey (Uni f t)+ alg (Bin op l r) = insertKey (Bin op l r)++ --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double)+ diff Add dx fx gy l r = pure (dx, dx)+ diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx)+ diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx)+ diff Div dx fx gy l r = do+ k <- gets (getKey (Bin Div l r))+ v <- gets (getVal k)+ pure (M.computeAs S $ M.zipWith (/) dx gy+ , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy))+ diff Power dx fx gy l r = do+ k <- gets (getKey (Bin Power l r))+ v <- gets (getVal k)+ pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v+ , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v)++ diff PowerAbs dx fx gy l r = do+ k <- gets (getKey (Bin PowerAbs l r))+ v <- gets (getVal k)+ let v2 = M.map abs fx+ v3 = M.computeAs S $ M.zipWith (*) fx gy+ pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2+ , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v)++ diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx+ dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl+ in pure (M.computeAs S $ dxl, M.computeAs S $ dxy)++ fixNaN x = if isNaN x then 0 else x++ insertKey key = do+ isCached <- gets ((key `Map.member`) . graph)+ when (not isCached) $ do+ ev <- evalKey key+ modify' (insKey key ev)+ gets (getKey key)++-- | Same as above, but using reverse mode with the tree encoded as an array, that is even faster.+reverseModeArr :: SRMatrix+ -> PVector+ -> Maybe PVector+ -> VS.Vector Double -- PVector+ -> [(Int, (Int, Int, Int, Double))] -- arity, opcode, ix, const val+ -> IntMap.IntMap Int+ -> (Array D Ix1 Double, Array S Ix1 Double)+reverseModeArr xss ys mYErr theta t j2ix =+ unsafePerformIO $ do+ fwd <- M.newMArray (Sz2 n m) 0+ partial <- M.newMArray (Sz2 n m) 0+ jacob <- M.newMArray (Sz p) 0+ val <- M.newMArray (Sz m) 0+ let+ stps = 2+ --delta = m `div` stps+ --rngs = [(i*delta, min m $ (i+1)*delta) | i <- [0..stps] ]+ (a, b) = (0, m)++ forward (a, b) fwd+ calculateYHat (a, b) fwd val+ reverseMode (a, b) fwd partial+ combine (a, b) partial jacob+ j <- UMA.unsafeFreeze (getComp xss) jacob+ v <- UMA.unsafeFreeze (getComp xss) val+ pure (delay v, j)++ where+ (Sz2 m _) = M.size xss+ p = VS.length theta+ n = length t+ toLin i j = i*m + j+ yErr = fromJust mYErr+ eps = 1e-8++ myForM_ [] _ = pure ()+ myForM_ (!x:xs) f = do f x+ myForM_ xs f+ {-# INLINE myForM_ #-}++ calculateYHat :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO ()+ calculateYHat (a, b) fwd yhat = myForM_ [a..b-1] $ \i -> do+ vi <- UMA.unsafeRead fwd (0 :. i)+ UMA.unsafeWrite yhat i vi+ {-# INLINE calculateYHat #-}++ forward :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> IO ()+ forward (a, b) fwd = do+ let t' = Prelude.reverse t+ myForM_ t' makeFwd+ where+ makeFwd (j, (0, 0, ix, _)) =+ do let j' = j2ix IntMap.! j+ myForM_ [a..b-1] $ \i -> do+ --let val = xss M.! (i :. ix)+ UMA.unsafeWrite fwd (j' :. i) $ case ix of+ (-1) -> ys M.! i+ (-2) -> yErr M.! i+ _ -> xss M.! (i :. ix)+ makeFwd (j, (0, 1, ix, _)) = do let j' = j2ix IntMap.! j+ v = theta VS.! ix+ myForM_ [a..b-1] $ \i -> do+ UMA.unsafeWrite fwd (j' :. i) v+ makeFwd (j, (0, 2, _, x)) = do let j' = j2ix IntMap.! j+ myForM_ [a..b-1] $ \i -> do+ UMA.unsafeWrite fwd (j' :. i) x+ makeFwd (j, (1, f, _, _)) = do let j' = j2ix IntMap.! j+ j2 = j2ix IntMap.! (2*j + 1)+ myForM_ [a..b-1] $ \i -> do+ v <- UMA.unsafeRead fwd (j2 :. i)+ UMA.unsafeWrite fwd (j' :. i) (evalFun (toEnum f) v)+ makeFwd (j, (2, op, _, _)) = do let j' = j2ix IntMap.! j+ j2 = j2ix IntMap.! (2*j + 1)+ j3 = j2ix IntMap.! (2*j + 2)+ myForM_ [a..b-1] $ \i -> do+ l <- UMA.unsafeRead fwd (j2 :. i)+ r <- UMA.unsafeRead fwd (j3 :. i)+ UMA.unsafeWrite fwd (j' :. i) (evalOp (toEnum op) l r)+ makeFwd _ = pure ()+ {-# INLINE makeFwd #-}+ {-# INLINE forward #-}++ reverseMode :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix2 Double -> IO ()+ reverseMode (a, b) fwd partial =+ do myForM_ [a..b-1] $ \i -> UMA.unsafeWrite partial (0 :. i) 1+ myForM_ t makeRev+ where+ makeRev (j, (1, f, _, _)) = do let dxj = j2ix IntMap.! j+ vj = j2ix IntMap.! (2*j + 1)+ myForM_ [a..b-1] $ \i -> do+ v <- UMA.unsafeRead fwd (vj :. i)+ dx <- UMA.unsafeRead partial (dxj :. i)+ --let val = dx * derivative (toEnum f) v+ UMA.unsafeWrite partial (vj :. i) (dx * derivative (toEnum f) v)+ makeRev (j, (2, op, _, _)) = do let dxj = j2ix IntMap.! j+ lj = j2ix IntMap.! (2*j + 1)+ rj = j2ix IntMap.! (2*j + 2)+ myForM_ [a..b-1] $ \i -> do+ l <- UMA.unsafeRead fwd (lj :. i)+ r <- UMA.unsafeRead fwd (rj :. i)+ dx <- UMA.unsafeRead partial (dxj :. i)+ let (dxl, dxr) = diff (toEnum op) dx l r+ UMA.unsafeWrite partial (lj :. i) dxl+ UMA.unsafeWrite partial (rj :. i) dxr+ makeRev _ = pure ()+ {-# INLINE makeRev #-}+ {-# INLINE reverseMode #-}++ --f(x)^g(x)+ --d f(x)^g(x) / d f(x) = f(x)^(g(x)-1)+ -- f(x) + g(x) = 1, 1+ -- f(x) - g(x) = 1, -1+ -- f(x) * g(x) = g(x), f(x)+ -- f(x) / g(x) = 1/g(x), -f(x)/g(x)^2+ -- f(x) ^ g(x) = g(x) * f(x) ^ (g(x) - 1), f(x) ^ g(x) * log f(x)+ -- |f(x)| ^ g(x) = g(x) * |f(x)| ^ (g(x) - 2) * f(x), |f(x)| ^ g(x) * log |f(x)|++ -- |f(x)| ^ g(x) = exp (log |f(x)| * g(x))+ -- => |f(x)| ^ (g(x) - 1) * g(x)+ -- => |f(x)| ^ g(x) * log |f(x)| * 1++ fixNaN x | isNaN x = 0+ | otherwise = x++ diff :: Op -> Double -> Double -> Double -> (Double, Double)+ diff Add dx fx gy = (dx, dx)+ diff Sub dx fx gy = (dx, negate dx)+ diff Mul dx fx gy = (dx * gy, dx * fx)+ diff Div dx fx gy = (dx / gy, dx * (negate fx / (gy * gy)))+ --diff Power dx fx gy = (fixNaN $ dx * ((fx+eps)**gy - fx**gy)/eps, fixNaN $ dx * (fx**(gy+eps) - fx**gy)/eps)+ --diff PowerAbs dx fx gy = (fixNaN $ dx * (abs (fx+eps)**gy - abs fx**gy)/eps, fixNaN $ dx * (abs fx**(gy+eps) - abs fx**gy)/eps)+ {--}+ diff Power 0 _ _ = (0, 0)+ diff Power dx 0 0 = (0, 0)+ diff Power dx fx 0 = (0, fixNaN $ dx * log fx)+ diff Power dx 0 gy = (fixNaN $ dx * gy * if gy < 1 then eps ** (gy - 1) else 0+ , 0) --dx * fx ** gy * log fx)+ diff Power dx fx gy = (fixNaN $ dx * gy * fx ** (gy - 1), fixNaN $ dx * fx ** gy * log fx)++ diff PowerAbs 0 fx gy = (0, 0)+ diff PowerAbs 0 0 0 = (0, 0)+ diff PowerAbs dx fx 0 = (0, fixNaN $ dx * log (abs fx))+ diff PowerAbs dx 0 gy = (0, fixNaN $ dx * if gy < 0 then eps ** gy else 0)+ diff PowerAbs dx fx gy = (fixNaN $ dx * gy * fx * abs fx ** (gy - 2), fixNaN $ dx * abs fx ** gy * log (abs fx))+ {--}+ diff AQ dx fx gy = let dxl = recip ((sqrt . (+1)) (gy * gy))+ dxy = fx * gy * (dxl^3) -- / (sqrt (gy*gy + 1))+ in (dxl * dx, dxy * dx)++ {-# INLINE diff #-}++ combine :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO ()+ combine (lo, hi) partial jacob = myForM_ t makeJacob+ where+ makeJacob (j, (0, 1, ix, _)) = do val <- UMA.unsafeRead jacob ix+ let j' = j2ix IntMap.! j+ addI a b acc = do v2 <- UMA.unsafeRead partial (b :. a)+ pure (v2 + acc)+ acc <- foldM (\a i -> addI i j' a) val [lo..hi-1]+ UMA.unsafeWrite jacob ix acc+ makeJacob _ = pure ()+ {-# INLINE combine #-}++-- | 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)
+ src/Algorithm/SRTree/ConfidenceIntervals.hs view
@@ -0,0 +1,448 @@+{-# 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+ ( minimizeNLL, 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 xss' -- $ 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 MSE y = y+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 PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> PVector -> [CI] -> Double -> [ProfileT]+getAllProfiles ptype dist mYerr xss ys tree theta stdErr estCIs alpha = reverse (getAll 0 [])+ where+ (A.Sz k) = A.size theta+ (A.Sz n) = A.size ys+ tau_max = sqrt $ quantile (fDistribution k (n - k)) (1 - 0.01)+ tau_max' = sqrt $ quantile (fDistribution k (n - k)) (1 - alpha)++ profFun ix = case ptype of+ Bates -> getProfile dist mYerr xss ys tree theta (stdErr A.! ix) tau_max ix+ ODE -> getProfileODE dist mYerr xss ys tree theta (stdErr A.! ix) (estCIs !! ix) tau_max ix+ Constrained -> getProfileCnstr dist mYerr xss ys tree theta (stdErr A.! ix) tau_max' ix++ getAll ix acc | ix == k = acc+ | otherwise = case profFun ix of+ Left t -> getAllProfiles ptype dist mYerr xss ys tree t stdErr estCIs alpha+ Right p -> getAll (ix + 1) (p : acc)++-- calculates the profile likelihood of a single parameter +getProfile :: Distribution+ -> Maybe PVector+ -> SRMatrix+ -> PVector+ -> Fix SRTree+ -> PVector+ -> Double+ -> Double+ -> Int+ -> Either PVector ProfileT+getProfile dist mYerr xss ys tree theta stdErr_i tau_max ix+ | stdErr_i == 0.0 = pure $ ProfileT (A.fromList compMode [-tau_max, tau_max]) (A.fromLists' compMode [theta', theta']) (theta A.! ix) (const (theta A.! ix)) (const tau_max)+ | 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 mYerr xss ys tree theta_opt+ (theta_opt, _, _) = minimizeNLL dist mYerr 100 xss ys tree theta+ optTh = theta_opt A.! ix+ minimizer = minimizeNLLWithFixedParam dist mYerr 100 xss ys tree ix++ -- 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 mYerr xss ys tree theta_t) A.! ix+ zvs = snd $ gradNLL dist mYerr xss ys tree theta_t+ inv_slope' = min 4.0 . max 0.0625 . abs $ (tau / (stdErr_i * zv))+ nll_cond = nll dist mYerr xss ys tree theta_t+ acc' = if nll_cond == nll_opt || ( (not.null) taus && tau == head taus ) || isNaN tau+ then acc+ else (tau:taus, theta_t:thetas)+ tau = signum delta * sqrt (2*nll_cond - 2*nll_opt)++-- Based on https://insysbio.github.io/LikelihoodProfiler.jl/latest/+-- Borisov, Ivan, and Evgeny Metelkin. "Confidence intervals by constrained optimization—An algorithm and software package for practical identifiability analysis in systems biology." PLOS Computational Biology 16.12 (2020): e1008495.+getProfileCnstr :: Distribution+ -> Maybe PVector+ -> SRMatrix+ -> PVector+ -> Fix SRTree+ -> PVector+ -> Double -> Double+ -> Int+ -> Either PVector ProfileT+getProfileCnstr dist mYerr xss ys tree theta stdErr_i tau_max ix+ | stdErr_i == 0.0 = pure $ ProfileT taus thetas theta_i (const theta_i) (const tau_max)+ | otherwise = pure $ ProfileT taus thetas theta_i tau2theta (const tau_max)+ where+ taus = A.fromList compMode [-tau_max, tau_max]+ theta' = A.toList theta+ thetas = A.fromLists' compMode [theta', theta']+ theta_i = theta A.! ix+ getPoint = getEndPoint dist mYerr xss ys tree theta tau_max ix+ leftPt = getPoint True+ rightPt = getPoint False+ tau2theta tau = if tau < 0 then leftPt else rightPt++getEndPoint :: Distribution -> Maybe PVector -> A.Array A.S Ix2 Double -> A.Array A.S A.Ix1 Double -> Fix SRTree -> A.Array A.S A.Ix1 Double -> Double -> Int -> Bool -> Double+getEndPoint dist mYerr xss ys tree theta tau_max ix isLeft =+ case minimizeAugLag problem (A.toStorableVector theta_opt) of+ Right sol -> solutionParams sol VS.! ix+ Left e -> traceShow e $ theta_opt A.! ix+ where+ (A.Sz1 n) = A.size theta++ (theta_opt, _, _) = minimizeNLL dist mYerr 100 xss ys tree theta+ nll_opt = nll dist mYerr xss ys tree theta_opt+ loss_crit = nll_opt + tau_max++ loss = subtract loss_crit . nll dist mYerr 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 PVector+ -> SRMatrix+ -> PVector+ -> Fix SRTree+ -> PVector+ -> Double+ -> CI+ -> Double+ -> Int+ -> Either PVector ProfileT+getProfileODE dist mYerr xss ys tree theta stdErr_i estCI tau_max ix+ | 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 = (\(x, _, _) -> x) . minimizeNLL dist mYerr 100 xss ys tree+ grader = snd . gradNLL dist mYerr xss ys tree+ theta_opt = minimizer theta+ theta' = A.toList theta+ nll_opt = nll dist mYerr xss ys tree theta_opt+ optTh = theta_opt A.! ix+ p' = p+1+ (A.Sz1 p) = A.size theta+ --sErr = fromMaybe 1 mSErr+ getHess = hessianNLL dist mYerr xss ys tree++ odeFun gamma _ u =+ let grad = grader u+ w = hessianNLL dist mYerr xss ys tree u+ m = A.makeArray compMode (A.Sz (p' :. p'))+ (\ (i :. j) -> if | i<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 mYerr xss ys tree $ snd t+ z = signum ((snd t A.! ix) - optTh) * sqrt (2 * nll_i - 2 * nll_opt)+ in if z == 0 || isNaN z then ([], []) else ([z], [snd t])++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 PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> BasicStats+getStatsFromModel dist mYerr xss ys tree theta = MkStats cov corr stdErr+ where+ (A.Sz1 k) = A.size theta+ (A.Sz1 n) = A.size ys+ 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 mYerr xss ys tree theta+ -- cov = catch (unsafePerformIO (invChol hess)) (\e -> trace "cov NegDef" $ pure ident)+ fexcept :: (A.PrimMonad m, A.MonadThrow m, A.MonadIO m) => A.SomeException -> m SRMatrix+ fexcept e = trace "cov NegDef" $ pure ident+ cov = unsafePerformIO $ catch (invChol hess) fexcept++ 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
+ src/Algorithm/SRTree/Likelihoods.hs view
@@ -0,0 +1,635 @@+{-# LANGUAGE ViewPatterns #-}+{-# LANGUAGE TypeApplications #-}++-----------------------------------------------------------------------------+-- |+-- 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+ , buildNLL+ , buildNLLEGraph+ , gradNLL+ , gradNLLArr+ , gradNLLGraph+ , gradNLLEGraph+ , fisherNLL+ , getSErr+ , hessianNLL+ , tree2arr+ )+ where++import Algorithm.SRTree.AD ( reverseModeArr, reverseModeGraph, reverseModeEGraph )+import Data.Massiv.Array hiding (all, map, read, replicate, tail, take, zip)+import qualified Data.Massiv.Array as M+import qualified Data.Massiv.Array.Mutable as Mut+import Data.Maybe (fromMaybe)+import Data.SRTree+import Data.SRTree.Recursion ( cata, accu )+import Data.SRTree.Derivative (deriveByParam, deriveByVar, derivative)+import Data.SRTree.Eval+import qualified Data.IntMap.Strict as IntMap+import qualified Data.Vector.Storable as VS+import GHC.IO (unsafePerformIO)+import Data.Maybe++import Debug.Trace+import Data.SRTree.Print+import Algorithm.EqSat.Egraph+import Algorithm.EqSat.Simplify+import Algorithm.EqSat.Build+import Control.Monad.State.Strict+import Control.Monad.Identity++import Data.SRTree.Print++-- | Supported distributions for negative log-likelihood+-- MSE refers to mean squared error+-- HGaussian is Gaussian with heteroscedasticity, where the error should be provided+data Distribution = MSE | Gaussian | HGaussian | Bernoulli | Poisson | ROXY | LOG10+ deriving (Show, Read, Enum, Bounded, Eq)++-- | Sum-of-square errors or Sum-of-square residues+sse :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double+sse xss ys tree theta = err+ where+ (Sz m) = M.size ys+ cmp = getComp xss+ yhat = evalTree xss theta tree+ err = M.sum $ (delay ys - yhat) ^ (2 :: Int)++sseError :: SRMatrix -> PVector -> PVector -> Fix SRTree -> PVector -> Double+sseError xss ys yErr tree theta = err+ where+ (Sz m) = M.size ys+ cmp = getComp xss+ yhat = evalTree xss theta tree+ err = M.sum $ ((delay ys - yhat) ^ (2 :: Int) / (delay yErr))++-- | Total Sum-of-squares+sseTot :: SRMatrix -> PVector -> Fix SRTree -> PVector -> Double+sseTot xss ys tree theta = err+ 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 PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> Double++-- | Mean Squared error (not a distribution)+nll MSE _ xss ys t theta = mse xss ys t theta++nll LOG10 _ xss ys t theta = M.sum $ (M.map (logBase 10) $ (f (delay ys) / f yhat)) ^ (2 :: Int)+ where+ yhat = evalTree xss theta t+ (Sz m) = M.size ys+ f :: Array D Ix1 Double -> Array D Ix1 Double+ f z = (z + M.map (\zi -> sqrt (zi^2 + 1e-10)) z)+ -- log ys - log y = log (ys/y)++-- | Gaussian distribution, theta must contain an additional parameter corresponding+-- to variance.+nll Gaussian mYerr xss ys t theta+ | nParams == (p'-1) = error "For Gaussian distribution theta must contain the variance as its last value."+ | otherwise = 0.5*(sse xss ys t theta / s + m*log (2*pi*s))+ where+ s = sqrt $ mse xss ys t theta -- theta M.! (p' - 1)+ (Sz m') = M.size ys + (Sz p') = M.size theta+ nParams = countParamsUniq t+ m = fromIntegral m'+ p = fromIntegral p'++-- | Gaussian with heteroscedasticity, it needs a valid mYerr+nll HGaussian mYerr xss ys t theta =+ case mYerr of+ Nothing -> error "For HGaussian, you must provide the measured error for the target variable."+ Just yErr -> 0.5*(sseError xss ys yErr t theta + M.sum (M.map (log . (2*) . (pi*)) yErr))+ where+ (Sz m') = M.size ys+ (Sz p') = M.size theta+ m = fromIntegral m'+ p = fromIntegral p'++-- | Bernoulli distribution of f(x; theta) is, given phi = 1 / (1 + exp (-f(x; theta))),+-- y log phi + (1-y) log (1 - phi), assuming y \in {0,1}+nll Bernoulli _ xss ys tree theta+ | notValid ys = error "For Bernoulli distribution the output must be either 0 or 1."+ | otherwise = M.sum $ (M.map (1-) (delay ys)) * yhat + log (M.map (1+) $ exp (M.map negate yhat))+ where+ (Sz m) = M.size ys+ yhat = evalTree xss theta tree+ notValid = M.any (\x -> x /= 0 && x /= 1)++nll Poisson _ xss ys tree theta+ | notValid ys = error "For Poisson distribution the output must be non-negative."+ -- | M.any isNaN yhat = error $ "NaN predictions " <> show theta+ | otherwise = negate . M.sum $ ys' * yhat - ys' * log ys' - exp yhat+ where+ ys' = delay ys+ yhat = evalTree xss theta tree+ notValid = M.any (<0)++nll ROXY mYerr xss ys tree theta+ | isNothing mYerr = error "Can't calculate ROXY nll without x,y-errors."+ | p < num_params + 3 = error "We need 3 additional parameters for ROXY."+ | n /= 1 && n/=5 = error "For ROXY dataset must contain a single variable, or 1 variable + 4 cached data."+ | otherwise = if isNaN negLL then (1.0/0.0) else negLL+ where+ (Sz p') = M.size theta+ (Sz2 m n) = M.size xss+ p = fromIntegral p'+ num_params = countParamsUniq tree++ x0 = xss <! 0+ logX = xss <! 1+ logY = xss <! 2+ logXErr = xss <! 3+ logYErr = xss <! 4+++ yErr = fromJust mYerr+ one = M.replicate compMode (Sz m) 1+ zero = M.replicate compMode (Sz m) 0++ (sig, mu_gauss, w_gauss) = (theta ! num_params, theta ! (num_params + 1), theta ! (num_params + 2))++ applyDer :: Op -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double -> Array D Ix1 Double+ applyDer Add l dl r dr = dl+dr+ applyDer Sub l dl r dr = dl-dr+ applyDer Mul l dl r dr = l*dr + r*dl+ applyDer Div l dl r dr = (dl*r - dr*l) / (r^2)+ applyDer Power l dl r dr = l ** (r.-1) * (r*dl + l * log l * dr)+ applyDer PowerAbs l dl r dr = (abs l ** r) * (dr * log (abs l) + r * dl / l)+ applyDer AQ l dl r dr = ((1 +. r*r) * dl - l * r * dr) / M.map (**1.5) (1 +. r*r)++ (yhat, grad) = cata alg tree+ where+ alg (Var ix) = (x0, one)+ alg (Param ix) = (M.replicate compMode (Sz m) (theta M.! ix), zero)+ alg (Const x) = (M.replicate compMode (Sz m) x, zero)+ alg (Uni f (val, der)) = (M.map (evalFun f) val, M.map (derivative f) val * der)+ alg (Bin op (valL, derL) (valR, derR)) = (M.zipWith (evalOp op) valL valR, applyDer op valL derL valR derR)++ f = M.map (logBase 10) (abs yhat)+ fprime = grad / (log 10 *. yhat) * x0 .* log 10++ -- nll+ w_gauss2 = w_gauss ^ 2+ s2 = delay $ logYErr .+ sig^2+ den = fprime ^ 2 .* w_gauss2 * logXErr + s2 * (w_gauss2 +. logXErr)++ neglogP = log (2 * pi)+ +. log den+ + (w_gauss2 *. (f - logY) * (f - logY)+ + logXErr * (fprime * (mu_gauss -. logX) + f - logY)^2+ + s2 * (logX .- mu_gauss)^2) / den+ negLL = 0.5 * M.sum neglogP++-- WARNING: pass tree with parameters+-- TODO: handle error similar to ROXY+buildNLL MSE m tree = ((tree - var (-1)) ** 2) / constv m+buildNLL LOG10 m tree = (((log (y / tree')) / log 10) ** 2) / constv m+ where+ tree' = (tree + sqrt(tree^2 + 1e-10))+ y = (var (-1) + sqrt(var (-1) ^ 2 + 1e-10))++buildNLL Gaussian m tree = (square(tree - var (-1)) / square (param p)) + log ((square (param p)))+ where+ square = Fix . Uni Square+ p = countParamsUniq tree+buildNLL HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2))+buildNLL Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree+buildNLL Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree+buildNLL ROXY m tree = neglogP+ where+ p = countParamsUniq tree+ f = log (abs tree) / log 10+ fprime = deriveByVar 0 tree / (log 10 * tree) * var 0 * log 10+ logX = var 1+ logY = var 2+ logXErr = var 3+ logYErr = var 4+ sig = param p+ mu_gauss = param (p+1)+ w_gauss = param (p+2)+ w_gauss2 = w_gauss ** 2+ s2 = logYErr + sig ** 2+ den = fprime ** 2 * w_gauss2 * logXErr + s2 * (w_gauss2 + logXErr)+ neglogP = log (2*pi)+ + log den+ + ( w_gauss2 * (f - logY) * (f - logY)+ + logXErr * (fprime *(mu_gauss - logX) + f - logY)**2+ + s2 * (logX - mu_gauss) ** 2+ ) / den++buildNLLEGraph MSE m egraph root = runIdentity $ addToEg `runStateT` egraph+ where+ addToEg :: EGraphST Identity EClassId+ addToEg = do v <- add myCost (Var (-1))+ c1 <- add myCost (Const 2)+ c2 <- add myCost (Const m)+ x <- add myCost (Bin Sub root v)+ y <- add myCost (Bin Power x c1)+ add myCost (Bin Div y c2)+buildNLLEGraph LOG10 m egraph root = runIdentity $ addToEg `runStateT` egraph+ where+ addToEg :: EGraphST Identity EClassId+ addToEg = do v <- add myCost (Var (-1))+ c1 <- add myCost (Const 2)+ c2 <- add myCost (Const m)+ c3 <- add myCost (Const 10)+ c4 <- add myCost (Const 1e-10)+ -- log (x + sqrt (x^2 + 1)) / log 10+ log10 <- add myCost (Uni Log c3)+ t2 <- add myCost (Uni Square root)+ t2p1 <- add myCost (Bin Add t2 c4)+ sqt <- add myCost (Uni Sqrt t2p1)+ tpt <- add myCost (Bin Add root sqt)++ -- same with y+ y2 <- add myCost (Uni Square v)+ y2p1 <- add myCost (Bin Add y2 c4)+ sqy <- add myCost (Uni Sqrt y2p1)+ ypy <- add myCost (Bin Add v sqy)++ tptypy <- add myCost (Bin Div ypy tpt)++ logy <- add myCost (Uni Log tptypy)+ log10y <- add myCost (Bin Div logy log10)++ --x <- add myCost (Bin Sub log10t v)+ y <- add myCost (Bin Power tptypy c1)+ add myCost (Bin Div y c2)++buildNLLEGraph Gaussian m egraph root = runIdentity (addToEg `runStateT` egraph)+ where+ p = countParamsUniqEg egraph root+ addToEg :: EGraphST Identity EClassId+ addToEg = do v <- add myCost (Var (-1))+ p <- add myCost (Param p)+ sp <- add myCost (Uni Square p)+ lsp <- add myCost (Uni Log sp)+ d <- add myCost (Bin Sub root v)+ sd <- add myCost (Uni Square d)+ x <- add myCost (Bin Div sd sp)+ add myCost (Bin Add x lsp)++buildNLLEGraph HGaussian m egraph root = runIdentity $ addToEg `runStateT` egraph+ where+ addToEg :: EGraphST Identity EClassId+ addToEg = do v1 <- add myCost (Var (-1))+ v2 <- add myCost (Var (-2))+ c1 <- add myCost (Const (2*pi))+ c2 <- add myCost (Const m)+ x <- add myCost (Bin Sub root v1)+ y <- add myCost (Uni Square x)+ z <- add myCost (Bin Div y v2)+ w <- add myCost (Bin Mul c1 v2)+ lw <- add myCost (Uni Log w)+ p <- add myCost (Bin Mul c2 lw)+ add myCost (Bin Add z p)+++buildNLLEGraph Poisson m egraph root = runIdentity $ addToEg `runStateT` egraph+ where+ addToEg :: EGraphST Identity EClassId+ addToEg = do v1 <- add myCost (Var (-1))+ lv <- add myCost (Uni Log v1)+ x <- add myCost (Bin Mul v1 lv)+ y <- add myCost (Uni Exp root)+ z <- add myCost (Bin Add x y)+ vt <- add myCost (Bin Mul v1 root)+ add myCost (Bin Sub z vt)++buildNLLEGraph Bernoulli m egraph root = runIdentity $ addToEg `runStateT` egraph+ where+ addToEg :: EGraphST Identity EClassId+ addToEg = do v <- add myCost (Var (-1))+ c1 <- add myCost (Const 1)+ c2 <- add myCost (Const (-1))+ mr <- add myCost (Bin Mul c2 root)+ er <- add myCost (Uni Exp mr)+ er1 <- add myCost (Bin Add c1 er)+ ler1 <- add myCost (Uni Log er1)+ v1 <- add myCost (Bin Sub c1 v)+ v1r <- add myCost (Bin Mul v1 root)+ add myCost (Bin Add ler1 v1r)++buildNLLEGraph ROXY m egraph root = error "ROXY not supported with cache"++-- | Prediction for different distributions+predict :: Distribution -> Fix SRTree -> PVector -> SRMatrix -> SRVector+predict MSE tree theta xss = evalTree xss theta tree+predict LOG10 tree theta xss = evalTree xss theta tree+predict Gaussian tree theta xss = evalTree xss theta tree+predict Bernoulli tree theta xss = logistic $ evalTree xss theta tree+predict Poisson tree theta xss = exp $ evalTree xss theta tree+predict ROXY tree theta xss = evalTree xss theta tree++-- | Gradient of the negative log-likelihood+gradNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (Double, SRVector)+gradNLL dist mYerr xss ys tree theta = (f, delay grad) -- gradNLLArr dist xss ys mYerr treeArr j2ix (toStorableVector theta)+ where+ grad :: PVector+ grad = M.fromList M.Seq [finitediff ix | ix <- [0..p-1]]+ (Sz p) = M.size theta++ disturb :: Int -> PVector+ disturb ix = M.fromList M.Seq $ Prelude.zipWith (\iy v -> if iy==ix then (v+eps) else v) [0..] (M.toList theta)+ eps :: Double+ eps = 1e-8+ f = (/ fromIntegral m) . M.sum . M.map (^2) $ (predict MSE tree theta xss) - delay ys+ finitediff ix = let t1 = disturb ix+ f' = (/ fromIntegral m) . M.sum . M.map (^2) $ (predict MSE tree t1 xss) - ys'+ in (f' - f)/eps+ (Sz2 m _) = M.size xss+ tree' = buildNLL dist (fromIntegral m) tree+ treeArr = IntMap.toAscList $ tree2arr tree'+ j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..]+ flog :: Array D Ix1 Double -> Array D Ix1 Double+ flog z = M.map (logBase 10) (z + M.map sqrt (z^2 + 1e-10))+ ys' = (if dist==LOG10 then id else id) (delay ys)+++nanTo0 x = x -- if isNaN x || isInfinite x then 0 else x+{-# INLINE nanTo0 #-}++-- | Gradient of the negative log-likelihood+gradNLLArr MSE xss ys mYerr tree j2ix theta =+ (M.sum yhat, delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr LOG10 xss ys mYerr tree j2ix theta =+ (M.sum yhat, delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr Gaussian xss ys mYerr tree j2ix theta =+ (M.sum yhat, delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr Bernoulli xss ys mYerr tree j2ix theta+ | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."+ | otherwise = (M.sum yhat, delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr Poisson xss ys mYerr tree j2ix theta+ | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."+ | otherwise = (M.sum yhat, delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad+gradNLLArr ROXY xss ys mYerr tree j2ix theta =+ ((*0.5) $ M.sum yhat, M.map (*(0.5)) $ delay grad')+ where+ (yhat, grad) = reverseModeArr xss ys mYerr theta tree j2ix+ grad' = M.map nanTo0 grad++-- | Gradient of the negative log-likelihood+gradNLLGraph MSE xss ys mYerr tree theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph LOG10 xss ys mYerr tree theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph Gaussian xss ys mYerr tree theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph Bernoulli xss ys mYerr tree theta+ | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."+ | otherwise = (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph Poisson xss ys mYerr tree theta+ | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."+ | otherwise = (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad+gradNLLGraph ROXY xss ys mYerr tree theta =+ ((*0.5) $ M.sum yhat, VS.map (*(0.5)) $ grad')+ where+ (yhat, grad) = reverseModeGraph xss ys mYerr theta tree+ grad' = VS.map nanTo0 grad++-- | e-graph support+gradNLLEGraph MSE xss ys mYerr egraph cache root theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad+gradNLLEGraph LOG10 xss ys mYerr egraph cache root theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad+ ys' :: PVector+ ys' = M.computeAs M.S $ M.map (logBase 10) (delay ys + M.map sqrt (delay ys^2 + 1e-10))+gradNLLEGraph Gaussian xss ys mYerr egraph cache root theta =+ (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad+gradNLLEGraph Bernoulli xss ys mYerr egraph cache root theta+ | M.any (\x -> x /= 0 && x /= 1) ys = error "For Bernoulli distribution the output must be either 0 or 1."+ | otherwise = (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad+gradNLLEGraph Poisson xss ys mYerr egraph cache root theta+ | M.any (<0) ys = error "For Poisson distribution the output must be non-negative."+ | otherwise = (M.sum yhat, grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad+gradNLLEGraph ROXY xss ys mYerr egraph cache root theta =+ ((*0.5) $ M.sum yhat, VS.map (*(0.5)) $ grad')+ where+ (yhat, grad) = reverseModeEGraph xss ys mYerr egraph cache root theta+ grad' = VS.map nanTo0 grad++-- | Fisher information of negative log-likelihood+fisherNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRVector+fisherNLL ROXY mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff+ where+ cmp = getComp xss+ (Sz m) = M.size ys+ (Sz p) = M.size theta+ f = nll ROXY mYerr xss ys tree theta+ eps = 1e-6+ finiteDiff ix = unsafePerformIO $ do+ theta' <- Mut.thaw theta+ v <- Mut.readM theta' ix+ Mut.writeM theta' ix (v + eps)+ thetaPlus <- Mut.freezeS theta'+ Mut.writeM theta' ix (v - eps)+ thetaMinus <- Mut.freezeS theta'+ let fPlus = nll ROXY mYerr xss ys tree thetaPlus+ fMinus = nll ROXY mYerr xss ys tree thetaMinus+ pure $ (fPlus + fMinus - 2*f)/(eps*eps)+fisherNLL Gaussian mYerr xss ys tree theta = makeArray cmp (Sz p) finiteDiff+ where+ cmp = getComp xss+ (Sz m) = M.size ys+ (Sz p) = M.size theta+ f = nll Gaussian mYerr xss ys tree theta+ eps = 1e-6+ finiteDiff ix = unsafePerformIO $ do+ theta' <- Mut.thaw theta+ v <- Mut.readM theta' ix+ Mut.writeM theta' ix (v + eps)+ thetaPlus <- Mut.freezeS theta'+ Mut.writeM theta' ix (v - eps)+ thetaMinus <- Mut.freezeS theta'+ let fPlus = nll Gaussian mYerr xss ys tree thetaPlus+ fMinus = nll Gaussian mYerr xss ys tree thetaMinus+ pure $ (fPlus + fMinus - 2*f)/(eps*eps)+fisherNLL dist mYerr xss ys tree theta = makeArray cmp (Sz p) build+ where+ build ix = let dtdix = deriveByParam ix t'+ d2tdix2 = deriveByParam ix dtdix + f' = eval dtdix + f'' = eval d2tdix2 + in M.sum $ phi' * f'^2 - res * f''+ --case dist of+ -- Gaussian -> M.sum . (/delay (theta M.! (p-1))) $ phi' * f'^2 - res * f''+ -- _ -> M.sum $ phi' * f'^2 - res * f''+ cmp = getComp xss + (Sz m) = M.size ys+ (Sz p) = M.size theta+ t' = fst $ floatConstsToParam tree+ eval = evalTree xss theta+ yhat = eval t'+ res = delay ys - phi+ yErr = case mYerr of+ Nothing -> M.replicate (getComp xss) (Sz m) est+ Just e -> e+ est = fromIntegral (m - p)++ (phi, phi') = case dist of+ MSE -> (yhat, M.replicate compMode (Sz m) 1)+ Gaussian -> (yhat, M.replicate compMode (Sz m) 1)+ Bernoulli -> (logistic yhat, phi*(M.replicate compMode (Sz m) 1 - phi))+ Poisson -> (exp yhat, phi)++-- | Hessian of negative log-likelihood+--+-- Note, though the Fisher is just the diagonal of the return of this function+-- it is better to keep them as different functions for efficiency+hessianNLL :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> Fix SRTree -> PVector -> SRMatrix+hessianNLL ROXY mYerr xss ys tree theta = undefined+hessianNLL dist mYerr xss ys tree theta = makeArray cmp (Sz (p :. p)) build+ where+ build (ix :. iy) = let dtdix = deriveByParam ix t' + dtdiy = deriveByParam iy t' + d2tdixy = deriveByParam iy dtdix+ fx = eval dtdix + fy = eval dtdiy + fxy = eval d2tdixy + in case dist of+ Gaussian -> M.sum . (/delay yErr) $ phi' * fx * fy - res * fxy+ _ -> M.sum $ phi' * fx * fy - res * fxy++ cmp = getComp xss+ (Sz m) = M.size ys+ (Sz p) = M.size theta+ t' = tree -- relabelParams tree -- $ floatConstsToParam tree+ eval = evalTree xss theta+ yErr = case mYerr of+ Nothing -> M.replicate compMode (Sz m) est+ Just e -> e+ est = fromIntegral (m - p)+ yhat = eval t'+ res = delay ys - phi++ (phi, phi') = case dist of+ MSE -> (yhat, M.replicate cmp (Sz m) 1)+ LOG10 -> (yhat, M.replicate cmp (Sz m) 1)+ Gaussian -> (yhat, M.replicate cmp (Sz m) 1)+ Bernoulli -> (logistic yhat, phi*(M.replicate cmp (Sz m) 1 - phi))+ Poisson -> (exp yhat, phi)++tree2arr :: Fix SRTree -> IntMap.IntMap (Int, Int, Int, Double)+tree2arr tree = IntMap.fromList listTree+ where+ height = cata alg+ where+ alg (Var ix) = 1+ alg (Const x) = 1+ alg (Param ix) = 1+ alg (Uni _ t) = 1 + t+ alg (Bin _ l r) = 1 + max l r+ listTree = accu indexer convert tree 0++ indexer (Var ix) iy = Var ix+ indexer (Const x) iy = Const x+ indexer (Param ix) iy = Param ix+ indexer (Bin op l r) iy = Bin op (l, 2*iy+1) (r, 2*iy+2)+ indexer (Uni f t) iy = Uni f (t, 2*iy+1)++ convert (Var ix) iy = [(iy, (0, 0, ix, -1))]+ convert (Const x) iy = [(iy, (0, 2, -1, x))]+ convert (Param ix) iy = [(iy, (0, 1, ix, -1))]+ convert (Uni f t) iy = (iy, (1, fromEnum f, -1, -1)) : t+ convert (Bin op l r) iy = (iy, (2, fromEnum op, -1, -1)) : (l <> r)+{-# INLINE tree2arr #-}
+ src/Algorithm/SRTree/ModelSelection.hs view
@@ -0,0 +1,190 @@+{-# 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++import Debug.Trace++-- | Bayesian information criterion+bic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+bic dist mYerr xss ys theta tree = p * log n + 2 * nll dist mYerr xss ys tree theta+ where+ (A.Sz (fromIntegral -> p)) = A.size theta+ (A.Sz (fromIntegral -> n)) = A.size ys+{-# INLINE bic #-}++-- | Akaike information criterion+aic :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+aic dist mYerr xss ys theta tree = 2 * p + 2 * nll dist mYerr xss ys tree theta+ where+ (A.Sz (fromIntegral -> p)) = A.size theta+ (A.Sz (fromIntegral -> n)) = A.size ys+{-# INLINE aic #-}++-- | Evidence +evidence :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+evidence dist mYerr xss ys theta tree = (1 - b) * nll dist mYerr xss ys tree theta - p / 2 * log b+ where+ (A.Sz (fromIntegral -> p)) = A.size theta+ (A.Sz (fromIntegral -> n)) = A.size ys+ b = 1 / sqrt n+{-# INLINE evidence #-}++fractionalBayesFactor :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+fractionalBayesFactor dist mYerr xss ys theta tree = (1 - b) * nll' - p / 2 * log b + f_compl + p / 2 * log(2*pi*nup)+ where+ nll_val = nll dist mYerr xss ys tree theta + nll_gaus = nll Gaussian mYerr xss ys tree theta+ nll' = if dist == MSE then nll_gaus else nll_val+ (A.Sz (fromIntegral -> p)) = A.size theta+ (A.Sz (fromIntegral -> n)) = A.size ys+ b = 1 / sqrt n+ nup = exp(1 - log 3)+ f_compl = countNodes tree * log (countUniqueTokens tree)+{-# INLINE fractionalBayesFactor #-}++-- | MDL as described in +-- Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023).+mdl :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+mdl dist mYerr xss ys theta tree = nll' dist mYerr xss ys theta tree+ + logFunctional tree+ + logParameters dist mYerr xss ys theta tree+ where+ fisher = fisherNLL dist mYerr xss ys tree theta+ theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher+ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1+{-# INLINE mdl #-}++-- | MDL Lattice as described in+-- Bartlett, Deaglan, Harry Desmond, and Pedro Ferreira. "Priors for symbolic regression." Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 2023.+mdlLatt :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+mdlLatt dist mYerr xss ys theta tree = nll' dist mYerr xss ys theta' tree+ + logFunctional tree+ + logParametersLatt dist mYerr xss ys theta tree+ where+ fisher = fisherNLL dist mYerr xss ys tree theta+ theta' = A.computeAs A.S $ A.zipWith (\t f -> if isSignificant t f then t else 0.0) theta fisher+ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1+{-# INLINE mdlLatt #-}++-- | same as `mdl` but weighting the functional structure by frequency calculated using a wiki information of+-- physics and engineering functions+mdlFreq :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+mdlFreq dist mYerr xss ys theta tree = nll dist mYerr xss ys tree theta+ + logFunctionalFreq tree+ + logParameters dist mYerr xss ys theta tree+{-# 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 PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+logParameters dist mYerr xss ys theta tree = -(p / 2) * log 3 + 0.5 * logFisher + logTheta+ where+ -- p = fromIntegral $ VS.length theta+ fisher = fisherNLL dist mYerr xss ys tree theta++ (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0)+ $ zip (A.toList theta) (A.toList fisher)++ addIfSignificant (v, f) (acc_v, acc_f, acc_p)+ | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1)+ | otherwise = (acc_v, acc_f, acc_p)++ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1++-- same as above but for the Lattice +logParametersLatt :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+logParametersLatt dist mYerr xss ys theta tree = 0.5 * p * (1 - log 3) + 0.5 * log detFisher+ where+ fisher = fisherNLL dist mYerr xss ys tree theta+ detFisher = det $ hessianNLL dist mYerr xss ys tree theta++ (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0)+ $ zip (A.toList theta) (A.toList fisher)++ addIfSignificant (v, f) (acc_v, acc_f, acc_p)+ | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1)+ | otherwise = (acc_v, acc_f, acc_p)++ isSignificant v f = abs (v / sqrt(12 / f) ) >= 1++-- flipped version of nll+nll' :: Distribution -> Maybe PVector -> SRMatrix -> PVector -> PVector -> Fix SRTree -> Double+nll' dist mYerr xss ys theta tree = nll dist mYerr xss ys tree theta+{-# INLINE nll' #-}++treeToNat :: Fix SRTree -> Double+treeToNat = cata $+ \case+ Uni f t -> funToNat f + t+ Bin op l r -> opToNat op + l + r+ _ -> 0.6610799229372109+ 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 #-}
+ src/Algorithm/SRTree/NonlinearOpt.hs view
@@ -0,0 +1,976 @@+{-# 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 nevals) <- N.optimize opt x0+ if (N.isSuccess outret)+ then return $ Solution outcost outx outret nevals+ else Ex.throw $ NloptException outret++minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution+minimizeGlobal' gp x0 = do+ opt <- setupGlobalProblem gp+ solveProblem opt x0++data LocalProblem = LocalProblem+ { lsize :: Word -- ^ The dimension of the+ -- parameter vector.+ , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping+ -- condition+ , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification+ }++-- | These are the local minimization algorithms provided by NLOPT. Please see+-- <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++ , nEvals :: Int -- ^ Number of evaluations until stop+ } deriving (Eq, Show, Read)++applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO ()+applyAugLagAlgorithm opt alg = go alg+ where+ ic = traverse_ (tryTo . applyConstraint opt)+ icd = traverse_ (tryTo . applyConstraint opt)+ -- AUGLAG won't work at all if you don't pass it the same+ -- objective as the subproblem -- here we pull out the subproblem+ -- objectives from the algorithm spec and set the same objective+ -- function so the user can't mess it up.+ local lp = tryTo $ do+ localopt <- setupLocalProblem lp+ applyLocalObjective opt (lalgorithm lp)+ N.set_local_optimizer opt localopt+ global gp = do+ tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt+ applyGlobalObjective opt (galgorithm gp)++ go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd+ go (AUGLAG_EQ_LOCAL lp) = local lp+ go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd+ go (AUGLAG_EQ_GLOBAL gp) = global gp++applyAugLagProblem :: N.Opt -> AugLagProblem -> IO ()+applyAugLagProblem opt (AugLagProblem eq eqd alg) = do+ traverse_ (tryTo . applyConstraint opt) eq+ traverse_ (tryTo . applyConstraint opt) eqd+ applyAugLagAlgorithm opt alg++minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution+minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do+ opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap)+ applyAugLagProblem opt ap+ solveProblem opt x0++-- |+-- == Example program+--+-- The following interactive session example enforces the same scalar+-- constraint as the nonlinear constraint example, but this time it+-- uses the augmented Lagrangian method to enforce the constraint and+-- the 'SBPLX' algorithm, which does not support nonlinear constraints+-- itself, to perform the minimization. As before, the parameters+-- must always sum to 1, and the minimizer finds the same constrained+-- minimum of 22.5 at @(0.5, 0.5)@.+--+-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList )+-- >>> let objf x = x `dot` x + 22+-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []+-- >>> let algorithm = SBPLX objf [] Nothing+-- >>> let subproblem = LocalProblem 2 stop algorithm+-- >>> let x0 = fromList [5, 10]+-- >>> minimizeLocal subproblem x0+-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED})+-- >>> -- define constraint function:+-- >>> let constraintf x = sum (toList x) - 1.0+-- >>> -- define constraint object to pass to the algorithm:+-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6+-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem)+-- >>> minimizeAugLag problem x0+-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED})++minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution+minimizeAugLag prob x0 =+ unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler+ where+ handler :: NloptException -> IO (Either N.Result a)+ handler (NloptException retcode) = return $ Left retcode
+ src/Algorithm/SRTree/Opt.hs view
@@ -0,0 +1,135 @@+{-# LANGUAGE BangPatterns #-}+-----------------------------------------------------------------------------+-- |+-- Module : Algorithm.SRTree.Opt +-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : ConstraintKinds+--+-- Functions to optimize the parameters of an expression.+--+-----------------------------------------------------------------------------+module Algorithm.SRTree.Opt+ where++import Algorithm.SRTree.Likelihoods+import Algorithm.SRTree.NonlinearOpt+import Data.Bifunctor (bimap, second)+import Data.Massiv.Array+import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams, countNodes, convertProtectedOps)+import Data.SRTree.Eval (evalTree, compMode)+import qualified Data.Vector.Storable as VS+import qualified Data.IntMap.Strict as IntMap+import Data.SRTree.Recursion+import Algorithm.EqSat.Egraph hiding ( size )+import Algorithm.EqSat.Build+import Control.Monad.State.Strict+import Control.Monad.Identity+import Algorithm.SRTree.AD (evalCache)++import Debug.Trace++-- | minimizes the negative log-likelihood of the expression+minimizeNLLEGraph :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> EGraph -> EClassId -> ECache -> PVector -> (PVector, Double, Int, ECache)+minimizeNLLEGraph alg dist mYerr niter xss ys egraph root cache t0+ | niter == 0 = (t0, f, 0, cache')+ | n == 0 = (t0, f, 0, cache')+ | otherwise = (t_opt', fst aa, nEvs, cache') -- (t_opt', nll dist mYerr xss ys tree t_opt', nEvs, cache')+ where+ (rt, eg) = buildNLLEGraph dist (fromIntegral m) egraph root -- convertProtectedOps+ t0' = toStorableVector t0+ (Sz n) = size t0+ (Sz m) = size ys+ tree = runIdentity $ getBestExpr root `evalStateT` egraph+ aa = gradNLLEGraph dist xss ys mYerr eg cache' rt t_opt++ funAndGrad = gradNLLEGraph dist xss ys mYerr eg cache' rt+ (f, _) = gradNLLEGraph dist xss ys mYerr eg cache' rt t0' -- if there's no parameter or no iterations+ cache' = evalCache xss egraph cache root t0'+++ algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n)+ stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)]+ problem = LocalProblem (fromIntegral n) stop algorithm+ (t_opt, nEvs) = case minimizeLocal problem t0' of+ Right sol -> (solutionParams sol, nEvals sol)+ Left e -> (t0', 0)+ t_opt' = fromStorableVector compMode t_opt+++-- | minimizes the negative log-likelihood of the expression+minimizeNLL' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeNLL' alg dist mYerr niter xss ys tree t0+ | niter == 0 = (t0, f, 0)+ | n == 0 = (t0, f, 0)+ | otherwise = (t_opt', nll dist mYerr xss ys tree t_opt', nEvs)+ where+ tree' = buildNLL dist (fromIntegral m) tree -- convertProtectedOps+ t0' = toStorableVector t0+ treeArr = IntMap.toAscList $ tree2arr tree'+ j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..]+ (Sz n) = size t0+ (Sz m) = size ys+ funAndGrad = gradNLLGraph dist xss ys mYerr tree' -- second (toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix++ (f, _) = gradNLLGraph dist xss ys mYerr tree' t0' -- if there's no parameter or no iterations+ -- gradNLL dist mYerr xss ys tree t0+ --debug1 = gradNLLArr dist msErr xss ys treeArr j2ix t0+ --debug2 = gradNLL dist msErr xss ys tree t0++ algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) -- alg funAndGrad Nothing -- PRAXIS (fst . funAndGrad) [] Nothing -- TNEWTON funAndGrad Nothing+ stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)]+ problem = LocalProblem (fromIntegral n) stop algorithm+ (t_opt, nEvs) = case minimizeLocal problem t0' of+ Right sol -> (solutionParams sol, nEvals sol) -- traceShow (">>>>>>>", nEvals sol) $+ Left e -> (t0', 0)+ t_opt' = fromStorableVector compMode t_opt+ debugGrad t = let g1 = gradNLL dist mYerr xss ys tree . fromStorableVector compMode $ t+ g2 = gradNLLArr dist xss ys mYerr treeArr j2ix t+ g3 = gradNLLGraph dist xss ys mYerr tree' t+ in traceShow (t, g1, g2, g3) $ g3 -- second (toStorableVector . computeAs S) g2++minimizeNLL :: Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeNLL = minimizeNLL' TNEWTON++-- | minimizes the function while keeping the parameter ix fixed (used to calculate the profile)+minimizeNLLWithFixedParam' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Distribution -> Maybe PVector -> Int -> SRMatrix -> PVector -> Fix SRTree -> Int -> PVector -> PVector+minimizeNLLWithFixedParam' alg dist mYerr niter xss ys tree ix t0+ | niter == 0 = t0+ | n == 0 = t0+ | otherwise = t_opt'+ where+ tree' = buildNLL dist (fromIntegral m) tree -- relabelParams+ t0' = toStorableVector t0+ treeArr = IntMap.toAscList $ tree2arr tree'+ j2ix = IntMap.fromList $ Prelude.zip (Prelude.map fst treeArr) [0..]+ (Sz n) = size t0+ (Sz m) = size ys+ setTo0 = (VS.// [(ix, 0.0)])+ funAndGrad = second (setTo0 . toStorableVector . computeAs S) . gradNLLArr dist xss ys mYerr treeArr j2ix++ (f, _) = gradNLL dist mYerr xss ys tree t0 -- if there's no parameter or no iterations++ algorithm = alg funAndGrad Nothing -- PRAXIS (fst . funAndGrad) [] Nothing -- TNEWTON funAndGrad Nothing+ stop = ObjectiveRelativeTolerance 1e-8 :| [ObjectiveAbsoluteTolerance 1e-8, MaximumEvaluations (fromIntegral niter)]+ problem = LocalProblem (fromIntegral n) stop algorithm+ (t_opt, nEvs) = case minimizeLocal problem t0' of+ Right sol -> (solutionParams sol, nEvals sol) -- traceShow (">>>>>>>", nEvals sol) $+ Left e -> (t0', 0)+ t_opt' = fromStorableVector compMode t_opt++minimizeNLLWithFixedParam = minimizeNLLWithFixedParam' TNEWTON++-- | minimizes using Gaussian likelihood +minimizeGaussian :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeGaussian = minimizeNLL Gaussian Nothing++-- | minimizes using Binomial likelihood +minimizeBinomial :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizeBinomial = minimizeNLL Bernoulli Nothing++-- | minimizes using Poisson likelihood +minimizePoisson :: Int -> SRMatrix -> PVector -> Fix SRTree -> PVector -> (PVector, Double, Int)+minimizePoisson = minimizeNLL Poisson Nothing
src/Data/SRTree.hs view
@@ -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,28 +16,29 @@ , Op(..) , param , var+ , constv , arity , getChildren+ , childrenOf+ , replaceChildren+ , getOperator , countNodes , countVarNodes , countConsts , countParams+ , countParamsUniq , countOccurrences- , deriveBy- , deriveByVar- , deriveByParam- , derivative- , forwardMode- , gradParamsFwd- , gradParamsRev- , evalFun- , evalOp- , inverseFunc- , evalTree+ , countUniqueTokens+ , numberOfVars+ , getIntConsts , relabelParams+ , relabelParamsOrder+ , relabelVars , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) ) where@@ -48,27 +49,28 @@ , Op(..) , param , var+ , constv , arity , getChildren+ , childrenOf+ , replaceChildren+ , getOperator , countNodes , countVarNodes , countConsts , countParams+ , countParamsUniq , countOccurrences- , deriveBy- , deriveByVar- , deriveByParam- , derivative- , forwardMode- , gradParamsFwd- , gradParamsRev- , evalFun- , evalOp- , inverseFunc- , evalTree+ , countUniqueTokens+ , numberOfVars+ , getIntConsts , relabelParams+ , relabelParamsOrder+ , relabelVars , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) )
+ src/Data/SRTree/Datasets.hs view
@@ -0,0 +1,281 @@+{-# 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, loadTrainingOnly, getX, splitData, DataSet(..) )+ 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)+import Data.Massiv.Array as MA hiding (forM_, forM, map, take, tail, zip, replicate, all, read)+import Control.Monad.State.Strict+import System.Random+import List.Shuffle ( shuffle )+++-- a dataset is a triple (X, y, y_error)+type DataSet = (SRMatrix, PVector, Maybe PVector)++-- | 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 6 params)+ where+ (fname : params') = B.split ':' filename+ -- fill up the empty parameters with an empty string+ params = params' <> replicate (6 - min 6 (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 -> B.ByteString -> ([Int], Int, Int)+getColumns headerMap target columns target_error = (ixs, iy, iy_error)+ where+ n_cols = length headerMap+ getIx c = case parseVal c of+ -- 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+ -- if the target PVector is ommitted, use the last one+ iy_error = if B.null target_error+ then (-1)+ else getIx $ B.unpack target_error+{-# inline getColumns #-}++-- | Given the start and end rows, it returns the +-- 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:y_err+-- it returns the X_train, y_train, X_test, y_test, varnames, target name +-- where varnames are a comma separated list of the name of the vars +-- and target name is the name of the target+--+-- 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), (Maybe PVector, Maybe 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), (Maybe PVector, Maybe PVector), String, String)+processData csv params hasHeader = ((x_train, y_train, x_val, y_val) , (y_err_train, y_err_val), varnames, targetname)+ where+ ncols = length $ head csv+ nrows = length csv - fromEnum hasHeader+ (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, iy_err) = getColumns header (params !! 2) (params !! 3) (params !! 4)++ -- load data and split sets+ datum = loadMtx content+ p = length ixs++ x = M.computeAs S $ M.throwEither $ M.stackInnerSlicesM $ map (datum <!) ixs+ y = datum <! iy+ y_err = datum <! iy_err++ x_train = M.computeAs S $ M.extractFromTo' (st :. 0) (end+1 :. p) x+ y_train = M.computeAs S $ M.extractFromTo' st (end+1) y + x_val = M.computeAs S $ M.throwEither $ M.deleteRowsM st (Sz1 $ end - st + 1) x+ y_val = M.computeAs S $ M.throwEither $ M.deleteColumnsM st (Sz1 $ end - st + 1) y++ y_err_train = if iy_err == -1 then Nothing else Just $ M.computeAs S $ M.extractFromTo' st (end+1) y_err+ y_err_val = if iy_err == -1 then Nothing else Just $ M.computeAs S $ M.throwEither $ M.deleteColumnsM st (Sz1 $ end - st + 1) y_err+{-# inline processData #-}++chunksOf :: Int -> [e] -> [[e]]+chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls))+ where+ splitter :: [e] -> ([e] -> a -> a) -> a -> a+ splitter [] _ n = n+ splitter l c n = l `c` splitter (Prelude.drop i l) c n+ build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a]+ build g = g (:) []++splitData :: DataSet ->Int -> State StdGen (DataSet, DataSet)+splitData (x, y, mYErr) k = do+ if k == 1+ then pure ((x, y, mYErr), (x, y, mYErr))+ else do+ ixs' <- (state . shuffle) [0 .. sz-1]+ let ixs = chunksOf k ixs'++ let (x_tr, x_te) = getX ixs x+ (y_tr, y_te) = getY ixs y+ mY = fmap (getY ixs) mYErr+ (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY)+ pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te))+ where+ (MA.Sz sz) = MA.size y+ comp_x = MA.getComp x+ comp_y = MA.getComp y++ getX :: [[Int]] -> SRMatrix -> (SRMatrix, SRMatrix)+ getX ixs xs' = let xs = MA.toLists xs' :: [MA.ListItem MA.Ix2 Double]+ in ( MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromLists' comp_x [xs !! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )+ getY :: [[Int]] -> PVector -> (PVector, PVector)+ getY ixs ys = ( MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i]+ , MA.fromList comp_y [ys MA.! ix | ixs_i <- ixs, let ix = Prelude.head ixs_i]+ )++getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2)+getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c)++getX :: DataSet -> SRMatrix+getX (a, _, _) = a++getTarget :: DataSet -> PVector+getTarget (_, b, _) = b++getError :: DataSet -> Maybe PVector+getError (_, _, c) = c++loadTrainingOnly fname b = getTrain <$> loadDataset fname b
+ src/Data/SRTree/Derivative.hs view
@@ -0,0 +1,127 @@+{-# 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) = (powabs (snd l) (snd r)) * (fst r * log (abs (snd l)) + snd r * fst l / snd l)+ alg1 (Bin AQ l r) = ((1 + snd r * snd r) * fst l - snd l * snd r * fst r) / (1 + snd r * snd r) ** 1.5++ alg2 (Var ix) = var ix+ 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))+ --(abs (snd l) ** (snd r))+ powabs l r = Fix (Bin PowerAbs l r)++-- | Derivative of each supported function+-- For a function h(f) it returns the derivative dh/df+--+-- >>> 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 #-}
+ src/Data/SRTree/Eval.hs view
@@ -0,0 +1,215 @@+{-# LANGUAGE LambdaCase #-}+-----------------------------------------------------------------------------+-- |+-- Module : Data.SRTree.Eval +-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables+--+-- Evaluation of SRTree expressions+--+-----------------------------------------------------------------------------+{-# LANGUAGE FlexibleInstances #-}+module Data.SRTree.Eval+ ( evalTree+ , evalOp+ , evalFun+ , cbrt+ , inverseFunc+ , invertibles+ , evalInverse+ , invright+ , invleft+ , replicateAs+ , SRVector, PVector, SRMatrix+ , compMode+ )+ where++import Data.Massiv.Array+import qualified Data.Massiv.Array as M+import Data.SRTree.Internal+import Data.SRTree.Recursion (Fix (..), cata)++-- | Vector of target values +type SRVector = M.Array D Ix1 Double+-- | Vector of parameter values. Needs to be strict to be readily accesible.+type PVector = M.Array S Ix1 Double+-- | Matrix of features values +type SRMatrix = M.Array S Ix2 Double++compMode :: M.Comp+compMode = M.Seq++-- Improve quality of life with Num and Floating instances for our matrices +instance Index ix => Num (M.Array D ix Double) where+ (+) = (!+!)+ (-) = (!-!)+ (*) = (!*!)+ abs = absA+ signum = signumA + fromInteger = fromInteger+ negate = negateA++instance Index ix => Floating (M.Array D ix Double) where+ pi = pi + exp = expA + log = logA + sqrt = sqrtA + sin = sinA + cos = cosA+ tan = tanA + asin = asinA + acos = acosA + atan = atanA + sinh = sinhA + cosh = coshA+ tanh = tanhA + asinh = asinhA + acosh = acoshA + atanh = atanhA + (**) = (.**)+instance Index ix => Fractional (M.Array D ix Double) where+ fromRational = fromRational+ (/) = (!/!)+ recip = recipA++-- returns a vector with the same number of rows as xss and containing a single repeated value.+replicateAs :: SRMatrix -> Double -> SRVector+replicateAs xss c = let (Sz (m :. _)) = M.size xss in M.replicate (getComp xss) (Sz m) c+{-# INLINE replicateAs #-}++-- | Evaluates the tree given a vector of variable values, a vector of parameter values and a function that takes a Double and change to whatever type the variables have. This is useful when working with datasets of many values per variables.+evalTree :: SRMatrix -> PVector -> Fix SRTree -> SRVector+evalTree xss params = cata $ + \case + Var ix -> xss <! ix+ Param ix -> replicateAs xss $ params ! ix+ Const c -> replicateAs xss c+ Uni g t -> evalFun g t+ Bin op l r -> evalOp op l r+{-# INLINE evalTree #-}++-- evaluates an operator +evalOp :: Floating a => Op -> a -> a -> a+evalOp Add = (+)+evalOp Sub = (-)+evalOp Mul = (*)+evalOp Div = (/)+evalOp Power = (**)+evalOp PowerAbs = \l r -> abs l ** r+evalOp AQ = \l r -> l / sqrt(1 + r*r)+{-# INLINE evalOp #-}++-- evaluates a function +evalFun :: Floating a => Function -> a -> a+evalFun Id = id+evalFun Abs = abs+evalFun Sin = sin+evalFun Cos = cos+evalFun Tan = tan+evalFun Sinh = sinh+evalFun Cosh = cosh+evalFun Tanh = tanh+evalFun ASin = asin+evalFun ACos = acos+evalFun ATan = atan+evalFun ASinh = asinh+evalFun ACosh = acosh+evalFun ATanh = atanh+evalFun Sqrt = sqrt+evalFun SqrtAbs = sqrt . abs+evalFun Cbrt = cbrt+evalFun Square = (^2)+evalFun Log = log+evalFun LogAbs = log . abs+evalFun Exp = exp+evalFun Recip = recip+evalFun Cube = (^3)+{-# INLINE evalFun #-}++-- Cubic root+cbrt :: Floating a => a -> a+cbrt x = signum x * abs x ** (1/3)+{-# INLINE cbrt #-}++-- | Returns the inverse of a function. This is a partial function.+inverseFunc :: Function -> Function+inverseFunc Id = Id+inverseFunc Sin = ASin+inverseFunc Cos = ACos+inverseFunc Tan = ATan+inverseFunc Sinh = ASinh+inverseFunc Cosh = ACosh+inverseFunc Tanh = ATanh+inverseFunc ASin = Sin+inverseFunc ACos = Cos+inverseFunc ATan = Tan+inverseFunc ASinh = Sinh+inverseFunc ACosh = Cosh+inverseFunc ATanh = Tanh+inverseFunc Sqrt = Square+inverseFunc Square = Sqrt+-- inverseFunc Cbrt = (^3)+inverseFunc Log = Exp+inverseFunc Exp = Log+inverseFunc Recip = Recip+-- inverseFunc Abs = Abs -- we assume abs(x) = sqrt(x^2) so y = sqrt(x^2) => x^2 = y^2 => x = sqrt(y^2) = x = abs(y)+inverseFunc x = error $ show x ++ " has no support for inverse function"+{-# INLINE inverseFunc #-}++-- | evals the inverse of a function+evalInverse :: Floating a => Function -> a -> a+evalInverse Id = id+evalInverse Sin = asin+evalInverse Cos = acos+evalInverse Tan = atan+evalInverse Sinh = asinh+evalInverse Cosh = acosh+evalInverse Tanh = atanh+evalInverse ASin = sin+evalInverse ACos = cos+evalInverse ATan = tan+evalInverse ASinh = sinh+evalInverse ACosh = cosh+evalInverse ATanh = tanh+evalInverse Sqrt = (^2)+evalInverse SqrtAbs = (^2)+evalInverse Square = sqrt+evalInverse Cbrt = (^3)+evalInverse Log = exp+evalInverse LogAbs = exp+evalInverse Exp = log+evalInverse Abs = abs -- we assume abs(x) = sqrt(x^2) so y = sqrt(x^2) => x^2 = y^2 => x = sqrt(y^2) = x = abs(y)+evalInverse Recip = recip+evalInverse Cube = cbrt+{-# INLINE evalInverse #-}++-- | evals the right inverse of an operator +invright :: Floating a => Op -> a -> (a -> a)+invright Add v = subtract v+invright Sub v = (+v)+invright Mul v = (/v)+invright Div v = (*v)+invright Power v = (**(1/v))+invright PowerAbs v = (**(1/v))+invright AQ v = (* sqrt (1 + v*v))+{-# INLINE invright #-}++-- | evals the left inverse of an operator +invleft :: Floating a => Op -> a -> (a -> a)+invleft Add v = subtract v+invleft Sub v = (+v) . negate -- y = v - r => r = v - y+invleft Mul v = (/v)+invleft Div v = (v/) -- y = v / r => r = v/y+invleft Power v = logBase v -- (/(log v)) . log -- y = v ^ r log y = r log v r = log y / log v+invleft PowerAbs v = logBase v . abs+invleft AQ v = (v/)+{-# INLINE invleft #-}++-- | List of invertible functions+invertibles :: [Function]+invertibles = [Id, Sin, Cos, Tan, Tanh, ASin, ACos, ATan, ATanh, Sqrt, Square, Log, Exp, Recip]+{-# INLINE invertibles #-}
src/Data/SRTree/Internal.hs view
@@ -1,11 +1,12 @@ {-# language FlexibleInstances, DeriveFunctor #-} {-# language ScopedTypeVariables #-} {-# language RankNTypes #-}-{-# language ViewPatterns #-}+{-# language OverloadedStrings #-}+{-# language LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Internal --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -21,41 +22,40 @@ , Op(..) , param , var+ , constv , arity , getChildren+ , childrenOf+ , replaceChildren+ , getOperator , countNodes , countVarNodes , countConsts , countParams+ , countParamsUniq , countOccurrences- , deriveBy- , deriveByVar- , deriveByParam- , derivative- , forwardMode- , gradParamsFwd- , gradParamsRev- , evalFun- , evalOp- , inverseFunc- , evalTree+ , countUniqueTokens+ , numberOfVars+ , getIntConsts , relabelParams+ , relabelParamsOrder+ , relabelVars , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) ) where -import Data.SRTree.Recursion ( Fix (..), cata, mutu, accu, cataM )--import qualified Data.Vector as V-import Data.Vector ((!))-import Control.Monad.State-import qualified Data.DList as DL-import Data.Bifunctor (second)--import Debug.Trace (trace)+import Control.Monad.State (MonadState (get), State, evalState, modify, put)+import Data.SRTree.Recursion (Fix (..), cata, cataM)+import qualified Data.Set as S+import Data.String (IsString (..))+import Text.Read (readMaybe)+import qualified Data.IntMap as IntMap+import Data.List ( nub ) -- | Tree structure to be used with Symbolic Regression algorithms. -- This structure is a fixed point of a n-ary tree. @@ -63,12 +63,14 @@ Var Int -- ^ index of the variables | Param Int -- ^ index of the parameter | Const Double -- ^ constant value, can be converted to a parameter+ -- | IConst Int -- TODO: integer constant+ -- | RConst Ratio -- TODO: rational constant | Uni Function val -- ^ univariate function | Bin Op val val -- ^ binary operator deriving (Show, Eq, Ord, Functor) -- | Supported operators-data Op = Add | Sub | Mul | Div | Power+data Op = Add | Sub | Mul | Div | Power | PowerAbs | AQ deriving (Show, Read, Eq, Ord, Enum) -- | Supported functions@@ -88,12 +90,42 @@ | ACosh | ATanh | Sqrt+ | SqrtAbs | Cbrt | Square | Log+ | LogAbs | Exp+ | Recip+ | Cube deriving (Show, Read, Eq, Ord, Enum) +removeProtectedOps :: Fix SRTree -> Fix SRTree +removeProtectedOps = cata alg + where + alg (Var ix) = var ix+ alg (Param ix) = param ix+ alg (Const x) = constv x+ alg (Bin PowerAbs l r) = l ** r+ alg (Bin op l r) = Fix $ Bin op l r+ alg (Uni SqrtAbs t) = Fix $ Uni Sqrt t+ alg (Uni LogAbs t) = Fix $ Uni Log t+ alg (Uni f t) = Fix $ Uni f t+{-# INLINE removeProtectedOps #-}++convertProtectedOps :: Fix SRTree -> Fix SRTree +convertProtectedOps = cata alg + where + alg (Var ix) = var ix+ alg (Param ix) = param ix+ alg (Const x) = constv x+ alg (Bin PowerAbs l r) = abs l ** r+ alg (Bin op l r) = Fix $ Bin op l r+ alg (Uni SqrtAbs t) = sqrt (abs t)+ alg (Uni LogAbs t) = log (abs t)+ alg (Uni f t) = Fix $ Uni f t+{-# INLINE convertProtectedOps #-}+ -- | create a tree with a single node representing a variable var :: Int -> Fix SRTree var ix = Fix (Var ix)@@ -102,6 +134,26 @@ param :: Int -> Fix SRTree param ix = Fix (Param ix) +-- | create a tree with a single node representing a constant value+constv :: Double -> Fix SRTree+constv x = Fix (Const x)++-- | the instance of `IsString` allows us to+-- create a tree using a more practical notation:+--+-- >>> :t "x0" + "t0" * sin("x1" * "t1")+-- Fix SRTree+--+instance IsString (Fix SRTree) where + fromString [] = error "empty string for SRTree"+ fromString ('x':ix) = case readMaybe ix of + Just iy -> Fix (Var iy)+ Nothing -> error "wrong format for variable. It should be xi where i is an index. Ex.: \"x0\", \"x1\"."+ fromString ('t':ix) = case readMaybe ix of + Just iy -> Fix (Param iy)+ Nothing -> error "wrong format for parameter. It should be ti where i is an index. Ex.: \"t0\", \"t1\"."+ fromString _ = error "A string can represent a variable or a parameter following the format xi or ti, respectivelly, where i is the index. Ex.: \"x0\", \"t0\"."+ instance Num (Fix SRTree) where Fix (Const 0) + r = r l + Fix (Const 0) = l@@ -142,6 +194,9 @@ l / r = Fix $ Bin Div l r {-# INLINE (/) #-} + recip = Fix . Uni Recip+ {-# INLINE recip #-}+ fromRational = Fix . Const . fromRational {-# INLINE fromRational #-} @@ -188,6 +243,29 @@ logBase l r = log l / log r {-# INLINE logBase #-} +instance Foldable SRTree where + foldMap f =+ \case+ Bin op l r -> f l <> f r+ Uni g t -> f t + _ -> mempty ++instance Traversable SRTree where + traverse f = + \case + Bin op l r -> Bin op <$> f l <*> f r + Uni g t -> Uni g <$> f t + Var ix -> pure (Var ix) + Param ix -> pure (Param ix) + Const x -> pure (Const x) + sequence =+ \case+ Bin op l r -> Bin op <$> l <*> r + Uni g t -> Uni g <$> t + Var ix -> pure (Var ix) + Param ix -> pure (Param ix) + Const x -> pure (Const x) + -- | Arity of the current node arity :: Fix SRTree -> Int arity = cata alg@@ -200,6 +278,10 @@ {-# INLINE arity #-} -- | Get the children of a node. Returns an empty list in case of a leaf node.+--+-- >>> map showExpr . getChildren $ "x0" + 2 +-- ["x0", 2]+-- getChildren :: Fix SRTree -> [Fix SRTree] getChildren (Fix (Var {})) = [] getChildren (Fix (Param {})) = []@@ -208,19 +290,53 @@ getChildren (Fix (Bin _ l r)) = [l, r] {-# INLINE getChildren #-} +-- | Get the children of an unfixed node +-- +childrenOf :: SRTree a -> [a] +childrenOf = + \case + Uni _ t -> [t] + Bin _ l r -> [l, r] + _ -> []++-- | replaces the children with elements from a list +replaceChildren :: [a] -> SRTree b -> SRTree a+replaceChildren [l, r] (Bin op _ _) = Bin op l r+replaceChildren [t] (Uni f _) = Uni f t+replaceChildren _ (Var ix) = Var ix+replaceChildren _ (Param ix) = Param ix+replaceChildren _ (Const x) = Const x+replaceChildren xs n = error "ERROR: trying to replace children with not enough elements."+{-# INLINE replaceChildren #-}++-- | returns a node containing the operator and () as children+getOperator :: SRTree a -> SRTree ()+getOperator (Bin op _ _) = Bin op () ()+getOperator (Uni f _) = Uni f ()+getOperator (Var ix) = Var ix+getOperator (Param ix) = Param ix+getOperator (Const x) = Const x+{-# INLINE getOperator #-}+ -- | Count the number of nodes in a tree.-countNodes :: Fix SRTree -> Int+--+-- >>> countNodes $ "x0" + 2+-- 3+countNodes :: Num a => Fix SRTree -> a countNodes = cata alg where- alg Var {} = 1- alg Param {} = 1- alg Const {} = 1- alg (Uni _ t) = 1 + t+ alg Var {} = 1+ alg Param {} = 1+ alg Const {} = 1+ alg (Uni _ t) = 1 + t alg (Bin _ l r) = 1 + l + r {-# INLINE countNodes #-} -- | Count the number of `Var` nodes-countVarNodes :: Fix SRTree -> Int+--+-- >>> countVarNodes $ "x0" + 2 * ("x0" - sin "x1")+-- 3+countVarNodes :: Num a => Fix SRTree -> a countVarNodes = cata alg where alg Var {} = 1@@ -231,7 +347,10 @@ {-# INLINE countVarNodes #-} -- | Count the number of `Param` nodes-countParams :: Fix SRTree -> Int+--+-- >>> countParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0"+-- 3+countParams :: Num a => Fix SRTree -> a countParams = cata alg where alg Var {} = 0@@ -241,8 +360,25 @@ alg (Bin _ l r) = 0 + l + r {-# INLINE countParams #-} +-- | Count the unique occurrences of `Param` nodes+--+-- >>> countParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0"+-- 2+countParamsUniq :: Fix SRTree -> Int+countParamsUniq t = length . nub $ cata alg t+ where+ alg Var {} = []+ alg (Param ix) = [ix]+ alg Const {} = []+ alg (Uni _ t) = t+ alg (Bin _ l r) = l <> r+{-# INLINE countParamsUniq #-}+ -- | Count the number of const nodes-countConsts :: Fix SRTree -> Int+--+-- >>> countConsts $ "x0"* 2 + 3 * sin "x0"+-- 2+countConsts :: Num a => Fix SRTree -> a countConsts = cata alg where alg Var {} = 0@@ -253,7 +389,10 @@ {-# INLINE countConsts #-} -- | Count the occurrences of variable indexed as `ix`-countOccurrences :: Int -> Fix SRTree -> Int+--+-- >>> countOccurrences 0 $ "x0"* 2 + 3 * sin "x0" + "x1"+-- 2+countOccurrences :: Num a => Int -> Fix SRTree -> a countOccurrences ix = cata alg where alg (Var iy) = if ix == iy then 1 else 0@@ -263,302 +402,170 @@ alg (Bin _ l r) = l + r {-# INLINE countOccurrences #-} --- | Evaluates the tree given a vector of variable values, a vector of parameter values and a function that takes a Double and change to whatever type the variables have. This is useful when working with datasets of many values per variables.-evalTree :: (Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> a-evalTree xss params f = cata alg+-- | counts the number of unique tokens +--+-- >>> countUniqueTokens $ "x0" + ("x1" * "x0" - sin ("x0" ** 2))+-- 8+countUniqueTokens :: Num a => Fix SRTree -> a+countUniqueTokens = len . cata alg where- alg (Var ix) = xss ! ix- alg (Param ix) = f $ params ! ix- alg (Const c) = f c- alg (Uni g t) = evalFun g t- alg (Bin op l r) = evalOp op l r-{-# INLINE evalTree #-}--evalOp :: Floating a => Op -> a -> a -> a-evalOp Add = (+)-evalOp Sub = (-)-evalOp Mul = (*)-evalOp Div = (/)-evalOp Power = (**)-{-# INLINE evalOp #-}--evalFun :: Floating a => Function -> a -> a-evalFun Id = id-evalFun Abs = abs-evalFun Sin = sin-evalFun Cos = cos-evalFun Tan = tan-evalFun Sinh = sinh-evalFun Cosh = cosh-evalFun Tanh = tanh-evalFun ASin = asin-evalFun ACos = acos-evalFun ATan = atan-evalFun ASinh = asinh-evalFun ACosh = acosh-evalFun ATanh = atanh-evalFun Sqrt = sqrt-evalFun Cbrt = cbrt-evalFun Square = (^2)-evalFun Log = log-evalFun Exp = exp-{-# INLINE evalFun #-}---- | Cubic root-cbrt :: Floating val => val -> val-cbrt x = signum x * abs x ** (1/3)-{-# INLINE cbrt #-}---- | Returns the inverse of a function. This is a partial function.-inverseFunc :: Function -> Function-inverseFunc Id = Id-inverseFunc Sin = ASin-inverseFunc Cos = ACos-inverseFunc Tan = ATan-inverseFunc Tanh = ATanh-inverseFunc ASin = Sin-inverseFunc ACos = Cos-inverseFunc ATan = Tan-inverseFunc ATanh = Tanh-inverseFunc Sqrt = Square-inverseFunc Square = Sqrt-inverseFunc Log = Exp-inverseFunc Exp = Log-inverseFunc x = error $ show x ++ " has no support for inverse function"-{-# INLINE inverseFunc #-}+ len (a, b, c, d, e) = fromIntegral $ length a + length b + length c + length d + length e+ alg (Var ix) = (mempty, mempty, S.singleton ix, mempty, mempty)+ alg (Param _) = (mempty, mempty, mempty, S.singleton 1, mempty)+ alg (Const _) = (mempty, mempty, mempty, mempty, S.singleton 1)+ alg (Uni f t) = (mempty, S.singleton f, mempty, mempty, mempty) <> t+ alg (Bin op l r) = (S.singleton op, mempty, mempty, mempty, mempty) <> l <> r+{-# INLINE countUniqueTokens #-} --- | Creates the symbolic partial derivative of a tree by variable `dx` (if `p` is `False`)--- or parameter `dx` (if `p` is `True`).-deriveBy :: Bool -> Int -> Fix SRTree -> Fix SRTree-deriveBy p dx = fst (mutu alg1 alg2)+-- | return the number of unique variables +-- +-- >>> numberOfVars $ "x0" + 2 * ("x0" - sin "x1")+-- 2+numberOfVars :: Num a => Fix SRTree -> a+numberOfVars = fromIntegral . S.size . cata alg where- alg1 (Var ix) = if not p && ix == dx then 1 else 0- alg1 (Param ix) = if p && ix == dx then 1 else 0- alg1 (Const _) = 0- alg1 (Uni f t) = derivative f (snd t) * fst t- alg1 (Bin Add l r) = fst l + fst r- alg1 (Bin Sub l r) = fst l - fst r- alg1 (Bin Mul l r) = fst l * snd r + snd l * fst r- alg1 (Bin Div l r) = (fst l * snd r - snd l * fst r) / snd r ** 2- alg1 (Bin Power l r) = snd l ** (snd r - 1) * (snd r * fst l + snd l * log (snd l) * fst r)-- alg2 (Var ix) = var ix- alg2 (Param ix) = param ix- alg2 (Const c) = Fix (Const c)- alg2 (Uni f t) = Fix (Uni f $ snd t)- alg2 (Bin f l r) = Fix (Bin f (snd l) (snd r))--newtype Tape a = Tape { untape :: [a] } deriving (Show, Functor)--instance Num a => Num (Tape a) where- (Tape x) + (Tape y) = Tape $ zipWith (+) x y- (Tape x) - (Tape y) = Tape $ zipWith (-) x y- (Tape x) * (Tape y) = Tape $ zipWith (*) x y- abs (Tape x) = Tape (map abs x)- signum (Tape x) = Tape (map signum x)- fromInteger x = Tape [fromInteger x]- negate (Tape x) = Tape $ map (*(-1)) x-instance Floating a => Floating (Tape a) where- pi = Tape [pi]- exp (Tape x) = Tape (map exp x)- log (Tape x) = Tape (map log x)- sqrt (Tape x) = Tape (map sqrt x)- sin (Tape x) = Tape (map sin x)- cos (Tape x) = Tape (map cos x)- tan (Tape x) = Tape (map tan x)- asin (Tape x) = Tape (map asin x)- acos (Tape x) = Tape (map acos x)- atan (Tape x) = Tape (map atan x)- sinh (Tape x) = Tape (map sinh x)- cosh (Tape x) = Tape (map cosh x)- tanh (Tape x) = Tape (map tanh x)- asinh (Tape x) = Tape (map asinh x)- acosh (Tape x) = Tape (map acosh x)- atanh (Tape x) = Tape (map atanh x)- (Tape x) ** (Tape y) = Tape $ zipWith (**) x y-instance Fractional a => Fractional (Tape a) where- fromRational x = Tape [fromRational x]- (Tape x) / (Tape y) = Tape $ zipWith (/) x y- recip (Tape x) = Tape $ map recip x+ alg (Uni f t) = t+ alg (Bin op l r) = l <> r+ alg (Var ix) = S.singleton ix+ alg _ = mempty+{-# INLINE numberOfVars #-} --- | Calculates the numerical derivative of a tree using forward mode--- provided a vector of variable values `xss`, a vector of parameter values `theta` and--- a function that changes a Double value to the type of the variable values.-forwardMode :: (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> [a]-forwardMode xss theta f = untape . fst (mutu alg1 alg2)+-- | returns the integer constants. We assume an integer constant +-- as those values in which `floor x == ceiling x`.+--+-- >>> getIntConsts $ "x0" + 2 * "x1" ** 3 - 3.14+-- [2.0,3.0]+getIntConsts :: Fix SRTree -> [Double]+getIntConsts = cata alg where- n = V.length theta- repMat v = Tape $ replicate n v- zeroes = repMat $ f 0- twos = repMat $ f 2- tapeXs = [repMat $ xss ! ix | ix <- [0 .. V.length xss - 1]]- tapeTheta = [repMat $ f (theta ! ix) | ix <- [0 .. n - 1]]- paramVec = [ Tape [if ix==iy then f 1 else f 0 | iy <- [0 .. n-1]] | ix <- [0 .. n-1] ]-- alg1 (Var ix) = zeroes- alg1 (Param ix) = paramVec !! ix- alg1 (Const _) = zeroes- alg1 (Uni f t) = derivative f (snd t) * fst t- alg1 (Bin Add l r) = fst l + fst r- alg1 (Bin Sub l r) = fst l - fst r- alg1 (Bin Mul l r) = (fst l * snd r) + (snd l * fst r)- alg1 (Bin Div l r) = ((fst l * snd r) - (snd l * fst r)) / snd r ** twos- alg1 (Bin Power l r) = snd l ** (snd r - 1) * ((snd r * fst l) + (snd l * log (snd l) * fst r))-- alg2 (Var ix) = tapeXs !! ix- alg2 (Param ix) = tapeTheta !! ix- alg2 (Const c) = repMat $ f c- alg2 (Uni g t) = fmap (evalFun g) (snd t)- alg2 (Bin op l r) = evalOp op (snd l) (snd r)+ alg (Uni f t) = t+ alg (Bin op l r) = l <> r+ alg (Var ix) = []+ alg (Param _) = []+ alg (Const x) = [x | floor x == ceiling x]+{-# INLINE getIntConsts #-} --- | The function `gradParams` calculates the numerical gradient of the tree and evaluates the tree at the same time. It assumes that each parameter has a unique occurrence in the expression. This should be significantly faster than `forwardMode`.-gradParamsFwd :: (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> (a, [a])-gradParamsFwd xss theta f = second DL.toList . cata alg+-- | Relabel the parameters indices incrementaly starting from 0+--+-- >>> showExpr . relabelParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0" +-- "x0" + "t0" * sin ("t1" + "x1") - "t2" +relabelParams :: Fix SRTree -> Fix SRTree+relabelParams t = cataM leftToRight alg t `evalState` 0 where- n = V.length theta-- alg (Var ix) = (xss ! ix, DL.empty)- alg (Param ix) = (f $ theta ! ix, DL.singleton 1)- alg (Const c) = (f c, DL.empty)- alg (Uni f (v, gs)) = let v' = evalFun f v- dv = derivative f v- in (v', DL.map (*dv) gs)- alg (Bin Add (v1, l) (v2, r)) = (v1+v2, DL.append l r)- alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, DL.append l (DL.map negate r))- alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, DL.append (DL.map (*v2) l) (DL.map (*v1) r))- alg (Bin Div (v1, l) (v2, r)) = let dv = (-v1/v2^2) - in (v1/v2, DL.append (DL.map (/v2) l) (DL.map (*dv) r))- alg (Bin Power (v1, l) (v2, r)) = let dv1 = v1 ** (v2 - 1)- dv2 = v1 * log v1- in (v1 ** v2, DL.map (*dv1) (DL.append (DL.map (*v2) l) (DL.map (*dv2) r)))+ -- | leftToRight (left to right) defines the sequence of processing+ leftToRight (Uni f mt) = Uni f <$> mt;+ leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+ leftToRight (Var ix) = pure (Var ix)+ leftToRight (Param ix) = pure (Param ix)+ leftToRight (Const c) = pure (Const c) -data TupleF a b = S a | T a b | B a b b deriving Functor -- hi, I'm a tree-type Tuple a = Fix (TupleF a)+ -- | any time we reach a Param ix, it replaces ix with current state+ -- and increments one to the state.+ alg :: SRTree (Fix SRTree) -> State Int (Fix SRTree)+ alg (Var ix) = pure $ var ix+ alg (Param ix) = do iy <- get; modify (+1); pure (param iy)+ alg (Const c) = pure $ Fix $ Const c+ alg (Uni f t) = pure $ Fix (Uni f t)+ alg (Bin f l r) = pure $ Fix (Bin f l r) -gradParamsRev :: forall a . (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> (a, [a])-gradParamsRev xss theta f t = (getTop fwdMode, DL.toList g)+-- | Reorder the labels of the parameters indices+--+-- >>> showExpr . relabelParamsOrder $ "x0" + "t1" * sin ("t3" + "x1") - "t1"+-- "x0" + "t0" * sin ("t1" + "x1") - "t0"+relabelParamsOrder :: Fix SRTree -> Fix SRTree+relabelParamsOrder t = cataM leftToRight alg t `evalState` (IntMap.empty, 0) where- fwdMode = cata forward t- g = accu reverse combine t (1, fwdMode)-- oneTpl x = Fix $ S x- tuple x y = Fix $ T x y- branch x y z = Fix $ B x y z- getTop (Fix (S x)) = x- getTop (Fix (T x y)) = x- getTop (Fix (B x y z)) = x- unCons (Fix (T x y)) = y- getBranches (Fix (B x y z)) = (y,z)-- forward (Var ix) = oneTpl (xss ! ix)- forward (Param ix) = oneTpl (f $ theta ! ix)- forward (Const c) = oneTpl (f c)- forward (Uni f t) = let v = getTop t- in tuple (evalFun f v) t- forward (Bin op l r) = let vl = getTop l- vr = getTop r- in branch (evalOp op vl vr) l r-- reverse (Var ix) (dx, _) = Var ix- reverse (Param ix) (dx, _) = Param ix- reverse (Const v) (dx, _) = Const v- reverse (Uni f t) (dx, unCons -> v) = Uni f (t, (dx * (derivative f $ getTop v), v))- reverse (Bin op l r) (dx, getBranches -> (vl, vr)) = let (dxl, dxr) = diff op dx (getTop vl) (getTop vr)- in Bin op (l, (dxl, vl)) (r, (dxr, vr))-- diff Add dx vl vr = (dx, dx)- diff Sub dx vl vr = (dx, negate dx)- diff Mul dx vl vr = (dx * vr, dx * vl)- diff Div dx vl vr = (dx / vr, dx * (-vl/vr^2))- diff Power dx vl vr = let dxl = dx * vl ** (vr - 1)- dv2 = vl * log vl- in (dxl * vr, dxl * dv2)-- combine (Var ix) s = DL.empty- combine (Param ix) s = DL.singleton $ fst s- combine (Const c) s = DL.empty- combine (Uni _ gs) s = gs- combine (Bin op l r) s = DL.append l r--derivative :: Floating a => Function -> a -> a-derivative Id = const 1-derivative Abs = \x -> x / abs x-derivative Sin = cos-derivative Cos = negate.sin-derivative Tan = recip . (**2.0) . cos-derivative Sinh = cosh-derivative Cosh = sinh-derivative Tanh = (1-) . (**2.0) . tanh-derivative ASin = recip . sqrt . (1-) . (^2)-derivative ACos = negate . recip . sqrt . (1-) . (^2)-derivative ATan = recip . (1+) . (^2)-derivative ASinh = recip . sqrt . (1+) . (^2)-derivative ACosh = \x -> 1 / (sqrt (x-1) * sqrt (x+1))-derivative ATanh = recip . (1-) . (^2)-derivative Sqrt = recip . (2*) . sqrt-derivative Cbrt = recip . (3*) . cbrt . (^2)-derivative Square = (2*)-derivative Exp = exp-derivative Log = recip-{-# INLINE derivative #-}---- | Symbolic derivative by a variable-deriveByVar :: Int -> Fix SRTree -> Fix SRTree-deriveByVar = deriveBy False+ -- | leftToRight (left to right) defines the sequence of processing+ leftToRight (Uni f mt) = Uni f <$> mt;+ leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+ leftToRight (Var ix) = pure (Var ix)+ leftToRight (Param ix) = pure (Param ix)+ leftToRight (Const c) = pure (Const c) --- | Symbolic derivative by a parameter-deriveByParam :: Int -> Fix SRTree -> Fix SRTree-deriveByParam = deriveBy True+ -- | any time we reach a Param ix, it replaces ix with current state+ -- and increments one to the state.+ alg :: SRTree (Fix SRTree) -> State (IntMap.IntMap Int, Int) (Fix SRTree)+ alg (Var ix) = pure $ var ix+ alg (Param ix) = do (m, iy) <- get+ if IntMap.member ix m+ then pure (param $ m IntMap.! ix)+ else do let m' = IntMap.insert ix iy m+ put (m', iy+1)+ pure (param iy)+ alg (Const c) = pure $ Fix $ Const c+ alg (Uni f t) = pure $ Fix (Uni f t)+ alg (Bin f l r) = pure $ Fix (Bin f l r) --- | Relabel the parameters incrementaly starting from 0-relabelParams :: Fix SRTree -> Fix SRTree-relabelParams t = cataM lTor alg t `evalState` 0+-- | Relabel the parameters indices incrementaly starting from 0+--+-- >>> showExpr . relabelParams $ "x0" + "t0" * sin ("t1" + "x1") - "t0"+-- "x0" + "t0" * sin ("t1" + "x1") - "t2"+relabelVars :: Fix SRTree -> Fix SRTree+relabelVars t = cataM leftToRight alg t `evalState` 0 where- lTor (Uni f mt) = Uni f <$> mt;- lTor (Bin f ml mr) = Bin f <$> ml <*> mr- lTor (Var ix) = pure (Var ix)- lTor (Param ix) = pure (Param ix)- lTor (Const c) = pure (Const c)+ -- | leftToRight (left to right) defines the sequence of processing+ leftToRight (Uni f mt) = Uni f <$> mt;+ leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr+ leftToRight (Var ix) = pure (Var ix)+ leftToRight (Param ix) = pure (Param ix)+ leftToRight (Const c) = pure (Const c) + -- | any time we reach a Param ix, it replaces ix with current state+ -- and increments one to the state. alg :: SRTree (Fix SRTree) -> State Int (Fix SRTree)- alg (Var ix) = pure $ var ix- alg (Param ix) = do iy <- get; modify (+1); pure (param iy)- alg (Const c) = pure $ Fix $ Const c- alg (Uni f t) = pure $ Fix (Uni f t)+ alg (Var ix) = do iy <- get; modify (+1); pure (var iy)+ alg (Param ix) = pure $ param ix+ alg (Const c) = pure $ Fix $ Const c+ alg (Uni f t) = pure $ Fix (Uni f t) alg (Bin f l r) = pure $ Fix (Bin f l r) -- | Change constant values to a parameter, returning the changed tree and a list -- of parameter values+--+-- >>> snd . constsToParam $ "x0" * 2 + 3.14 * sin (5 * "x1")+-- [2.0,3.14,5.0] constsToParam :: Fix SRTree -> (Fix SRTree, [Double]) constsToParam = first relabelParams . cata alg where first f (x, y) = (f x, y) - alg (Var ix) = (Fix $ Var ix, [])- alg (Param ix) = (Fix $ Param ix, [1.0])- alg (Const c) = (Fix $ Param 0, [c])- alg (Uni f t) = (Fix $ Uni f (fst t), snd t)+ -- | If the tree already contains a parameter+ -- it will return a default value of 1.0+ -- whenever it finds a constant, it changes that+ -- to a parameter and adds its content to the singleton list+ alg (Var ix) = (Fix $ Var ix, [])+ alg (Param ix) = (Fix $ Param ix, [1.0])+ alg (Const c) = (Fix $ Param 0, [c])+ alg (Uni f t) = (Fix $ Uni f (fst t), snd t) alg (Bin f l r) = (Fix (Bin f (fst l) (fst r)), snd l <> snd r) -- | Same as `constsToParam` but does not change constant values that -- can be converted to integer without loss of precision+--+-- >>> snd . floatConstsToParam $ "x0" * 2 + 3.14 * sin (5 * "x1")+-- [3.14] floatConstsToParam :: Fix SRTree -> (Fix SRTree, [Double]) floatConstsToParam = first relabelParams . cata alg where- first f (x, y) = (f x, y)+ first f (x, y) = (f x, y)+ combine f (x, y) (z, w) = (f x z, y <> w)+ isInt x = floor x == ceiling x - alg (Var ix) = (Fix $ Var ix, [])- alg (Param ix) = (Fix $ Param ix, [1.0])- alg (Const c) = if floor c == ceiling c then (Fix $ Const c, []) else (Fix $ Param 0, [c])- alg (Uni f t) = (Fix $ Uni f (fst t), snd t)- alg (Bin f l r) = (Fix (Bin f (fst l) (fst r)), snd l <> snd r)+ alg (Var ix) = (var ix, [])+ alg (Param ix) = (param ix, [1.0])+ alg (Const c) = if isInt c then (constv c, []) else (param 0, [c])+ alg (Uni f t) = first (Fix . Uni f) t -- (Fix $ Uni f (fst t), snd t)+ alg (Bin f l r) = combine ((Fix .) . Bin f) l r -- (Fix (Bin f (fst l) (fst r)), snd l <> snd r) -- | Convert the parameters into constants in the tree+--+-- >>> showExpr . paramsToConst [1.1, 2.2, 3.3] $ "x0" + "t0" * sin ("t1" * "x0" - "t2")+-- x0 + 1.1 * sin(2.2 * x0 - 3.3) paramsToConst :: [Double] -> Fix SRTree -> Fix SRTree paramsToConst theta = cata alg where- alg (Var ix) = Fix $ Var ix- alg (Param ix) = Fix $ Const (theta !! ix)- alg (Const c) = Fix $ Const c- alg (Uni f t) = Fix $ Uni f t+ alg (Var ix) = Fix $ Var ix+ alg (Param ix) = Fix $ Const (theta !! ix)+ alg (Const c) = Fix $ Const c+ alg (Uni f t) = Fix $ Uni f t alg (Bin f l r) = Fix $ Bin f l r
src/Data/SRTree/Print.hs view
@@ -1,7 +1,9 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Print --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -12,52 +14,103 @@ ----------------------------------------------------------------------------- module Data.SRTree.Print ( showExpr+ , showExprWithVars , printExpr+ , printExprWithVars , showTikz , printTikz , showPython , printPython , showLatex+ , showLatexWithVars , printLatex+ , showOp ) where -import Control.Monad.Reader ( asks, runReader, Reader )-import Data.Char ( toLower )-+import Control.Monad.Reader (Reader, asks, runReader)+import Data.Char (toLower) import Data.SRTree.Internal-import Data.SRTree.Recursion+import Data.SRTree.Recursion (cata) +-- | converts a tree with protected operators to+-- a conventional math tree+removeProtection :: Fix SRTree -> Fix SRTree+removeProtection = cata $+ \case+ Var ix -> Fix (Var ix)+ Param ix -> Fix (Param ix)+ Const x -> Fix (Const x)+ Uni SqrtAbs t -> sqrt (abs t)+ Uni LogAbs t -> log (abs t)+ Uni Cube t -> t ** 3+ Uni f t -> Fix (Uni f t)+ Bin AQ l r -> l / sqrt (1 + r*r)+ Bin PowerAbs l r -> abs l ** r+ Bin op l r -> Fix (Bin op l r)++-- | convert a tree into a string in math notation +--+-- >>> showExpr $ "x0" + sin ( tanh ("t0" + 2) )+-- "(x0 + Sin(Tanh((t0 + 2.0))))" showExpr :: Fix SRTree -> String-showExpr = cata alg- where- alg (Var ix) = 'x' : show ix- alg (Param ix) = 't' : show ix- alg (Const c) = show c- alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"]- alg (Uni f t) = concat [show f, "(", t, ")"]+showExpr = cata alg . removeProtection+ where alg = \case+ Var ix -> 'x' : show ix+ Param ix -> 't' : show ix+ Const c -> show c+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [show f, "(", t, ")"] +-- | convert a tree into a string in math notation+-- given named vars.+--+-- >>> showExprWithVar ["mu", "eps"] $ "x0" + sin ( "x1" * tanh ("t0" + 2) )+-- "(mu + Sin(Tanh(eps * (t0 + 2.0))))"+showExprWithVars :: [String] -> Fix SRTree -> String+showExprWithVars varnames = cata alg . removeProtection+ where alg = \case+ Var ix -> varnames !! ix+ Param ix -> 't' : show ix+ Const c -> show c+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [show f, "(", t, ")"]++-- | prints the expression printExpr :: Fix SRTree -> IO () printExpr = putStrLn . showExpr +-- | prints the expression+printExprWithVars :: [String] -> Fix SRTree -> IO ()+printExprWithVars varnames = putStrLn . showExprWithVars varnames++-- how to display an operator +showOp :: Op -> String showOp Add = "+" showOp Sub = "-" showOp Mul = "*" showOp Div = "/" showOp Power = "^"+showOp AQ = "aq"+showOp PowerAbs = "|^|" {-# INLINE showOp #-} -- | Displays a tree as a numpy compatible expression.+--+-- >>> showPython $ "x0" + sin ( tanh ("t0" + 2) )+-- "(x[:, 0] + np.sin(np.tanh((t[:, 0] + 2.0))))" showPython :: Fix SRTree -> String-showPython = cata alg+showPython = cata alg . removeProtection where- alg (Var ix) = concat ["x[:, ", show ix, "]"]- alg (Param ix) = concat ["t[:, ", show ix, "]"]- alg (Const c) = show c- alg (Bin Power l r) = concat [l, " ** ", r]- alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"]- alg (Uni f t) = concat [pyFun f, "(", t, ")"]+ alg = \case+ Var ix -> concat ["x[:, ", show ix, "]"]+ Param ix -> concat ["t[", show ix, "]"]+ Const c -> show c+ Bin Power l r -> concat [l, " ** ", r]+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [pyFun f, "(", t, ")"] + pyFun Id = "" pyFun Abs = "np.abs" pyFun Sin = "np.sin"@@ -76,39 +129,69 @@ pyFun Square = "np.square" pyFun Log = "np.log" pyFun Exp = "np.exp"+ pyFun Cbrt = "np.cbrt"+ pyFun Recip = "np.reciprocal" +-- | print the expression in numpy notation printPython :: Fix SRTree -> IO () printPython = putStrLn . showPython --- | Displays a tree as a sympy compatible expression.+-- | Displays a tree as a LaTeX compatible expression.+--+-- >>> showLatex $ "x0" + sin ( tanh ("t0" + 2) )+-- "\\left(x_{, 0} + \\operatorname{sin}(\\operatorname{tanh}(\\left(\\theta_{, 0} + 2.0\\right)))\\right)" showLatex :: Fix SRTree -> String-showLatex = cata alg+showLatex = cata alg . removeProtection where- alg (Var ix) = concat ["x_{, ", show ix, "}"]- alg (Param ix) = concat ["\\theta_{, ", show ix, "}"]- alg (Const c) = show c- alg (Bin Power l r) = concat [l, "^{", r, "}"]- alg (Bin op l r) = concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]- alg (Uni Abs t) = concat ["\\left |", t, "\\right |"]- alg (Uni f t) = concat [showLatexFun f, "(", t, ")"]-+ alg = \case+ Var ix -> concat ["x_{", show ix, "}"]+ Param ix -> concat ["\\theta_{", show ix, "}"]+ Const c -> show c+ Bin Power l r -> concat ["{", l, "^{", r, "}}"]+ Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"]+ Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"]+ Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"]+ Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]+ Uni Abs t -> concat ["\\left |", t, "\\right |"]+ Uni Recip t -> concat ["\\frac{1}{", t, "}"]+ Uni f t -> concat [showLatexFun f, "(", t, ")"]+ +showLatexWithVars :: [String] -> Fix SRTree -> String+showLatexWithVars varnames = cata alg . removeProtection+ where + alg = \case+ Var ix -> concat ["\\operatorname{", varnames !! ix, "}"]+ Param ix -> concat ["\\theta_{", show ix, "}"]+ Const c -> show c+ Bin Power l r -> concat ["{", l, "^{", r, "}}"]+ Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"]+ Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"]+ Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"]+ Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]+ Uni Abs t -> concat ["\\left |", t, "\\right |"]+ Uni Recip t -> concat ["\\frac{1}{", t, "}"]+ Uni f t -> concat [showLatexFun f, "(", t, ")"]+ showLatexFun :: Function -> String showLatexFun f = mconcat ["\\operatorname{", map toLower $ show f, "}"] {-# INLINE showLatexFun #-} +-- | prints expression in LaTeX notation. printLatex :: Fix SRTree -> IO () printLatex = putStrLn . showLatex -- | Displays a tree in Tikz format showTikz :: Fix SRTree -> String-showTikz = cata alg+showTikz = cata alg . removeProtection where+ alg = \case+ Var ix -> concat ["[$x_{, ", show ix, "}$]\n"]+ Param ix -> concat ["[$\\theta_{, ", show ix, "}$]\n"]+ Const c -> concat ["[$", show (roundN 2 c), "$]\n"]+ Bin op l r -> concat ["[", showOpTikz op, l, r, "]\n"]+ Uni f t -> concat ["[", map toLower $ show f, t, "]\n"]+ roundN n x = let ten = 10^n in (/ ten) . fromIntegral . round $ x*ten- alg (Var ix) = concat ["[$x_{, ", show ix, "}$]\n"]- alg (Param ix) = concat ["[$\\theta_{, ", show ix, "}$]\n"]- alg (Const c) = concat ["[$", show (roundN 2 c), "$]\n"]- alg (Bin op l r) = concat ["[", showOpTikz op, l, r, "]\n"]- alg (Uni f t) = concat ["[", map toLower $ show f, t, "]\n"] showOpTikz Add = "+\n" showOpTikz Sub = "-\n"@@ -116,4 +199,6 @@ showOpTikz Div = "÷\n" showOpTikz Power = "\\^{}\n" +-- | prints the tree in TikZ format +printTikz :: Fix SRTree -> IO () printTikz = putStrLn . showTikz
src/Data/SRTree/Random.hs view
@@ -2,7 +2,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Random --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -18,25 +18,35 @@ , HasEverything , FullParams(..) , RndTree+ , Rng(..) , randomVar , randomConst , randomPow , randomFunction , randomNode , randomNonTerminal+ , randomRange+ , randomTreeTemplate , randomTree , randomTreeBalanced+ , toss+ , tossBiased+ , randomVal+ , randomVec+ , randomFrom ) where -import System.Random -import Control.Monad.State -import Control.Monad.Reader +import Control.Monad.Reader (ReaderT, asks, runReaderT)+import Control.Monad.State.Strict ( MonadState(state), MonadTrans(lift), StateT ) import Data.Maybe (fromJust)- import Data.SRTree.Internal-import Data.SRTree.Recursion+import System.Random (Random (random, randomR), StdGen, mkStdGen)+import Data.Massiv.Array as MA hiding (forM_, forM, P)+import Data.SRTree.Eval+import Control.Monad + -- * Class definition of properties that a certain parameter type has. -- -- HasVars: does `p` provides a list of the variable indices?@@ -67,19 +77,28 @@ instance HasFuns FullParams where _funs (P _ _ _ fs) = fs +type Rng m a = StateT StdGen m a+ -- auxiliary function to sample between False and True-toss :: StateT StdGen IO Bool+toss :: Monad m => Rng m Bool toss = state random {-# INLINE toss #-} +tossBiased :: Monad m => Double -> Rng m Bool+tossBiased p = do r <- state random+ pure (r < p)++randomVal :: Monad m => Rng m Double+randomVal = state random+ -- returns a random element of a list-randomFrom :: [a] -> StateT StdGen IO a+randomFrom :: Monad m => [a] -> Rng m a randomFrom funs = do n <- randomRange (0, length funs - 1) pure $ funs !! n {-# INLINE randomFrom #-} -- returns a random element within a range-randomRange :: (Ord val, Random val) => (val, val) -> StateT StdGen IO val+randomRange :: (Ord val, Random val, Monad m) => (val, val) -> Rng m val randomRange rng = state (randomR rng) {-# INLINE randomRange #-} @@ -90,38 +109,38 @@ {-# INLINE replaceChild #-} -- Replace the children of a binary tree.-replaceChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)-replaceChildren (Fix (Bin f _ _)) l r = Just $ Fix (Bin f l r)-replaceChildren _ _ _ = Nothing-{-# INLINE replaceChildren #-}+replaceFixChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)+replaceFixChildren (Fix (Bin f _ _)) l r = Just $ Fix (Bin f l r)+replaceFixChildren _ _ _ = Nothing+{-# INLINE replaceFixChildren #-} -- | RndTree is a Monad Transformer to generate random trees of type `SRTree ix val` -- given the parameters `p ix val` using the random number generator `StdGen`.-type RndTree p = ReaderT p (StateT StdGen IO) (Fix SRTree)+type RndTree m p = ReaderT p (StateT StdGen m) (Fix SRTree) -- | Returns a random variable, the parameter `p` must have the `HasVars` property-randomVar :: HasVars p => RndTree p+randomVar :: Monad m => HasVars p => RndTree m p randomVar = do vars <- asks _vars lift $ Fix . Var <$> randomFrom vars -- | Returns a random constant, the parameter `p` must have the `HasConst` property-randomConst :: HasVals p => RndTree p+randomConst :: (HasVals p, Monad m) => RndTree m p randomConst = do rng <- asks _range lift $ Fix . Const <$> randomRange rng -- | Returns a random integer power node, the parameter `p` must have the `HasExps` property-randomPow :: HasExps p => RndTree p+randomPow :: (HasExps p, Monad m) => RndTree m p randomPow = do rng <- asks _exponents lift $ Fix . Bin Power 0 . Fix . Const . fromIntegral <$> randomRange rng -- | Returns a random function, the parameter `p` must have the `HasFuns` property-randomFunction :: HasFuns p => RndTree p+randomFunction :: (HasFuns p, Monad m) => RndTree m p randomFunction = do funs <- asks _funs f <- lift $ randomFrom funs lift $ pure $ Fix (Uni f 0) -- | Returns a random node, the parameter `p` must have every property.-randomNode :: HasEverything p => RndTree p+randomNode :: (HasEverything p, Monad m) => RndTree m p randomNode = do choice <- lift $ randomRange (0, 8 :: Int) case choice of@@ -136,7 +155,7 @@ 8 -> pure . Fix $ Bin Power 0 0 -- | Returns a random non-terminal node, the parameter `p` must have every property.-randomNonTerminal :: HasEverything p => RndTree p+randomNonTerminal :: (HasEverything p, Monad m) => RndTree m p randomNonTerminal = do choice <- lift $ randomRange (0, 6 :: Int) case choice of@@ -149,21 +168,31 @@ 6 -> pure . Fix $ Bin Power 0 0 -- | Returns a random tree with a limited budget, the parameter `p` must have every property.-randomTree :: HasEverything p => Int -> RndTree p-randomTree 0 = do+--+-- >>> let treeGen = runReaderT (randomTree 12) (P [0,1] (-10, 10) (2, 3) [Log, Exp])+-- >>> tree <- evalStateT treeGen (mkStdGen 52)+-- >>> showExpr tree+-- "(-2.7631152121655838 / Exp((x0 / ((x0 * -7.681722660704317) - Log(3.378309080134594)))))"+randomTreeTemplate :: (HasEverything p, Monad m) => Int -> RndTree m p+randomTreeTemplate 0 = do coin <- lift toss if coin then randomVar else randomConst-randomTree budget = do +randomTreeTemplate budget = do node <- randomNode fromJust <$> case arity node of 0 -> pure $ Just node- 1 -> replaceChild node <$> randomTree (budget - 1)- 2 -> replaceChildren node <$> randomTree (budget `div` 2) <*> randomTree (budget `div` 2)+ 1 -> replaceChild node <$> randomTreeTemplate (budget - 1)+ 2 -> replaceFixChildren node <$> randomTreeTemplate (budget `div` 2) <*> randomTreeTemplate (budget `div` 2) -- | Returns a random tree with a approximately a number `n` of nodes, the parameter `p` must have every property.-randomTreeBalanced :: HasEverything p => Int -> RndTree p+--+-- >>> let treeGen = runReaderT (randomTreeBalanced 10) (P [0,1] (-10, 10) (2, 3) [Log, Exp])+-- >>> tree <- evalStateT treeGen (mkStdGen 42)+-- >>> showExpr tree+-- "Exp(Log((((7.784360517385774 * x0) - (3.6412224491658223 ^ x1)) ^ ((x0 ^ -4.09764995657091) + Log(-7.710216839988497)))))"+randomTreeBalanced :: (HasEverything p, Monad m) => Int -> RndTree m p randomTreeBalanced n | n <= 1 = do coin <- lift toss if coin@@ -173,4 +202,29 @@ node <- randomNonTerminal fromJust <$> case arity node of 1 -> replaceChild node <$> randomTreeBalanced (n - 1)- 2 -> replaceChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) + 2 -> replaceFixChildren node <$> randomTreeBalanced (n `div` 2) <*> randomTreeBalanced (n `div` 2) +++randomVec :: Monad m => Int -> Rng m PVector+randomVec n = MA.fromList compMode <$> replicateM n (randomRange (-1, 1))++randomTree :: Monad m => Int -> Int -> Int -> Rng m (Fix SRTree) -> Rng m (SRTree ()) -> Bool -> Rng m (Fix SRTree)+randomTree minDepth maxDepth maxSize genTerm genNonTerm grow+ | noSpaceLeft = genTerm+ | needNonTerm = genRecursion+ | otherwise = do r <- toss+ if r+ then genTerm+ else genRecursion+ where+ noSpaceLeft = maxDepth <= 1 || maxSize <= 2+ needNonTerm = (minDepth >= 0 || (maxDepth > 2 && not grow)) -- && maxSize > 2++ genRecursion = do+ node <- genNonTerm+ case node of+ Uni f _ -> Fix . Uni f <$> randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1) genTerm genNonTerm grow+ Bin op _ _ -> do l <- randomTree (minDepth - 1) (maxDepth - 1) (if grow then maxSize - 2 else maxSize `div` 2) genTerm genNonTerm grow+ r <- randomTree (minDepth - 1) (maxDepth - 1) (maxSize - 1 - countNodes l) genTerm genNonTerm grow+ pure . Fix $ Bin op l r+{-# INLINE randomTree #-}
src/Data/SRTree/Recursion.hs view
@@ -1,5 +1,17 @@ {-# language RankNTypes #-} {-# language DeriveFunctor #-}+-----------------------------------------------------------------------------+-- |+-- Module : Data.SRTree.Recursion +-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables+--+-- Recursion schemes+--+----------------------------------------------------------------------------- module Data.SRTree.Recursion where import Control.Monad ( (>=>) )
+ src/Numeric/Optimization/NLOPT/Bindings.hs view
@@ -0,0 +1,1075 @@+{-# OPTIONS_GHC -Wall #-}+{-# LANGUAGE ForeignFunctionInterface #-}+{-# LANGUAGE NoMonomorphismRestriction #-}++{- |+Module : Numeric.Optimization.NLOPT.Bindings+Copyright : (c) Matthew Peddie 2017+License : BSD3+Maintainer : Matthew Peddie <mpeddie@gmail.com>+Stability : provisional+Portability : GHC++Low-level interface to the NLOPT library. Please see+<http://ab-initio.mit.edu/wiki/index.php/NLopt_Reference the NLOPT reference manual>+for detailed information; the Haskell functions in this module closely+follow the interface to the C library in @nlopt.h@.++Differences between this module and the C interface are documented+here; functions with identical interfaces are not. In general:++ ['Opt'] corresponds to an @nlopt_opt@ object++ ['Result'] corresponds to @nlopt_result@++ ['V.Vector' 'Double'] corresponds to a @const double *@ input or a+ @double *@ output++ ['ScalarFunction'] corresponds to @nlopt_func@++ ['VectorFunction'] corresponds to @nlopt_mfunc@++ ['PreconditionerFunction'] corresponds to @nlopt_precond@++User data that is handled by @void *@ in the C bindings can be any+Haskell value.++-}++module Numeric.Optimization.NLOPT.Bindings (+ -- * C enums+ Algorithm(..)+ , algorithm_name+ , Result(..)+ , isSuccess+ -- * Optimizer object+ , Opt+ , create+ , destroy+ , copy+ -- * Random number generator seeding+ , srand+ , srand_time+ -- * Metadata+ , Version(..)+ , version+ , get_algorithm+ , get_dimension+ -- * Callbacks+ , ScalarFunction+ , VectorFunction+ , PreconditionerFunction+ -- * Running the optimizer+ , Output(..)+ , optimize+ -- * Objective function configuration+ , set_min_objective+ , set_max_objective+ , set_precond_min_objective+ , set_precond_max_objective+ -- * Bound configuration+ , set_lower_bounds+ , set_lower_bounds1+ , get_lower_bounds+ , set_upper_bounds+ , set_upper_bounds1+ , get_upper_bounds+ -- * Constraint configuration+ , remove_inequality_constraints+ , add_inequality_constraint+ , add_precond_inequality_constraint+ , add_inequality_mconstraint+ , remove_equality_constraints+ , add_equality_constraint+ , add_precond_equality_constraint+ , add_equality_mconstraint+ -- * Stopping criterion configuration+ , set_stopval+ , get_stopval+ , set_ftol_rel+ , get_ftol_rel+ , set_ftol_abs+ , get_ftol_abs+ , set_xtol_rel+ , get_xtol_rel+ , set_xtol_abs1+ , set_xtol_abs+ , get_xtol_abs+ , set_maxeval+ , get_maxeval+ , set_maxtime+ , get_maxtime+ , force_stop+ , set_force_stop+ , get_force_stop+ -- * Algorithm-specific configuration+ , set_local_optimizer+ , set_population+ , get_population+ , set_vector_storage+ , get_vector_storage+ , set_default_initial_step+ , set_initial_step+ , set_initial_step1+ , get_initial_step+ ) where++import Foreign hiding (void)+import Foreign.C.String+import Foreign.C.Types+import qualified Foreign.Concurrent as CFP++import qualified Data.Vector.Storable.Mutable as MV+import qualified Data.Vector.Storable as V++{- C enums -}++-- | The NLOPT algorithm names, apart from the names of the actual+-- optimization methods, follow this scheme:+--+-- [@G@] means a global method+-- [@L@] means a local method+-- [@D@] means a method that requires the derivative+-- [@N@] means a method that does not require the derivative+-- [@*_RAND@] means the algorithm involves some randomization.+-- [@*_NOSCAL@] means the algorithm is *not* scaled to a unit+-- hypercube (i.e. it is sensitive to the units of x)+data Algorithm+ = GN_DIRECT -- ^ DIviding RECTangles+ | GN_DIRECT_L -- ^ DIviding RECTangles,+ -- locally-biased variant+ | GN_DIRECT_L_RAND -- ^ DIviding RECTangles, "slightly+ -- randomized"+ | GN_DIRECT_NOSCAL -- ^ DIviding RECTangles, unscaled version+ | GN_DIRECT_L_NOSCAL -- ^ DIviding RECTangles,+ -- locally-biased and unscaled+ | GN_DIRECT_L_RAND_NOSCAL -- ^ DIviding RECTangles, locally-biased,+ -- unscaled and "slightly randomized"+ | GN_ORIG_DIRECT -- ^ DIviding RECTangles, original FORTRAN+ -- implementation+ | GN_ORIG_DIRECT_L -- ^ DIviding RECTangles,+ -- locally-biased, original FORTRAN+ -- implementation+ | GD_STOGO -- ^ Stochastic Global Optimization+ | GD_STOGO_RAND -- ^ Stochastic Global Optimization,+ -- randomized variant+ | LD_LBFGS_NOCEDAL -- ^ Limited-memory BFGS+ | LD_LBFGS -- ^ Limited-memory BFGS+ | LN_PRAXIS -- ^ PRincipal AXIS gradient-free local+ -- optimization+ | LD_VAR2 -- ^ Shifted limited-memory+ -- variable-metric, rank-2+ | LD_VAR1 -- ^ Shifted limited-memory+ -- variable-metric, rank-1+ | LD_TNEWTON -- ^ Truncated Newton's method+ | LD_TNEWTON_RESTART -- ^ Truncated Newton's method with+ -- automatic restarting+ | LD_TNEWTON_PRECOND -- ^ Preconditioned truncated Newton's+ -- method+ | LD_TNEWTON_PRECOND_RESTART -- ^ Preconditioned truncated Newton's+ -- method with automatic restarting+ | GN_CRS2_LM -- ^ Controlled Random Search with+ -- Local Mutation+ | GN_MLSL -- ^ Original Multi-Level+ -- Single-Linkage+ | GD_MLSL -- ^ Original Multi-Level+ -- Single-Linkage, user-provided+ -- derivative+ | GN_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points+ | GD_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points, user-provided+ -- derivative+ | LD_MMA -- ^ Method of moving averages+ | LN_COBYLA -- ^ Constrained Optimization BY Linear+ -- Approximations+ | LN_NEWUOA -- ^ Powell's NEWUOA algorithm+ | LN_NEWUOA_BOUND -- ^ Powell's NEWUOA algorithm with+ -- bounds by SGJ+ | LN_NELDERMEAD -- ^ Nelder-Mead Simplex gradient-free+ -- method+ | LN_SBPLX -- ^ NLOPT implementation of Rowan's+ -- Subplex algorithm+ | LN_AUGLAG -- ^ AUGmented LAGrangian+ | LD_AUGLAG -- ^ AUGmented LAGrangian,+ -- user-provided derivative+ | LN_AUGLAG_EQ -- ^ AUGmented LAGrangian with penalty+ -- functions only for equality+ -- constraints+ | LD_AUGLAG_EQ -- ^ AUGmented LAGrangian with+ -- penalty functions only for equality+ -- constraints, user-provided+ -- derivative+ | LN_BOBYQA -- ^ Bounded Optimization BY Quadratic+ -- Approximations+ | GN_ISRES -- ^ Improved Stochastic Ranking+ -- Evolution Strategy++ | AUGLAG -- ^ AUGmented LAGrangian, requires+ -- local_optimizer to be set+ | AUGLAG_EQ -- ^ AUGmented LAGrangian with penalty+ -- functions only for equality+ -- constraints, requires+ -- local_optimizer to be set+ | G_MLSL -- ^ Original Multi-Level+ -- Single-Linkage, user-provided+ -- derivative, requires local_optimizer+ -- to be set+ | G_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points, requires+ -- local_optimizer to be set+ | LD_SLSQP -- ^ Sequential Least-SQuares Programming+ | LD_CCSAQ -- ^ Conservative Convex Separable+ -- Approximation+ | GN_ESCH -- ^ Evolutionary Algorithm+ deriving (Eq, Show, Read, Bounded)++instance Enum Algorithm where+ fromEnum GN_DIRECT = 0+ fromEnum GN_DIRECT_L = 1+ fromEnum GN_DIRECT_L_RAND = 2+ fromEnum GN_DIRECT_NOSCAL = 3+ fromEnum GN_DIRECT_L_NOSCAL = 4+ fromEnum GN_DIRECT_L_RAND_NOSCAL = 5+ fromEnum GN_ORIG_DIRECT = 6+ fromEnum GN_ORIG_DIRECT_L = 7+ fromEnum GD_STOGO = 8+ fromEnum GD_STOGO_RAND = 9+ fromEnum LD_LBFGS_NOCEDAL = 10+ fromEnum LD_LBFGS = 11+ fromEnum LN_PRAXIS = 12+ fromEnum LD_VAR2 = 13+ fromEnum LD_VAR1 = 14+ fromEnum LD_TNEWTON = 15+ fromEnum LD_TNEWTON_RESTART = 16+ fromEnum LD_TNEWTON_PRECOND = 17+ fromEnum LD_TNEWTON_PRECOND_RESTART = 18+ fromEnum GN_CRS2_LM = 19+ fromEnum GN_MLSL = 20+ fromEnum GD_MLSL = 21+ fromEnum GN_MLSL_LDS = 22+ fromEnum GD_MLSL_LDS = 23+ fromEnum LD_MMA = 24+ fromEnum LN_COBYLA = 25+ fromEnum LN_NEWUOA = 26+ fromEnum LN_NEWUOA_BOUND = 27+ fromEnum LN_NELDERMEAD = 28+ fromEnum LN_SBPLX = 29+ fromEnum LN_AUGLAG = 30+ fromEnum LD_AUGLAG = 31+ fromEnum LN_AUGLAG_EQ = 32+ fromEnum LD_AUGLAG_EQ = 33+ fromEnum LN_BOBYQA = 34+ fromEnum GN_ISRES = 35+ fromEnum AUGLAG = 36+ fromEnum AUGLAG_EQ = 37+ fromEnum G_MLSL = 38+ fromEnum G_MLSL_LDS = 39+ fromEnum LD_SLSQP = 40+ fromEnum LD_CCSAQ = 41+ fromEnum GN_ESCH = 42+ toEnum 0 = GN_DIRECT+ toEnum 1 = GN_DIRECT_L+ toEnum 2 = GN_DIRECT_L_RAND+ toEnum 3 = GN_DIRECT_NOSCAL+ toEnum 4 = GN_DIRECT_L_NOSCAL+ toEnum 5 = GN_DIRECT_L_RAND_NOSCAL+ toEnum 6 = GN_ORIG_DIRECT+ toEnum 7 = GN_ORIG_DIRECT_L+ toEnum 8 = GD_STOGO+ toEnum 9 = GD_STOGO_RAND+ toEnum 10 = LD_LBFGS_NOCEDAL+ toEnum 11 = LD_LBFGS+ toEnum 12 = LN_PRAXIS+ toEnum 13 = LD_VAR2+ toEnum 14 = LD_VAR1+ toEnum 15 = LD_TNEWTON+ toEnum 16 = LD_TNEWTON_RESTART+ toEnum 17 = LD_TNEWTON_PRECOND+ toEnum 18 = LD_TNEWTON_PRECOND_RESTART+ toEnum 19 = GN_CRS2_LM+ toEnum 20 = GN_MLSL+ toEnum 21 = GD_MLSL+ toEnum 22 = GN_MLSL_LDS+ toEnum 23 = GD_MLSL_LDS+ toEnum 24 = LD_MMA+ toEnum 25 = LN_COBYLA+ toEnum 26 = LN_NEWUOA+ toEnum 27 = LN_NEWUOA_BOUND+ toEnum 28 = LN_NELDERMEAD+ toEnum 29 = LN_SBPLX+ toEnum 30 = LN_AUGLAG+ toEnum 31 = LD_AUGLAG+ toEnum 32 = LN_AUGLAG_EQ+ toEnum 33 = LD_AUGLAG_EQ+ toEnum 34 = LN_BOBYQA+ toEnum 35 = GN_ISRES+ toEnum 36 = AUGLAG+ toEnum 37 = AUGLAG_EQ+ toEnum 38 = G_MLSL+ toEnum 39 = G_MLSL_LDS+ toEnum 40 = LD_SLSQP+ toEnum 41 = LD_CCSAQ+ toEnum 42 = GN_ESCH+ toEnum e = error $+ "Algorithm.toEnum: invalid C value '" ++ show e ++ "' received."++foreign import ccall "nlopt.h nlopt_algorithm_name"+ nlopt_algorithm_name :: CInt -> CString++algorithm_name :: Algorithm -> IO String+algorithm_name = peekCString . nlopt_algorithm_name . fromIntegral . fromEnum++-- | Mostly self-explanatory.+data Result+ = FAILURE -- ^ Generic failure code+ | INVALID_ARGS+ | OUT_OF_MEMORY+ | ROUNDOFF_LIMITED+ | FORCED_STOP+ | SUCCESS -- ^ Generic success code+ | STOPVAL_REACHED+ | FTOL_REACHED+ | XTOL_REACHED+ | MAXEVAL_REACHED+ | MAXTIME_REACHED+ deriving (Eq, Read, Show, Bounded)++instance Enum Result where+ fromEnum FAILURE = -1+ fromEnum INVALID_ARGS = -2+ fromEnum OUT_OF_MEMORY = -3+ fromEnum ROUNDOFF_LIMITED = -4+ fromEnum FORCED_STOP = -5+ fromEnum SUCCESS = 1+ fromEnum STOPVAL_REACHED = 2+ fromEnum FTOL_REACHED = 3+ fromEnum XTOL_REACHED = 4+ fromEnum MAXEVAL_REACHED = 5+ fromEnum MAXTIME_REACHED = 6+ toEnum (-1) = FAILURE+ toEnum (-2) = INVALID_ARGS+ toEnum (-3) = OUT_OF_MEMORY+ toEnum (-4) = ROUNDOFF_LIMITED+ toEnum (-5) = FORCED_STOP+ toEnum 1 = SUCCESS+ toEnum 2 = STOPVAL_REACHED+ toEnum 3 = FTOL_REACHED+ toEnum 4 = XTOL_REACHED+ toEnum 5 = MAXEVAL_REACHED+ toEnum 6 = MAXTIME_REACHED+ toEnum e = error $+ "Result.toEnum: invalid C value '" ++ show e ++ "' received."++isSuccess :: Result -> Bool+isSuccess SUCCESS = True+isSuccess STOPVAL_REACHED = True+isSuccess FTOL_REACHED = True+isSuccess XTOL_REACHED = True+isSuccess MAXEVAL_REACHED = True+isSuccess MAXTIME_REACHED = True+isSuccess _ = False++parseEnum :: (Integral a, Enum b) => a -> b+parseEnum = toEnum . fromIntegral++{- NLOPT optimizer object -}++type NloptOpt = Ptr ()++-- | An optimizer object which must be created, configured and then+-- passed to 'optimize' to solve a problem+newtype Opt = Opt { pointerFromOpt :: ForeignPtr () }++withOpt :: Opt -> (NloptOpt -> IO a) -> IO a+withOpt (Opt p) f = do+ ret <- withForeignPtr p f+ touchForeignPtr p -- This is critical! Otherwise the GC might+ -- think it's done with everything in the middle+ -- of the problem.+ return ret++useOpt :: (NloptOpt -> IO a) -> Opt -> IO a+useOpt = flip withOpt++-- Every time we make a "wrapper" call, the runtime allocates a new+-- function pointer and won't release it until we explicitly tell it+-- to. This doesn't mesh well with NLOPT's "object-oriented" design,+-- wherein we have to allocate an object and make a bunch of setup+-- calls before we run the problem, so what we do is add a finalizer+-- to the 'Opt' object's 'ForeignPtr' every time we need to create a+-- function pointer for C to use.+addFunPtrFinalizer :: Opt -> FunPtr a -> IO ()+addFunPtrFinalizer (Opt p) funptr =+ CFP.addForeignPtrFinalizer p (freeHaskellFunPtr funptr)++foreign import ccall "nlopt.h nlopt_create"+ nlopt_create :: CInt -> CUInt -> IO (NloptOpt)++-- | Create a new 'Opt' object+create :: Algorithm -- ^ Choice of algorithm+ -> Word -- ^ Parameter vector dimension+ -> IO (Maybe Opt) -- ^ Optimizer object+create alg dimension = do+ outp <- nlopt_create (fromIntegral $ fromEnum alg) (fromIntegral dimension)+ if (outp == nullPtr)+ then return Nothing+ else Just . Opt <$> CFP.newForeignPtr outp (nlopt_destroy outp)++foreign import ccall "nlopt.h nlopt_destroy"+ nlopt_destroy :: NloptOpt -> IO ()++-- It shouldn't be strictly necessary to call this by hand since we've+-- already put a call to 'nlopt_destroy' into the 'ForeignPtr', but+-- it's available in the C interface.+destroy :: Opt -> IO ()+destroy = finalizeForeignPtr . pointerFromOpt++foreign import ccall "nlopt.h nlopt_copy"+ nlopt_copy :: NloptOpt -> IO (NloptOpt)++copy :: Opt -> IO Opt+copy = useOpt $ \inp -> do+ outp <- nlopt_copy inp+ Opt <$> CFP.newForeignPtr outp (nlopt_destroy outp)++{- Random seeding functions -}++foreign import ccall "nlopt.h nlopt_srand"+ nlopt_srand :: CUInt -> IO ()++srand :: Integral a => a -> IO ()+srand = nlopt_srand . fromIntegral++foreign import ccall "nlopt.h nlopt_srand_time"+ nlopt_srand_time :: IO ()++srand_time :: IO ()+srand_time = nlopt_srand_time++{- Metadata -}++foreign import ccall "nlopt.h nlopt_version"+ nlopt_version :: Ptr CInt -> Ptr CInt -> Ptr CInt -> IO ()++-- | NLOPT library version, e.g. @2.4.2@+data Version = Version+ { major :: Int+ , minor :: Int+ , bugfix :: Int+ } deriving (Eq, Ord, Read, Show)++version :: IO Version+version =+ alloca $ \majptr ->+ alloca $ \minptr ->+ alloca $ \bfptr -> do+ nlopt_version majptr minptr bfptr+ Version <$> pk majptr <*> pk minptr <*> pk bfptr+ where+ pk = fmap fromIntegral . peek++foreign import ccall "nlopt.h nlopt_get_algorithm"+ nlopt_get_algorithm :: NloptOpt -> IO CInt++get_algorithm :: Opt -> IO Algorithm+get_algorithm = useOpt $ fmap parseEnum . nlopt_get_algorithm++foreign import ccall "nlopt.h nlopt_get_dimension"+ nlopt_get_dimension :: NloptOpt -> IO CUInt++get_dimension :: Opt -> IO Word+get_dimension = useOpt $ fmap fromIntegral . nlopt_get_dimension++{- Callback functions -}++asMVector :: CUInt -> Ptr CDouble -> IO (MV.IOVector Double)+asMVector dim ptr =+ MV.unsafeCast . flip MV.unsafeFromForeignPtr0 (fromIntegral dim) <$>+ newForeignPtr_ ptr++asVector :: CUInt -> Ptr CDouble -> IO (V.Vector Double)+asVector dim ptr =+ V.unsafeCast . flip V.unsafeFromForeignPtr0 (fromIntegral dim) <$>+ newForeignPtr_ ptr++type CFunc a = CUInt -> Ptr CDouble -> Ptr CDouble -> StablePtr a -> IO CDouble++-- | This function type corresponds to @nlopt_func@ in C and is used+-- for scalar functions of the parameter vector. You may pass data of+-- any type @a@ to the functions in this module that take a+-- 'ScalarFunction' as an argument; this data will be supplied to your+-- your function when it is called.+type ScalarFunction a+ = V.Vector Double -- ^ Parameter vector+ -> Maybe (MV.IOVector Double) -- ^ Gradient vector to be filled in+ -> a -- ^ User data+ -> IO Double -- ^ Scalar result++-- | This function type corresponds to @nlopt_mfunc@ in C and is used+-- for vector functions of the parameter vector. You may pass data of+-- any type @a@ to the functions in this module that take a+-- 'VectorFunction' as an argument; this data will be supplied to your+-- function when it is called.+type VectorFunction a+ = V.Vector Double -- ^ Parameter vector+ -> MV.IOVector Double -- ^ Output vector to be filled in+ -> Maybe (MV.IOVector Double) -- ^ Gradient vector to be filled in+ -> a -- ^ User data+ -> IO ()++-- | This function type corresponds to @nlopt_precond@ in C and is+-- used for functions that precondition a vector at a given point in+-- the parameter space. You may pass data of any type @a@ to the+-- functions in this module that take a 'PreconditionerFunction' as an+-- argument; this data will be supplied to your function when it is+-- called.+type PreconditionerFunction a+ = V.Vector Double -- ^ Parameter vector+ -> V.Vector Double -- ^ Vector @v@ to precondition+ -> MV.IOVector Double -- ^ Output vector @vpre@ to be filled in+ -> a -- ^ User data+ -> IO ()++wrapCFunction :: ScalarFunction a -> CFunc a+wrapCFunction cfunc dim stateptr gradientptr userptr = do+ nloptgradient <- asMVector dim gradientptr+ statevec <- asVector dim stateptr+ userdata <- deRefStablePtr userptr+ let+ gradptr = if gradientptr /= nullPtr+ then Just nloptgradient+ else Nothing+ realToFrac <$> cfunc statevec gradptr userdata++foreign import ccall safe "wrapper"+ mkCFunction :: CFunc a -> IO (FunPtr (CFunc a))++type CMFunc a = CUInt -> Ptr CDouble -> CUInt -> Ptr CDouble+ -> Ptr CDouble -> StablePtr a -> IO ()++wrapMFunction :: VectorFunction a -> CMFunc a+wrapMFunction mfunc constrdim constrptr dim stateptr gradientptr userptr+ = do+ nloptgradient <- asMVector (dim * constrdim) gradientptr+ nloptconstraint <- asMVector constrdim constrptr+ statevec <- asVector dim stateptr+ userdata <- deRefStablePtr userptr+ let+ gradptr = if gradientptr /= nullPtr+ then Just nloptgradient+ else Nothing+ mfunc statevec nloptconstraint gradptr userdata++foreign import ccall safe "wrapper"+ mkMFunction :: CMFunc a -> IO (FunPtr (CMFunc a))++type CPrecond a = CUInt -> Ptr CDouble -> Ptr CDouble+ -> Ptr CDouble -> StablePtr a -> IO ()++wrapPreconditioner :: PreconditionerFunction a -> CPrecond a+wrapPreconditioner prec dim stateptr vptr preptr userptr = do+ nloptpre <- asMVector dim preptr+ statevec <- asVector dim stateptr+ vvec <- asVector dim vptr+ userdata <- deRefStablePtr userptr+ prec statevec vvec nloptpre userdata++foreign import ccall safe "wrapper"+ mkPreconditionerFunction :: CPrecond a -> IO (FunPtr (CPrecond a))++-- We have to do the same silly dance with our user-data 'StablePtr's+-- as we do with function pointer wrappers: because NLOPT expects+-- these pointers before the actual optimization run, we have to+-- attach finalizers for them to the 'Opt' object so that they get+-- cleaned up properly.+addStablePtrFinalizer :: Opt -> StablePtr a -> IO ()+addStablePtrFinalizer (Opt p) sp =+ CFP.addForeignPtrFinalizer p (freeStablePtr sp)++getStablePtr :: Opt -> a -> IO (StablePtr a)+getStablePtr opt a = do+ aptr <- newStablePtr a+ addStablePtrFinalizer opt aptr+ return aptr++exportFunPtr :: (t1 -> IO (FunPtr a)) -> (t -> t1) -> t -> Opt -> IO (FunPtr a)+exportFunPtr mk wrap fun opt = do+ funptr <- mk $ wrap fun+ addFunPtrFinalizer opt funptr+ return funptr++{- Invoking the optimizer -}++-- | The output of an NLOPT optimizer run.+data Output = Output+ { resultCode :: Result -- ^ Return code+ , resultCost :: Double -- ^ Minimum of the objective+ -- function if optimization+ -- succeeded+ , resultParameters :: V.Vector Double -- ^ Parameters corresponding+ -- to the minimum if+ -- optimization succeeded++ , nEvals :: Int -- ^ number of evaluations+ }++foreign import ccall "nlopt.h nlopt_optimize"+ nlopt_optimize :: NloptOpt -> Ptr CDouble -> Ptr CDouble -> IO CInt++-- | This function is very similar to the C function @nlopt_optimize@,+-- but it does not use mutable vectors and returns an 'Output'+-- structure.+optimize :: Opt -- ^ Optimizer object set up to solve the problem+ -> V.Vector Double -- ^ Initial-guess parameter vector+ -> IO Output -- ^ Results of the optimization run+optimize optimizer x0 = withOpt optimizer $ \opt -> do+ vmut <- V.thaw $ V.unsafeCast x0+ (result, outputCost, iceout) <- alloca $ \costPtr -> do+ result <- MV.unsafeWith vmut $ \xptr ->+ parseEnum <$> nlopt_optimize opt xptr costPtr+ outputCost <- peek . castPtr $ costPtr+ iceout <- V.unsafeFreeze (MV.unsafeCast vmut)+ return (result, outputCost, iceout)+ nEvals <- fromIntegral <$> get_numevals optimizer+ return $ Output result outputCost iceout nEvals++{- Objective function setup -}++foreign import ccall "nlopt.h nlopt_set_min_objective"+ nlopt_set_min_objective :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> IO CInt++foreign import ccall "nlopt.h nlopt_set_max_objective"+ nlopt_set_max_objective :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> IO CInt++set_min_objective :: Opt -> ScalarFunction a -> a -> IO Result+set_min_objective opt objf userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_set_min_objective o objfunptr userptr++set_max_objective :: Opt -> ScalarFunction a -> a -> IO Result+set_max_objective opt objf userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_max_objective o objfunptr userptr++foreign import ccall "nlopt.h nlopt_set_precond_min_objective"+ nlopt_set_precond_min_objective :: NloptOpt+ -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a)+ -> StablePtr a+ -> IO CInt++foreign import ccall "nlopt.h nlopt_set_precond_max_objective"+ nlopt_set_precond_max_objective :: NloptOpt+ -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a)+ -> StablePtr a+ -> IO CInt++set_precond_min_objective :: Opt+ -> ScalarFunction a+ -> PreconditionerFunction a+ -> a+ -> IO Result+set_precond_min_objective opt objf pref userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ prefunptr <- exportFunPtr mkPreconditionerFunction wrapPreconditioner pref opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o -> parseEnum <$>+ nlopt_set_precond_min_objective o objfunptr prefunptr userptr++set_precond_max_objective :: Opt+ -> ScalarFunction a+ -> PreconditionerFunction a+ -> a+ -> IO Result+set_precond_max_objective opt objf pref userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ prefunptr <- exportFunPtr mkPreconditionerFunction wrapPreconditioner pref opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o -> parseEnum <$>+ nlopt_set_precond_max_objective o objfunptr prefunptr userptr++{- Working with bounds -}++foreign import ccall "nlopt.h nlopt_set_lower_bounds"+ nlopt_set_lower_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_lower_bounds1"+ nlopt_set_lower_bounds1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_lower_bounds"+ nlopt_get_lower_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_upper_bounds"+ nlopt_set_upper_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_upper_bounds1"+ nlopt_set_upper_bounds1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_upper_bounds"+ nlopt_get_upper_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++set_lower_bounds :: Opt -> V.Vector Double -> IO Result+set_lower_bounds opt bounds =+ withForeignPtr (fst . V.unsafeToForeignPtr0 . V.unsafeCast $ bounds) $+ \bptr -> withOpt opt $ \o ->+ parseEnum <$> nlopt_set_lower_bounds o bptr++set_lower_bounds1 :: Opt -> Double -> IO Result+set_lower_bounds1 opt bound =+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_lower_bounds1 o (realToFrac bound)++get_lower_bounds :: Opt -> IO (V.Vector Double, Result)+get_lower_bounds opt = do+ v <- get_dimension opt >>= MV.new . fromIntegral+ MV.unsafeWith (MV.unsafeCast v) $ \vptr -> withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_lower_bounds o vptr+ retv <- V.unsafeFreeze v+ return (retv, result)++set_upper_bounds :: Opt -> V.Vector Double -> IO Result+set_upper_bounds opt bounds =+ withForeignPtr (fst . V.unsafeToForeignPtr0 . V.unsafeCast $ bounds) $+ \bptr -> withOpt opt $ \o ->+ parseEnum <$> nlopt_set_upper_bounds o bptr++set_upper_bounds1 :: Opt -> Double -> IO Result+set_upper_bounds1 opt bound =+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_upper_bounds1 o (realToFrac bound)++get_upper_bounds :: Opt -> IO (V.Vector Double, Result)+get_upper_bounds opt = do+ v <- get_dimension opt >>= MV.new . fromIntegral+ MV.unsafeWith (MV.unsafeCast v) $ \vptr -> withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_upper_bounds o vptr+ retv <- V.unsafeFreeze v+ return (retv, result)++{- Working with constraints -}++foreign import ccall "nlopt.h nlopt_remove_inequality_constraints"+ nlopt_remove_inequality_constraints :: NloptOpt -> IO CInt++foreign import ccall "nlopt.h nlopt_add_inequality_constraint"+ nlopt_add_inequality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_precond_inequality_constraint"+ nlopt_add_precond_inequality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a) -> StablePtr a+ -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_inequality_mconstraint"+ nlopt_add_inequality_mconstraint :: NloptOpt -> CUInt -> FunPtr (CMFunc a)+ -> StablePtr a -> CDouble -> IO CInt++remove_inequality_constraints :: Opt -> IO Result+remove_inequality_constraints =+ useOpt $ fmap parseEnum . nlopt_remove_inequality_constraints++add_inequality_constraint :: Opt -> ScalarFunction a+ -> a -> Double -> IO Result+add_inequality_constraint opt objfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_inequality_constraint o objfunptr userptr (realToFrac tol)++add_precond_inequality_constraint :: Opt -> ScalarFunction a+ -> PreconditionerFunction a -> a -> Double+ -> IO Result+add_precond_inequality_constraint opt objfun precfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ precfunptr <-+ exportFunPtr mkPreconditionerFunction wrapPreconditioner precfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_precond_inequality_constraint o objfunptr+ precfunptr userptr (realToFrac tol)++add_inequality_mconstraint :: Opt -> Word -> VectorFunction a -> a+ -> Double -> IO Result+add_inequality_mconstraint opt constraintsize constrfun userdata tol = do+ constrfunptr <- exportFunPtr mkMFunction wrapMFunction constrfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_inequality_mconstraint o (fromIntegral constraintsize)+ constrfunptr userptr (realToFrac tol)++foreign import ccall "nlopt.h nlopt_remove_equality_constraints"+ nlopt_remove_equality_constraints :: NloptOpt -> IO CInt++foreign import ccall "nlopt.h nlopt_add_equality_constraint"+ nlopt_add_equality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_precond_equality_constraint"+ nlopt_add_precond_equality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a) -> StablePtr a+ -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_equality_mconstraint"+ nlopt_add_equality_mconstraint :: NloptOpt -> CUInt -> FunPtr (CMFunc a)+ -> StablePtr a -> CDouble -> IO CInt++remove_equality_constraints :: Opt -> IO Result+remove_equality_constraints =+ useOpt $ fmap parseEnum . nlopt_remove_equality_constraints++add_equality_constraint :: Opt -> ScalarFunction a+ -> a -> Double -> IO Result+add_equality_constraint opt objfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_equality_constraint o objfunptr userptr (realToFrac tol)++add_precond_equality_constraint :: Opt -> ScalarFunction a+ -> PreconditionerFunction a -> a -> Double+ -> IO Result+add_precond_equality_constraint opt objfun precfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ precfunptr <-+ exportFunPtr mkPreconditionerFunction wrapPreconditioner precfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_precond_equality_constraint o objfunptr+ precfunptr userptr (realToFrac tol)++add_equality_mconstraint :: Opt -> Word -> VectorFunction a -> a+ -> Double -> IO Result+add_equality_mconstraint opt constraintsize constrfun userdata tol = do+ constrfunptr <- exportFunPtr mkMFunction wrapMFunction constrfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_equality_mconstraint o (fromIntegral constraintsize)+ constrfunptr userptr (realToFrac tol)++{- Stopping criteria -}++withInputVector :: (Storable c, Storable a)+ => V.Vector c -> (Ptr a -> IO b) -> IO b+withInputVector = withForeignPtr . fst . V.unsafeToForeignPtr0 . V.unsafeCast+withOutputVector :: (Storable c, Storable a)+ => V.MVector s c -> (Ptr a -> IO b) -> IO b+withOutputVector = withForeignPtr . fst . MV.unsafeToForeignPtr0 . MV.unsafeCast++setScalar :: (Enum a, Integral b) => (NloptOpt -> t1 -> IO b)+ -> (t -> t1) -> Opt -> t -> IO a+setScalar setter conv opt val = withOpt opt $ \o ->+ parseEnum <$> setter o (conv val)++getScalar :: (NloptOpt -> IO b) -> (b -> a) -> Opt -> IO a+getScalar getter conv = useOpt $ fmap conv . getter++foreign import ccall "nlopt.h nlopt_set_stopval"+ nlopt_set_stopval :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_stopval"+ nlopt_get_stopval :: NloptOpt -> IO CDouble++set_stopval :: Opt -> Double -> IO Result+set_stopval = setScalar nlopt_set_stopval realToFrac++get_stopval :: Opt -> IO Double+get_stopval = getScalar nlopt_get_stopval realToFrac++foreign import ccall "nlopt.h nlopt_set_ftol_rel"+ nlopt_set_ftol_rel :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_ftol_rel"+ nlopt_get_ftol_rel :: NloptOpt -> IO CDouble++set_ftol_rel :: Opt -> Double -> IO Result+set_ftol_rel = setScalar nlopt_set_ftol_rel realToFrac++get_ftol_rel :: Opt -> IO Double+get_ftol_rel = getScalar nlopt_get_ftol_rel realToFrac++foreign import ccall "nlopt.h nlopt_set_ftol_abs"+ nlopt_set_ftol_abs :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_ftol_abs"+ nlopt_get_ftol_abs :: NloptOpt -> IO CDouble++set_ftol_abs :: Opt -> Double -> IO Result+set_ftol_abs = setScalar nlopt_set_ftol_abs realToFrac++get_ftol_abs :: Opt -> IO Double+get_ftol_abs = getScalar nlopt_get_ftol_abs realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_rel"+ nlopt_set_xtol_rel :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_xtol_rel"+ nlopt_get_xtol_rel :: NloptOpt -> IO CDouble++set_xtol_rel :: Opt -> Double -> IO Result+set_xtol_rel = setScalar nlopt_set_xtol_rel realToFrac++get_xtol_rel :: Opt -> IO Double+get_xtol_rel = getScalar nlopt_get_xtol_rel realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_abs1"+ nlopt_set_xtol_abs1 :: NloptOpt -> CDouble -> IO CInt++set_xtol_abs1 :: Opt -> Double -> IO Result+set_xtol_abs1 = setScalar nlopt_set_xtol_abs1 realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_abs"+ nlopt_set_xtol_abs :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_xtol_abs"+ nlopt_get_xtol_abs :: NloptOpt -> Ptr CDouble -> IO CInt++set_xtol_abs :: Opt -> V.Vector Double -> IO Result+set_xtol_abs opt tolvec =+ withInputVector tolvec $ \tolptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_xtol_abs o tolptr++get_xtol_abs :: Opt -> IO (Result, V.Vector Double)+get_xtol_abs opt = do+ mutv <- get_dimension opt >>= MV.new . fromIntegral+ withOutputVector mutv $ \vecptr ->+ withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_xtol_abs o vecptr+ outvec <- V.unsafeFreeze mutv+ return (result, outvec)++foreign import ccall "nlopt.h nlopt_set_maxeval"+ nlopt_set_maxeval :: NloptOpt -> CInt -> IO CInt++foreign import ccall "nlopt.h nlopt_get_maxeval"+ nlopt_get_maxeval :: NloptOpt -> IO CInt++set_maxeval :: Opt -> Word -> IO Result+set_maxeval = setScalar nlopt_set_maxeval fromIntegral++get_maxeval :: Opt -> IO Word+get_maxeval = getScalar nlopt_get_maxeval fromIntegral++foreign import ccall "nlopt.h nlopt_get_numevals"+ nlopt_get_numevals :: NloptOpt -> IO CInt++get_numevals :: Opt -> IO CInt+get_numevals = getScalar nlopt_get_numevals fromIntegral++foreign import ccall "nlopt.h nlopt_set_maxtime"+ nlopt_set_maxtime :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_maxtime"+ nlopt_get_maxtime :: NloptOpt -> IO CDouble++set_maxtime :: Opt -> Double -> IO Result+set_maxtime = setScalar nlopt_set_maxtime realToFrac++get_maxtime :: Opt -> IO Double+get_maxtime = getScalar nlopt_get_maxtime realToFrac++foreign import ccall "nlopt.h nlopt_force_stop"+ nlopt_force_stop :: NloptOpt -> IO CInt++force_stop :: Opt -> IO Result+force_stop = useOpt $ fmap parseEnum . nlopt_force_stop++foreign import ccall "nlopt.h nlopt_set_force_stop"+ nlopt_set_force_stop :: NloptOpt -> CInt -> IO CInt++foreign import ccall "nlopt.h nlopt_get_force_stop"+ nlopt_get_force_stop :: NloptOpt -> IO CInt++set_force_stop :: Opt -> Word -> IO Result+set_force_stop = setScalar nlopt_set_force_stop fromIntegral++get_force_stop :: Opt -> IO Word+get_force_stop = getScalar nlopt_get_force_stop fromIntegral++{- Algorithm-specific configuration -}++foreign import ccall "nlopt.h nlopt_set_local_optimizer"+ nlopt_set_local_optimizer :: NloptOpt -> NloptOpt -> IO CInt++set_local_optimizer :: Opt -- ^ Primary optimizer+ -> Opt -- ^ Subsidiary (local) optimizer+ -> IO Result+set_local_optimizer p s =+ withOpt p $ \primary -> withOpt s $ \secondary ->+ parseEnum <$> nlopt_set_local_optimizer primary secondary++foreign import ccall "nlopt.h nlopt_set_population"+ nlopt_set_population :: NloptOpt -> Word -> IO CInt++foreign import ccall "nlopt.h nlopt_get_population"+ nlopt_get_population :: NloptOpt -> IO Word++set_population :: Opt -> Word -> IO Result+set_population = setScalar nlopt_set_population fromIntegral++get_population :: Opt -> IO Word+get_population = getScalar nlopt_get_population fromIntegral++foreign import ccall "nlopt.h nlopt_set_vector_storage"+ nlopt_set_vector_storage :: NloptOpt -> Word -> IO CInt++foreign import ccall "nlopt.h nlopt_get_vector_storage"+ nlopt_get_vector_storage :: NloptOpt -> IO Word++set_vector_storage :: Opt -> Word -> IO Result+set_vector_storage = setScalar nlopt_set_vector_storage fromIntegral++get_vector_storage :: Opt -> IO Word+get_vector_storage = getScalar nlopt_get_vector_storage fromIntegral++foreign import ccall "nlopt.h nlopt_set_default_initial_step"+ nlopt_set_default_initial_step :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_initial_step"+ nlopt_set_initial_step :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_initial_step1"+ nlopt_set_initial_step1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_initial_step"+ nlopt_get_initial_step :: NloptOpt -> Ptr CDouble -> Ptr CDouble -> IO CInt++set_default_initial_step :: Opt -> V.Vector Double -> IO Result+set_default_initial_step opt stepvec =+ withInputVector stepvec $ \stepptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_default_initial_step o stepptr++set_initial_step :: Opt -> V.Vector Double -> IO Result+set_initial_step opt stepvec =+ withInputVector stepvec $ \stepptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_initial_step o stepptr++set_initial_step1 :: Opt -> Double -> IO Result+set_initial_step1 = setScalar nlopt_set_initial_step1 realToFrac++get_initial_step :: Opt -> V.Vector Double -> IO (Result, V.Vector Double)+get_initial_step opt xvec = do+ mutv <- get_dimension opt >>= MV.new . fromIntegral+ withOutputVector mutv $ \outptr ->+ withInputVector xvec $ \inptr ->+ withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_initial_step o inptr outptr+ outvec <- V.unsafeFreeze mutv+ return (result, outvec)
+ src/Text/ParseSR.hs view
@@ -0,0 +1,402 @@+{-# language OverloadedStrings #-}+-----------------------------------------------------------------------------+-- |+-- Module : Text.ParseSR+-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : ConstraintKinds+--+-- Functions to parse a string representing an expression+--+-----------------------------------------------------------------------------+module Text.ParseSR ( parseSR, parsePat, parseNonTerms, showOutput, SRAlgs(..), Output(..) )+ where++import Control.Applicative ((<|>))+import Data.Attoparsec.ByteString.Char8+import Data.Attoparsec.Expr+import qualified Data.ByteString.Char8 as B+import Data.Char (toLower)+import Data.List (sortOn)+import Data.SRTree+import Algorithm.EqSat.DB+import qualified Data.SRTree.Print as P+import qualified Data.Map.Strict as Map+import Data.List.Split ( splitOn )++import Debug.Trace (trace, traceShow)++-- * Data types++-- | Parser of a symbolic regression tree with `Int` variable index and+-- numerical values represented as `Double`. The numerical values type+-- can be changed with `fmap`.+type ParseTree = Parser (Fix SRTree)+type ParsePat = Parser Pattern++-- * Data types and caller functions++-- | Supported algorithms.+data SRAlgs = TIR | HL | OPERON | BINGO | GOMEA | PYSR | SBP | EPLEX deriving (Show, Read, Enum, Bounded)++-- | Supported outputs.+data Output = PYTHON | MATH | TIKZ | LATEX deriving (Show, Read, Enum, Bounded)++-- | Returns the corresponding function from Data.SRTree.Print for a given `Output`.+showOutput :: Output -> Fix SRTree -> String+showOutput PYTHON = P.showPython+showOutput MATH = P.showExpr+showOutput TIKZ = P.showTikz+showOutput LATEX = P.showLatex++-- | Calls the corresponding parser for a given `SRAlgs`+--+-- >>> fmap (showOutput MATH) $ parseSR OPERON "lambda,theta" False "lambda ^ 2 - sin(theta*3*lambda)"+-- Right "((x0 ^ 2.0) - Sin(((x1 * 3.0) * x0)))"+parseSR :: SRAlgs -> B.ByteString -> Bool -> B.ByteString -> Either String (Fix SRTree)+parseSR HL header reparam = eitherResult . (`feed` "") . parse (parseHL True reparam $ splitHeader header) . putEOL . B.strip+parseSR BINGO header reparam = eitherResult . (`feed` "") . parse (parseBingo True reparam $ splitHeader header) . putEOL . B.strip+parseSR TIR header reparam = eitherResult . (`feed` "") . parse (parseTIR True reparam $ splitHeader header) . putEOL . B.strip+parseSR OPERON header reparam = eitherResult . (`feed` "") . parse (parseOperon True reparam $ splitHeader header) . putEOL . B.strip+parseSR GOMEA header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR SBP header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR EPLEX header reparam = eitherResult . (`feed` "") . parse (parseGOMEA True reparam $ splitHeader header) . putEOL . B.strip+parseSR PYSR header reparam = eitherResult . (`feed` "") . parse (parsePySR True reparam $ splitHeader header) . putEOL . B.strip++parsePat :: B.ByteString -> Either String Pattern+parsePat = eitherResult . (`feed` "") . parse parsePatExpr . putEOL . B.strip++eitherResult' :: Show r => Result r -> Either String r+eitherResult' res = trace (show res) $ eitherResult res++-- * Parsers++-- | Creates a parser for a binary operator+binary :: B.ByteString -> (a -> a -> a) -> Assoc -> Operator B.ByteString a+binary name fun = Infix (do{ string (B.cons ' ' (B.snoc name ' ')) <|> string name; pure fun })++-- | Creates a parser for a unary function+prefix :: B.ByteString -> (a -> a) -> Operator B.ByteString a+prefix name fun = Prefix (do{ string name; pure fun })++-- | Envelopes the parser in parens+parens :: Parser a -> Parser a+parens e = do{ string "("; e' <- e; string ")"; pure e' } <?> "parens"++-- | Parse an expression using a user-defined parser given by the `Operator` lists containing+-- the name of the functions and operators of that SR algorithm, a list of parsers `binFuns` for binary functions+-- a parser `var` for variables, a boolean indicating whether to change floating point values to free+-- parameters variables, and a list of variable names with their corresponding indexes.+parseExpr :: Bool -> [[Operator B.ByteString (Fix SRTree)]] -> [ParseTree -> ParseTree] -> ParseTree -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseExpr relabel table binFuns var reparam header =+ do e <- if relabel then (relabelParams <$> expr) else expr+ many1' space+ pure e+ where+ term = parens expr <|> enclosedAbs expr <|> choice (map ($ expr) binFuns) <|> coef <|> varC <?> "term"+ expr = buildExpressionParser table term+ coef = if reparam + then do eNumber <- intOrDouble+ case eNumber of+ Left x -> pure $ fromIntegral x+ Right _ -> pure $ param 0+ else Fix . Const <$> signed double <?> "const"+ varC = if null header+ then var+ else var <|> varHeader++ varHeader = choice $ map (uncurry getParserVar) $ sortOn (negate . B.length . fst) header+ getParserVar k v = (string k <|> enveloped k) >> pure (Fix $ Var v)+ enveloped s = (char ' ' <|> char '(') >> string s >> (char ' ' <|> char ')') >> pure ""++enumerate :: [a] -> [(a, Int)]+enumerate = (`zip` [0..])++splitHeader :: B.ByteString -> [(B.ByteString, Int)]+splitHeader = enumerate . B.split ','++-- | Tries to parse as an `Int`, if it fails, +-- parse as a Double.+intOrDouble :: Parser (Either Int Double)+intOrDouble = eitherP parseInt (signed double)+ where+ parseInt :: Parser Int+ parseInt = do x <- signed decimal+ c <- peekChar+ case c of + Just '.' -> digit >> pure 0+ Just 'e' -> digit >> pure 0+ Just 'E' -> digit >> pure 0+ _ -> pure x++putEOL :: B.ByteString -> B.ByteString+putEOL bs | B.last bs == '\n' = bs+ | otherwise = B.snoc bs '\n'++-- * Special case functions++-- | analytic quotient+aq :: Fix SRTree -> Fix SRTree -> Fix SRTree+aq x y = x / sqrt (1 + y ** 2)++log1p :: Fix SRTree -> Fix SRTree+log1p x = log (1 + x)++log10 :: Fix SRTree -> Fix SRTree+log10 x = log x / log 10++log2 :: Fix SRTree -> Fix SRTree+log2 x = log x / log 2++cbrt :: Fix SRTree -> Fix SRTree+cbrt x = x ** (1/3)++cube :: Fix SRTree -> Fix SRTree+cube x = Fix $ Uni Cube x++sqrtabs :: Fix SRTree -> Fix SRTree+sqrtabs x = Fix $ Uni SqrtAbs x++logabs :: Fix SRTree -> Fix SRTree+logabs x = Fix $ Uni LogAbs x++-- Parse `abs` functions as | x |+enclosedAbs :: Num a => Parser a -> Parser a+enclosedAbs expr = do char '|'+ e <- expr+ char '|'+ pure $ abs e++-- | Parser for binary functions+binFun :: B.ByteString -> (a -> a -> a) -> Parser a -> Parser a+binFun name f expr = do string name+ many' space >> char '(' >> many' space+ e1 <- expr+ many' space >> char ',' >> many' space -- many' space >> char ',' >> many' space+ e2 <- expr+ many' space >> char ')'+ pure $ f e1 e2 ++-- * Custom parsers for SR algorithms++-- | parser for Transformation-Interaction-Rational.+parseTIR :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseTIR b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = [ ]+ prefixOps = map (uncurry prefix)+ [ ("id", id), ("abs", abs)+ , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh)+ , ("sin", sin), ("cos", cos), ("tan", tan)+ , ("asinh", asinh), ("acosh", acosh), ("atanh", atanh)+ , ("asin", asin), ("acos", acos), ("atan", atan)+ , ("sqrtabs", sqrtabs), ("sqrt", sqrt), ("cbrt", cbrt), ("square", (**2))+ , ("logabs", logabs), ("log", log), ("exp", exp), ("cube", cube), ("recip", recip)+ , ("Id", id), ("Abs", abs)+ , ("Sinh", sinh), ("Cosh", cosh), ("Tanh", tanh)+ , ("Sin", sin), ("Cos", cos), ("Tan", tan)+ , ("ASinh", asinh), ("ACosh", acosh), ("ATanh", atanh)+ , ("ASin", asin), ("ACos", acos), ("ATan", atan)+ , ("SqrtAbs", sqrtabs), ("Sqrt", sqrt), ("Cbrt", cbrt), ("Square", (**2))+ , ("LogAbs", logabs), ("Log", log), ("Exp", exp), ("Recip", recip), ("Cube", cube)+ ]+ binOps = [[binary "^" (**) AssocLeft], [binary "**" (**) AssocLeft]+ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ , [binary "|**|" powabs AssocLeft], [binary "aq" aq AssocLeft]+ ]+ powabs l r = Fix $ Bin PowerAbs l r+ aq l r = Fix $ Bin AQ l r++ var = do char 'x'+ ix <- decimal+ pure $ Fix $ Var ix+ <|> do char 't'+ ix <- decimal+ pure $ Fix $ Param ix+ <?> "var"++-- | parser for Operon.+parseOperon :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseOperon b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = [ binFun "pow" (**) ]+ prefixOps = map (uncurry prefix)+ [ ("abs", abs), ("cbrt", cbrt)+ , ("acos", acos), ("cosh", cosh), ("cos", cos)+ , ("asin", asin), ("sinh", sinh), ("sin", sin)+ , ("exp", exp), ("log", log)+ , ("sqrt", sqrt), ("square", (**2))+ , ("atan", atan), ("tanh", tanh), ("tan", tan)+ ]+ binOps = [[binary "^" (**) AssocLeft]+ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ ]+ var = do char 'X' <|> char 'x'+ ix <- decimal+ pure $ Fix $ Var (ix - 1) -- Operon is not 0-based+ <?> "var"++-- | parser for HeuristicLab.+parseHL :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseHL b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = [ binFun "aq" aq ]+ prefixOps = map (uncurry prefix)+ [ ("logabs", log.abs), ("sqrtabs", sqrt.abs) -- the longer versions should come first+ , ("abs", abs), ("exp", exp), ("log", log)+ , ("sqrt", sqrt), ("sqr", (**2)), ("cube", (**3))+ , ("cbrt", cbrt), ("sin", sin), ("cos", cos)+ , ("tan", tan), ("tanh", tanh)+ ]+ binOps = [[binary "^" (**) AssocLeft]+ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ ]+ var = do char 'x'+ ix <- decimal+ pure $ Fix $ Var ix+ <?> "var"++-- | parser for Bingo+parseBingo :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseBingo b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = []+ prefixOps = map (uncurry prefix)+ [ ("abs", abs), ("exp", exp), ("log", log.abs)+ , ("sqrt", sqrt.abs)+ , ("sinh", sinh), ("cosh", cosh)+ , ("sin", sin), ("cos", cos)+ ]+ binOps = [[binary "^" (**) AssocLeft]+ , [binary "/" (/) AssocLeft, binary "" (*) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ ]+ var = do string "X_"+ ix <- decimal+ pure $ Fix $ Var ix+ <?> "var"++-- | parser for GOMEA+parseGOMEA :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parseGOMEA b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = []+ prefixOps = map (uncurry prefix)+ [ ("exp", exp), ("plog", log.abs)+ , ("sqrt", sqrt.abs)+ , ("sin", sin), ("cos", cos)+ ]+ binOps = [[binary "^" (**) AssocLeft]+ , [binary "/" (/) AssocLeft, binary "*" (*) AssocLeft, binary "aq" aq AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ ]+ var = do string "x"+ ix <- decimal+ pure $ Fix $ Var ix+ <?> "var"++-- | parser for PySR+parsePySR :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree+parsePySR b = parseExpr b (prefixOps : binOps) binFuns var+ where+ binFuns = [ binFun "pow" (**) ]+ prefixOps = map (uncurry prefix)+ [ ("abs", abs), ("exp", exp)+ , ("square", (**2)), ("cube", (**3)), ("neg", negate)+ , ("acosh_abs", acosh . (+1) . abs), ("acosh", acosh), ("asinh", asinh)+ , ("acos", acos), ("asin", asin), ("atan", atan)+ , ("sqrt_abs", sqrt.abs), ("sqrt", sqrt)+ , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh)+ , ("sin", sin), ("cos", cos), ("tan", tan)+ , ("log10", log10), ("log2", log2), ("log1p", log1p) + , ("log_abs", log.abs), ("log10_abs", log10 . abs)+ , ("log", log)+ ]+ binOps = [[binary "^" (**) AssocLeft]+ , [binary "/" (/) AssocLeft, binary "*" (*) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ ]+ var = do string "x"+ ix <- decimal+ pure $ Fix $ Var ix+ <?> "var"++-- parse a pattern expression+parsePatExpr :: ParsePat+parsePatExpr = parsePattern (prefixOps : binOps) binFuns var+ where+ binFuns = [ ]+ prefixOps = map (uncurry prefix)+ [ ("id", id), ("abs", abs)+ , ("sinh", sinh), ("cosh", cosh), ("tanh", tanh)+ , ("sin", sin), ("cos", cos), ("tan", tan)+ , ("asinh", asinh), ("acosh", acosh), ("atanh", atanh)+ , ("asin", asin), ("acos", acos), ("atan", atan)+ , ("sqrtabs", sqrtabs'), ("sqrt", sqrt), ("cbrt", cbrt'), ("square", (**2))+ , ("logabs", logabs'), ("log", log), ("exp", exp), ("cube", cube'), ("recip", recip')+ , ("Id", id), ("Abs", abs)+ , ("Sinh", sinh), ("Cosh", cosh), ("Tanh", tanh)+ , ("Sin", sin), ("Cos", cos), ("Tan", tan)+ , ("ASinh", asinh), ("ACosh", acosh), ("ATanh", atanh)+ , ("ASin", asin), ("ACos", acos), ("ATan", atan)+ , ("SqrtAbs", sqrtabs'), ("Sqrt", sqrt), ("Cbrt", cbrt'), ("Square", (**2))+ , ("LogAbs", logabs'), ("Log", log), ("Exp", exp), ("Recip", recip'), ("Cube", cube')+ , ("|log|", logabs'), ("|Log|", logabs'), ("|sqrt|", sqrtabs'), ("|Sqrt|", sqrtabs')+ , ("√", sqrt), ("|√|", sqrtabs')+ ]+ binOps = [[binary "^" (**) AssocLeft], [binary "**" (**) AssocLeft]+ , [binary "*" (*) AssocLeft, binary "/" (/) AssocLeft]+ , [binary "+" (+) AssocLeft, binary "-" (-) AssocLeft]+ , [binary "|**|" powabs AssocLeft], [binary "|^|" powabs AssocLeft]+ , [binary "aq" aq AssocLeft], [binary "|/|" aq AssocLeft]+ ]+ powabs l r = Fixed $ Bin PowerAbs l r+ aq l r = Fixed $ Bin AQ l r+ logabs' t = Fixed $ Uni LogAbs t+ sqrtabs' t = Fixed $ Uni SqrtAbs t+ cbrt' t = Fixed $ Uni Cbrt t+ cube' t = Fixed $ Uni Cube t+ recip' t = Fixed $ Uni Recip t++ var = do char 'x'+ ix <- decimal+ pure $ Fixed $ Var ix+ <|> do char 't'+ ix <- decimal+ pure $ Fixed $ Param ix+ <|> do char 'v'+ ix <- decimal+ pure $ VarPat (toEnum $ ix+65)+ <?> "var"++parsePattern :: [[Operator B.ByteString Pattern]] -> [ParsePat -> ParsePat] -> ParsePat -> ParsePat+parsePattern table binFuns var =+ do e <- expr+ many1' space+ pure e+ where+ term = parens expr <|> enclosedAbs expr <|> choice (map ($ expr) binFuns) <|> coef <|> var <?> "term"+ expr = buildExpressionParser table term+ coef = Fixed . Const <$> signed double <?> "const"++ getParserVar k v = (string k <|> enveloped k) >> pure (Fix $ Var v)+ enveloped s = (char ' ' <|> char '(') >> string s >> (char ' ' <|> char ')') >> pure ""+++-- * Parse the non-terminal nodes into a SRTree () value+parseNonTerms :: String -> [SRTree ()]+parseNonTerms = Prelude.map toNonTerm . splitOn ","+ where+ binTerms = Map.fromList [ (Prelude.map toLower (show op), op) | op <- [Add .. AQ]]+ uniTerms = Map.fromList [ (Prelude.map toLower (show f), f) | f <- [Abs .. Cube]]+ toNonTerm xs' = let xs = Prelude.map toLower xs'+ in case binTerms Map.!? xs of+ Just op -> Bin op () ()+ Nothing -> case uniTerms Map.!? xs of+ Just f -> Uni f ()+ Nothing -> error $ "invalid non-terminal " <> show xs
+ src/Text/ParseSR/IO.hs view
@@ -0,0 +1,73 @@+{-# language LambdaCase #-}+-----------------------------------------------------------------------------+-- |+-- Module : Text.ParseSR.IO+-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : ConstraintKinds+--+-- Functions to parse multiple expressions from stdin or a text file.+--+-----------------------------------------------------------------------------+module Text.ParseSR.IO ( withInput, withOutput, withOutputDebug )+ where++-- import Data.SRTree.EqSat1+import Algorithm.EqSat.Simplify ( simplifyEqSatDefault )+import Control.Monad (forM_, unless)+import qualified Data.ByteString.Char8 as B+import Data.SRTree+import Data.SRTree.Recursion (Fix (..))+import System.IO+import Text.ParseSR (Output, SRAlgs, parseSR, showOutput)++-- | given a filename, the symbolic regression algorithm, a string of variables name, +-- and two booleans indicating whether to convert float values to parameters and +-- whether to simplify the expression or not, it will read the file and parse everything +-- returning a list of either an error message or a tree.+--+-- empty filename defaults to stdin +withInput :: String -> SRAlgs -> String -> Bool -> Bool -> IO [Either String (Fix SRTree)]+withInput fname sr hd param simpl = do+ h <- if null fname then pure stdin else openFile fname ReadMode+ contents <- hGetLines h + let myParserFun = parseSR sr (B.pack hd) param . B.pack+ -- myParser = if simpl then fmap simplifyEqSat . myParserFun else myParserFun+ myParser = if simpl then fmap simplifyEqSatDefault . myParserFun else myParserFun+ es = map myParser $ filter (not . null) contents+ unless (null fname) $ hClose h+ pure es++-- | outputs a list of either error or trees to a file using the Output format. +--+-- empty filename defaults to stdout +withOutput :: String -> Output -> [Either String (Fix SRTree)] -> IO ()+withOutput fname output exprs = do+ h <- if null fname then pure stdout else openFile fname WriteMode+ forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err+ Right ex -> hPutStrLn h (showOutput output ex)+ unless (null fname) $ hClose h++-- | debug version of output function to check the invalid parsers+withOutputDebug :: String -> Output -> [Either String (Fix SRTree, Fix SRTree)] -> IO ()+withOutputDebug fname output exprs = do+ h <- if null fname then pure stdout else openFile fname WriteMode+ forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err+ Right (t1, t2) -> do + hPutStrLn h ("First: " <> showOutput output t1)+ hPutStrLn h ("Second: " <> showOutput output t2)+ hFlush h+ unless (null fname) $ hClose h++hGetLines :: Handle -> IO [String]+hGetLines h = do+ done <- hIsEOF h+ if done+ then return []+ else do+ line <- hGetLine h+ (line :) <$> hGetLines h
srtree.cabal view
@@ -1,12 +1,12 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.35.2.+-- This file has been generated from package.yaml by hpack version 0.38.1. -- -- see: https://github.com/sol/hpack name: srtree-version: 1.0.0.5-synopsis: A general framework to work with Symbolic Regression expression trees.+version: 2.0.1.8+synopsis: A general library to work with Symbolic Regression expression trees. description: A Symbolic Regression Tree data structure to work with mathematical expressions with support to first order derivative and simplification; category: Math, Data, Data Structures homepage: https://github.com/folivetti/srtree#readme@@ -27,24 +27,185 @@ library exposed-modules:+ Algorithm.EqSat+ Algorithm.EqSat.Build+ Algorithm.EqSat.DB+ Algorithm.EqSat.Egraph+ Algorithm.EqSat.Info+ Algorithm.EqSat.Queries+ Algorithm.EqSat.SearchSR+ Algorithm.EqSat.SearchSRCache+ Algorithm.EqSat.Simplify+ Algorithm.Massiv.Utils+ Algorithm.SRTree.AD+ Algorithm.SRTree.ConfidenceIntervals+ Algorithm.SRTree.Likelihoods+ Algorithm.SRTree.ModelSelection+ Algorithm.SRTree.NonlinearOpt+ Algorithm.SRTree.Opt Data.SRTree+ Data.SRTree.Datasets+ Data.SRTree.Derivative+ Data.SRTree.Eval Data.SRTree.Internal Data.SRTree.Print Data.SRTree.Random Data.SRTree.Recursion+ Numeric.Optimization.NLOPT.Bindings+ Text.ParseSR+ Text.ParseSR.IO other-modules: Paths_srtree hs-source-dirs: src+ ghc-options: -fwarn-incomplete-patterns -threaded+ extra-libraries:+ nlopt build-depends:- base >=4.16 && <5- , containers ==0.6.*+ attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.9 , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.1 && <1.1 , mtl >=2.2 && <2.4- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5 && <0.3+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.* , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8 default-language: Haskell2010 +executable srsimplify+ main-is: Main.hs+ other-modules:+ Paths_srtree+ hs-source-dirs:+ apps/srsimplify+ ghc-options: -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.9+ , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.1 && <1.1+ , mtl >=2.2 && <2.4+ , optparse-applicative >=0.18 && <0.20+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5 && <0.3+ , srtree+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.*+ , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8+ default-language: Haskell2010++executable srtools+ main-is: Main.hs+ other-modules:+ Args+ IO+ Report+ Paths_srtree+ hs-source-dirs:+ apps/srtools+ ghc-options: -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.9+ , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.1 && <1.1+ , mtl >=2.2 && <2.4+ , optparse-applicative >=0.18 && <0.20+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5 && <0.3+ , srtree+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.*+ , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8+ default-language: Haskell2010++executable tinygp+ main-is: Main.hs+ other-modules:+ GP+ Initialization+ Util+ Paths_srtree+ hs-source-dirs:+ apps/tinygp+ ghc-options: -threaded -rtsopts -with-rtsopts=-N+ build-depends:+ attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.9+ , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.1 && <1.1+ , mtl >=2.2 && <2.4+ , optparse-applicative >=0.18 && <0.20+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5 && <0.3+ , srtree+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.*+ , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8+ default-language: Haskell2010+ test-suite srtree-test type: exitcode-stdio-1.0 main-is: Spec.hs@@ -56,11 +217,30 @@ build-depends: HUnit , ad- , base >=4.16 && <5- , containers ==0.6.*+ , attoparsec >=0.14.4 && <0.15+ , attoparsec-expr >=0.1.1.2 && <0.2+ , base >=4.19 && <5+ , binary >=0.8.9.1 && <0.9+ , bytestring >=0.11 && <0.13+ , containers >=0.6.7 && <0.9 , dlist ==1.0.*+ , exceptions >=0.10.7 && <0.11+ , filepath >=1.4.0.0 && <1.6+ , hashable >=1.4.4.0 && <1.6+ , ieee754 >=0.8.0 && <0.9+ , lens >=5.2.3 && <5.4+ , list-shuffle >=1.0.0.1 && <1.1+ , massiv >=1.0.4.1 && <1.1 , mtl >=2.2 && <2.4- , random ==1.2.*+ , random >=1.2 && <1.4+ , scheduler >=2.0.0.1 && <3+ , split >=0.2.5 && <0.3 , srtree+ , statistics >=0.16.2.1 && <0.17+ , transformers >=0.6.1.0 && <0.7+ , unliftio >=0.2.10 && <1+ , unliftio-core >=0.2.1 && <1+ , unordered-containers ==0.2.* , vector >=0.12 && <0.14+ , zlib >=0.6.3 && <0.8 default-language: Haskell2010
test/Spec.hs view
@@ -1,76 +1,4 @@-import Data.SRTree--import qualified Data.Vector as V-import Numeric.AD.Double ( grad ) import Test.HUnit --- test expressions-exprs = [- param 0 * sin ( param 1)- , sin (param 0) + cos (param 1)- , 0.5 * sin (param 0) + 0.7 * cos (param 1)- , log (param 0) + param 0 * param 1 - sin (param 1)- , 1 / param 0 * param 1- , param 0 + param 1 + param 0 * param 1 + sin (param 0) + sin (param 1) + cos (param 0) + cos (param 1) + sin (param 0 * param 1) + cos (param 0 * param 1)- , sin (exp (param 0) + param 1)- ]---- autodiff with multiple occurrences of vars-autoDiffMult :: [[Double]]-autoDiffMult = [ grad (\[x,y] -> x * sin y) [2,3]- , grad (\[x,y] -> sin x + cos y) [2,3]- , grad (\[x,y] -> 0.5 * sin x + 0.7 * cos y) [2,3]- , grad (\[x,y] -> log x + x*y - sin y) [2,3]- , grad (\[x,y] -> 1 / x * y) [2,3]- , grad (\[x,y] -> x + y + x * y + sin x + sin y + cos x + cos y + sin (x * y) + cos (x * y)) [2,3]- , grad (\[x,y] -> sin (exp x + y)) [2,3]- ]---- autodiff with single occurrences of vars-autoDiffSingle :: [[Double]]-autoDiffSingle = [ grad (\[x,y] -> x * sin y) [2,3]- , grad (\[x,y] -> sin x + cos y) [2,3]- , grad (\[x,y] -> 0.5 * sin x + 0.7 * cos y) [2,3]- , grad (\[x,y,v,w] -> log x + y*v - sin w) [2,3,2,3]- , grad (\[x,y] -> 1 / x * y) [2,3]- , grad (\[a,b,c,d,e,f,g,h,i,j,k,l] -> a + b + c * d + sin e + sin f + cos g + cos h + sin (i * j) + cos (k * l)) [2,3,2,3,2,3,2,3,2,3,2,3]- , grad (\[x,y] -> sin (exp x + y)) [2,3]- ]---- xs is empty since we are interested in theta-xs :: V.Vector a-xs = V.empty--- theta values-thetaMulti, thetaSingle :: V.Vector Double-thetaMulti = V.fromList [2.0, 3.0]-thetaSingle = V.fromList [2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0, 2.0, 3.0]---- values from forward mode-forwardVals :: [[Double]]-forwardVals = map (forwardMode xs thetaMulti id) exprs---- values from grad--- we must relabel the parameters of the expression to sequence values-gradVals :: [(Double, [Double])]-gradVals = map (gradParamsFwd xs thetaSingle id . relabelParams) exprs---- values of the evaluated expressions-exprVals :: [Double]-exprVals = map (evalTree xs thetaSingle id . relabelParams) exprs--refGrad :: [(Double, [Double])]-refGrad = zip exprVals autoDiffSingle--testDiff :: (Eq a, Show a) => String -> String -> a -> a -> Test-testDiff lbl name a b = TestLabel lbl $ TestCase (assertEqual name a b)--tests :: Test-tests = TestList $- zipWith (testDiff "forward mode" "autodiff x forward mode") autoDiffMult forwardVals- <> zipWith (testDiff "opt. grad. parameters" "(evalTree, autodiff) x gradVals") refGrad gradVals- <> zipWith (testDiff "deriveByParam" "deriveByParam x autodiff") (map head autoDiffSingle) (map (head.snd) gradVals)- main :: IO ()-main = do- result <- runTestTT tests- putStrLn $ showCounts result+main = pure ()