diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -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 diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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 ". + +### `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 diff --git a/apps/srsimplify/Main.hs b/apps/srsimplify/Main.hs new file mode 100644 --- /dev/null +++ b/apps/srsimplify/Main.hs @@ -0,0 +1,103 @@ +module Main (main) where + +import Options.Applicative +import Text.ParseSR.IO ( withInput, withOutput ) +import Text.ParseSR ( SRAlgs (..), Output (..) ) +import System.Random ( getStdGen, mkStdGen ) +import Text.Read ( readMaybe ) +import Data.Char ( toLower, toUpper ) +import Data.List ( intercalate ) + +-- Data type to store command line arguments +data Args = Args + { from :: SRAlgs + , to :: Output + , infile :: String + , outfile :: String + , varnames :: String + } deriving Show + +-- parser of command line arguments +opt :: Parser Args +opt = Args + <$> option sralgsReader + ( long "from" + <> short 'f' + <> metavar ("[" <> intercalate "|" sralgsHelp <> "]") + <> help "Input expression format" ) + <*> option srtoReader -- TODO + ( long "to" + <> short 't' + <> metavar ("[" <> intercalate "|" srtoHelp <> "]") + <> help "Output expression format" ) + <*> strOption + ( long "input" + <> short 'i' + <> metavar "INPUT-FILE" + <> showDefault + <> value "" + <> help "Input file containing expressions. \ + \ Empty string gets expression from stdin." ) + <*> strOption + ( long "output" + <> short 'o' + <> metavar "OUTPUT-FILE" + <> showDefault + <> value "" + <> help "Output file to store the stats in CSV format. \ + \ Empty string prints expressions to stdout." ) + <*> strOption + ( long "varnames" + <> short 'v' + <> metavar "VARNAMES" + <> showDefault + <> value "" + <> help "Comma separated string of variable names. \ + \ Empty string defaults to the algorithm default (x0, x1,..)." ) + +-- helper functions to show the possible options +mkDescription :: Show a => [a] -> [String] +mkDescription = map (envelope '\'' . map toLower . show) + where + envelope :: a -> [a] -> [a] + envelope c xs = c : xs <> [c] +{-# INLINE mkDescription #-} + +sralgsHelp :: [String] +sralgsHelp = mkDescription [toEnum 0 :: SRAlgs ..] +{-# INLINE sralgsHelp #-} + +srtoHelp :: [String] +srtoHelp = mkDescription [toEnum 0 :: Output ..] +{-# INLINE srtoHelp #-} + +-- helper functions to parse the options +mkReader :: Read a => String -> (a -> b) -> String -> ReadM b +mkReader err val sr = eitherReader + $ case readMaybe sr of + Nothing -> pure (Left err) + Just x -> pure (Right (val x)) + +sralgsReader :: ReadM SRAlgs +sralgsReader = + str >>= (mkReader errMsg id . map toUpper) + where + errMsg = "unknown algorithm. Available options are " <> intercalate "," sralgsHelp + +srtoReader :: ReadM Output +srtoReader = + str >>= (mkReader errMsg id . map toUpper) + where + errMsg = "unknown algorithm. Available options are " <> intercalate "," srtoHelp + +main :: IO () +main = do + args <- execParser opts + withInput (infile args) (from args) (varnames args) False True + >>= withOutput (outfile args) (to args) + where + opts = info (opt <**> helper) + ( fullDesc <> progDesc "Simplify an expression\ + \ using equality saturation." + <> header "srsimplify - a CLI tool to simplify\ + \ symbolic regression expressions with equality saturation." ) diff --git a/apps/srtools/Args.hs b/apps/srtools/Args.hs new file mode 100644 --- /dev/null +++ b/apps/srtools/Args.hs @@ -0,0 +1,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 diff --git a/apps/srtools/IO.hs b/apps/srtools/IO.hs new file mode 100644 --- /dev/null +++ b/apps/srtools/IO.hs @@ -0,0 +1,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 "=============================================================" diff --git a/apps/srtools/Main.hs b/apps/srtools/Main.hs new file mode 100644 --- /dev/null +++ b/apps/srtools/Main.hs @@ -0,0 +1,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" + ) diff --git a/apps/srtools/Report.hs b/apps/srtools/Report.hs new file mode 100644 --- /dev/null +++ b/apps/srtools/Report.hs @@ -0,0 +1,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 diff --git a/apps/tinygp/GP.hs b/apps/tinygp/GP.hs new file mode 100644 --- /dev/null +++ b/apps/tinygp/GP.hs @@ -0,0 +1,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' diff --git a/apps/tinygp/Initialization.hs b/apps/tinygp/Initialization.hs new file mode 100644 --- /dev/null +++ b/apps/tinygp/Initialization.hs @@ -0,0 +1,50 @@ +module Initialization where + +data InitiMethod = GROW | FULL | BTC | HALFHALF +{- +btc = undefined + +def btc(pset_, depth_, length_, type_=None): + if type_ is None: + type_ = pset_.ret + + expr = [] + + arities = list(map(lambda x: x.arity, pset_.primitives[type_])) + minFunctionArity = min(arities) + maxFunctionArity = max(arities) + + # adapt length to restrictions of the primitive set + if length_ % 2 == 0 and minFunctionArity > 1: + length_ = length_ + 1 if np.random.random_sample(1) > 0.5 else length_ - 1 + + targetLength = length_ - 1 # don't count the root node + maxFunctionArity = min(maxFunctionArity, targetLength) + minFunctionArity = min(minFunctionArity, targetLength) + root = sampleChild(pset_, minFunctionArity, maxFunctionArity, type_) + + # inner lists of the form [node, depth, childIndex] + # childIndex is only used at the end to transform + # the representation from breadth to prefix + expr.append([root, 0, 1]) + + openSlots = root.arity + + for i in range(0, length_): + (node, nodeDepth, childIndex) = expr[i] + childDepth = nodeDepth + 1 + + for j in range(0, getArity(node)): + maxArity = 0 if childDepth == depth_ - 1 else min(maxFunctionArity, targetLength - openSlots) + minArity = min(minFunctionArity, maxArity) + child = sampleChild(pset_, minArity, maxArity, type_) + + if j == 0: + expr[i][2] = len(expr) + + expr.append([child, childDepth, 0]) + openSlots += getArity(child) + + nodes = breadthToPrefix(expr) + return nodes + -} diff --git a/apps/tinygp/Main.hs b/apps/tinygp/Main.hs new file mode 100644 --- /dev/null +++ b/apps/tinygp/Main.hs @@ -0,0 +1,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." ) diff --git a/apps/tinygp/Util.hs b/apps/tinygp/Util.hs new file mode 100644 --- /dev/null +++ b/apps/tinygp/Util.hs @@ -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 + diff --git a/src/Algorithm/EqSat.hs b/src/Algorithm/EqSat.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat.hs @@ -0,0 +1,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 diff --git a/src/Algorithm/EqSat/Build.hs b/src/Algorithm/EqSat/Build.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Build.hs @@ -0,0 +1,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 #-} diff --git a/src/Algorithm/EqSat/DB.hs b/src/Algorithm/EqSat/DB.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/DB.hs @@ -0,0 +1,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 <> " | " + +-- 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 #-} diff --git a/src/Algorithm/EqSat/Egraph.hs b/src/Algorithm/EqSat/Egraph.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Egraph.hs @@ -0,0 +1,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 #-} diff --git a/src/Algorithm/EqSat/Info.hs b/src/Algorithm/EqSat/Info.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Info.hs @@ -0,0 +1,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 + diff --git a/src/Algorithm/EqSat/Queries.hs b/src/Algorithm/EqSat/Queries.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Queries.hs @@ -0,0 +1,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) + diff --git a/src/Algorithm/EqSat/SearchSR.hs b/src/Algorithm/EqSat/SearchSR.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/SearchSR.hs @@ -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'' diff --git a/src/Algorithm/EqSat/SearchSRCache.hs b/src/Algorithm/EqSat/SearchSRCache.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/SearchSRCache.hs @@ -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'' diff --git a/src/Algorithm/EqSat/Simplify.hs b/src/Algorithm/EqSat/Simplify.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Simplify.hs @@ -0,0 +1,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 diff --git a/src/Algorithm/Massiv/Utils.hs b/src/Algorithm/Massiv/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/Massiv/Utils.hs @@ -0,0 +1,278 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +module Algorithm.Massiv.Utils where + +import Data.Massiv.Array hiding ( forM_, unzip, map, init, zipWith, zip, tail, replicate, take ) +import qualified Data.Massiv.Array as A +import qualified Data.Massiv.Array.Unsafe as UMA +import qualified Data.Massiv.Array.Mutable as MMA +import Control.Monad +import Data.Vector.Storable ((//)) +import System.IO.Unsafe + +-- taken from https://hackage.haskell.org/package/cubicspline-0.1.2 +import Control.Arrow +import Data.List(unfoldr) + +import Data.SRTree.Eval + +type MMassArray m = MMA.MArray (PrimState m) S Ix2 Double + +getRows :: SRMatrix -> Array B Ix1 PVector +getRows = computeAs B . outerSlices +{-# INLINE getRows #-} +getCols :: SRMatrix -> Array B Ix1 PVector +getCols = computeAs B . A.map (computeAs S) . innerSlices +{-# INLINE getCols #-} + +appendRow :: MonadThrow m => SRMatrix -> PVector -> m SRMatrix +appendRow xs v = computeAs S <$> (stackOuterSlicesM . toList . computeAs B $ snoc (outerSlices xs) v) +{-# INLINE appendRow #-} + +appendCol :: MonadThrow m => SRMatrix -> PVector -> m SRMatrix +appendCol xs v = computeAs S <$> (stackInnerSlicesM . toList . computeAs B $ snoc (A.map (computeAs S) $ innerSlices xs) v) +{-# INLINE appendCol #-} + +updateS :: Array S Ix1 Double -> [(Int, Double)] -> Array S Ix1 Double +updateS vec new = fromStorableVector compMode $ toStorableVector vec // new + +linSpace :: Int -> (Double, Double) -> [Double] +linSpace num (lo, hi) = Prelude.take num $ iterate (\x -> x + step) lo + where + step = (hi - lo) / (fromIntegral num - 1) +{-# INLINE linSpace #-} + +outer :: (MonadThrow m) + => PVector + -> PVector + -> m SRMatrix +outer arr1 arr2 + | isEmpty arr1 || isEmpty arr2 = pure $ setComp comp empty + | otherwise = + pure $ makeArray comp (Sz2 m1 m2) $ \(i :. j) -> + UMA.unsafeIndex arr1 i * UMA.unsafeIndex arr2 j + where + comp = getComp arr1 <> getComp arr2 + Sz1 m1 = size arr1 + Sz1 m2 = size arr2 +{-# INLINE outer #-} + +det :: SRMatrix -> Double +det mtx + | m==0 || n==0 = 1 + | otherwise = (^2) $ Prelude.product [l ! (i :. i) | i <- [0 .. m-1]] + where + Sz (m :. n) = size mtx + (l, _) = unsafePerformIO (lu mtx) + +detChol :: SRMatrix -> Double +detChol mtx + | m==0 || n==0 = 1 + | otherwise = (^2) $ Prelude.product [cho ! (i :. i) | i <- [0 .. m-1]] + where + Sz (m :. n) = size mtx + cho = unsafePerformIO (cholesky mtx) +{-# INLINE det #-} + +rangedLinearDotProd :: PrimMonad m => Int -> Int -> Int -> MMassArray m -> m Double +rangedLinearDotProd r1 r2 len arr = go 0 0 + where + go !acc k + | k < len = do x <- UMA.unsafeLinearRead arr (r1 + k) + y <- UMA.unsafeLinearRead arr (r2 + k) + go (acc + x*y) (k + 1) + | otherwise = pure acc +{-# INLINE rangedLinearDotProd #-} + +data NegDef = NegDef + deriving Show + +instance Exception NegDef + +cholesky :: (PrimMonad m, MonadThrow m, MonadIO m) + => SRMatrix + -> m SRMatrix +cholesky arr + | m /= n = throwM $ SizeMismatchException (size arr) (size arr) + | isEmpty arr = pure $ setComp comp empty + | otherwise = MMA.createArrayS_ (size arr) create + where + comp = getComp arr + (Sz2 m n) = size arr + create l = Prelude.mapM_ (update l) [i :. j | i <- [0..m-1], j <- [0..m-1]] + + update l ix@(i :. j) + | i < j = UMA.unsafeWrite l ix 0 + | otherwise = do let cur = UMA.unsafeIndex arr ix + rowI = i*m + rowJ = j*m + xjj <- UMA.unsafeLinearRead l (rowJ + j) + tot <- rangedLinearDotProd rowI rowJ j l + let delta = cur - tot + if i == j + then if delta <= 0 + then throwM NegDef -- SizeMismatchException (size arr) (size arr) -- look at a better exception + else UMA.unsafeLinearWrite l (rowI + j) (sqrt delta) + else UMA.unsafeLinearWrite l (rowI + j) (delta / xjj) +{-# INLINE cholesky #-} + +invChol :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> m SRMatrix +invChol arr = do l <- cholesky arr -- lower diag + mtx <- thawS l + forM_ [0 .. m-1] $ \i -> do + lII <- UMA.unsafeRead mtx (i :. i) + UMA.unsafeWrite mtx (i :. i) (1 / lII) + forM_ [0 .. i-1] $ \j -> do + tot <- rangedLinearDotProd (i*m + j) (j*m + j) (i-j) mtx + UMA.unsafeWrite mtx (j :. i) ((-tot)/lII) + UMA.unsafeWrite mtx (i :. j) 0 + mm <- newMArray (Sz2 m m) 0 + forM_ [0 .. m-1] $ \i -> do + dii <- rangedLinearDotProd (i*m + i) (i*m + i) (m - i) mtx + UMA.unsafeWrite mm (i :. i) dii + forM_ [i+1 .. m-1] $ \j -> do + dij <- rangedLinearDotProd (i*m + j) (j*m + j) (m - j) mtx + UMA.unsafeWrite mm (i :. j) dij + UMA.unsafeWrite mm (j :. i) dij + freezeS mm + + where + Sz2 m _ = size arr +{-# INLINE invChol #-} + +-- LU decomposition and solver taken from https://hackage.haskell.org/package/linear-1.23/docs/src/Linear.Matrix.html +lu :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> m (SRMatrix, SRMatrix) +lu mtx = do + let (Sz2 m n) = size mtx + u <- thawS $ computeAs S $ identityMatrix (Sz m) + l <- thawS $ A.replicate compMode (Sz2 m n) 0 + + let buildLVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do lik <- UMA.unsafeRead l (i :. k) + ukj <- UMA.unsafeRead u (k :. j) + go (k+1) ( s + (lik * ukj) ) + s' <- go 0 0 + UMA.unsafeWrite l (i :. j) ((mtx ! (i :. j)) - s') + -- pure l + buildL !i !j + = when (i /= n) $ do buildLVal i j + buildL (i+1) j + buildUVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do ljk <- UMA.unsafeRead l (j :. k) + uki <- UMA.unsafeRead u (k :. i) + go (k+1) (s + ljk * uki) + + s' <- go 0 0 + ljj <- UMA.unsafeRead l (j :. j) + UMA.unsafeWrite u (j :. i) (((mtx ! (j :. i)) - s') / (ljj)) + -- pure u + + buildU !i !j + = when (i /= n) $ do buildUVal i j + buildU (i+1) j + buildLU !j + = when (j /= n) $ + do buildL j j + buildU j j + buildLU (j+1) + buildLU 0 + finalL <- freezeS l + finalU <- freezeS u + pure (finalL, finalU) + +forwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector +forwardSub a b = do + let (Sz m) = size b + x <- thawS $ A.replicate compMode (Sz1 m) 0 + let coeff !i !j !s + | j == i = pure s + | otherwise = do let aij = a ! (i :. j) + xj <- UMA.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i/= m) $ + do let bi = b ! i + aii = a ! (i :. i) + c <- coeff i 0 0 + UMA.unsafeWrite x i ((bi - c)/aii) + go (i+1) + go 0 + freezeS x + +backwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector +backwardSub a b = do + let (Sz m) = size b + x <- thawS $ A.replicate compMode (Sz1 m) 0 + let coeff !i !j !s + | j == m = pure s + | otherwise = do let aij = a ! (i :. j) + xj <- UMA.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i >= 0) $ + do let bi = b ! i + aii = a ! (i :. i) + c <- coeff i (i+1) 0 + UMA.unsafeWrite x i ((bi - c)/aii) + go (i-1) + go (m-1) + freezeS x + +luSolve :: (PrimMonad m, MonadThrow m, MonadIO m) => SRMatrix -> PVector -> m PVector +luSolve a b = do (l, u) <- lu a + forwardSub l b >>= backwardSub u + +type PolyCos = (Double, Double, Double) + +-- | Given a list of (x,y) co-ordinates, produces a list of coefficients to cubic equations, with knots at each of the initially provided x co-ordinates. Natural cubic spline interpololation is used. See: . +cubicSplineCoefficients :: [(Double, Double)] -> [PolyCos] +cubicSplineCoefficients xs = Prelude.zip3 x y z' + where + x = map fst xs + y = map snd xs + xdiff = zipWith (-) (tail x) x + xdiff' = fromList compMode xdiff :: Vector S Double + dydx :: Vector S Double + dydx = fromList compMode $ Prelude.zipWith3 (\y0 y1 xd -> (y0-y1)/xd) (tail y) y xdiff + n = length x + + w :: [Double] + w = 0 : nextW 1 w + where + nextW ix (wi : t) + | ix == n-1 = [] + | otherwise = let m = (xdiff' ! (ix-1)) * (2 - wi) + 2 * (xdiff' ! ix) + wn = (xdiff' ! ix) / m + in wn : nextW (ix+1) t + z :: [Double] + z = 0 : nextZ 1 z + where + nextZ ix (zi : t) + | ix == n-1 = [0] + | otherwise = let m = (xdiff' ! (ix-1)) * (2 - (w !! (ix-1))) + 2 * (xdiff' ! ix) + zn = (6*((dydx ! ix) - (dydx ! (ix-1))) - (xdiff' ! (ix-1)) * zi) / m + in zn : nextZ (ix+1) t + + z' :: [Double] + z' = Prelude.reverse $ 0 : [z !! i - w !! i * z !! (i+1) | i <- [n-2,n-3 .. 0]] + +chunkBy :: Int -> [t] -> [[t]] +chunkBy n = unfoldr go + where go [] = Nothing + go x = Just $ splitAt n x + +genSplineFun :: [(Double, Double)] -> Double -> Double +genSplineFun pts x = go xs $ zip coefs (tail coefs) + where + xs = map fst pts + coefs = cubicSplineCoefficients pts + evalAt (a1,b1,c1) (a2,b2,c2) y = let hi1 = a2 - a1 + in c1/(6*hi1)*(a2-y)^3 + c2/(6*hi1)*(y-a1)^3 + (b2/hi1 - c2*hi1/6)*(y-a1) + (b1/hi1 - c1*hi1/6)*(a2-y) + + go [x1,x2] [(c1,c2)] = evalAt c1 c2 x + go (x1:x2:xs) ((c1,c2):cs) + | x < x1 = evalAt c1 c2 x + | x >= x1 && x <= x2 = evalAt c1 c2 x + | otherwise = go (x2:xs) cs diff --git a/src/Algorithm/SRTree/AD.hs b/src/Algorithm/SRTree/AD.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD.hs @@ -0,0 +1,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 evalKey n + Const v -> evalKey n + Param ix -> evalKey n + _ -> getFromCache rt + getFromCache rt = do + global <- gets ((IntMap.!? rt) . fst . fst) + local <- gets ((IntMap.!? rt) . snd . fst) + if | isJust global -> pure (fromJust global, False) + | isJust local -> pure (fromJust local, True) + | otherwise -> insertKey rt + +-- reverse mode applied directly on an e-graph. Supports caching. +-- assumes root points to the loss function, so for an expression +-- f(x) and the loss (y - (f(x))^2), root will point to "^" +reverseModeEGraph :: SRMatrix -> PVector -> Maybe PVector -> EGraph -> ECache -> EClassId -> VS.Vector Double -> (Array D Ix1 Double, VS.Vector Double) +reverseModeEGraph xss ys mYErr egraph cache root' theta = + (delay $ rootVal + , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]] + ) + where + rootVal = extractCache (cache'' IntMap.!? root', localcache' IntMap.!? root') + root = canon root' + yErr = fromJust mYErr + m = M.size ys + p = VS.length theta + comp = M.getComp xss + one :: Array S Ix1 Double + one = M.replicate comp m 1 + + canon rt = case _canonicalMap egraph IntMap.!? rt of + Nothing -> error "wrong canon" + Just rt' -> if rt == rt' then rt else canon rt' + + getNode rt' = let rt = canon rt' + cls = _eClass egraph IntMap.! rt + in (_best . _info) cls + + getId n' = let n = runIdentity $ canonize n' `evalStateT` egraph + in if n `Map.member` _eNodeToEClass egraph then _eNodeToEClass egraph Map.! n else _eNodeToEClass egraph Map.! n' + + ((cache', localcache), _) = evalCached root `execState` ((cache, IntMap.empty), Map.empty) + where + evalCached :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalCached rt = insertKey rt + + insertKey :: EClassId -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + insertKey key' = do + let key = canon key' + isCachedGlobal <- gets ((key `IntMap.member`) . fst . fst) + isCachedLocal <- gets ((key `IntMap.member`) . snd . fst) + when (not isCachedLocal && not isCachedGlobal) $ do + let node = getNode key + (ev, toLocal) <- evalKey node + modify' (insKey node ev toLocal) + getVal key + + evalKey :: ENode -> State ((ECache, ECache), Map.Map ENode PVector) (PVector, Bool) + evalKey (Var ix) = pure $ if | ix == -1 -> (ys, False) + | ix == -2 -> (yErr, False) + | otherwise -> (M.computeAs S $ xss evalKey n + Const v -> evalKey n + Param ix -> evalKey n + _ -> getFromCache rt + getFromCache rt = do + global <- gets ((IntMap.!? rt) . fst . fst) + local <- gets ((IntMap.!? rt) . snd . fst) + if | isJust global -> pure (fromJust global, False) + | isJust local -> pure (fromJust local, True) + | otherwise -> insertKey rt + + extractCache (Nothing, Nothing) = error "no root info" + extractCache (Just r, _) = r + extractCache (_, Just r) = r + + ((cache'', localcache'), cachedGrad) = calcGrad root one `execState` ((cache', localcache), Map.empty) + + calcGrad :: Int -> Array S Ix1 Double -> State ((IntMap.IntMap (Array S Ix1 Double), IntMap.IntMap (Array S Ix1 Double)), Map.Map (SRTree Int) (Array S Ix1 Double)) () + calcGrad rt v = do let node = getNode rt + case node of + Bin op l r -> do xl <- fst <$> getVal l + xr <- fst <$> getVal r + (dl, dr) <- diff op v xl xr l r + calcGrad l dl + calcGrad r dr + Uni f t -> do x <- fst <$> getVal t + calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x)) + Param ix -> modify' (insertGrad v (Param ix)) + _ -> pure () + where + insertGrad v k ((a, b), g) = ((a, b), Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g) + + --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double) + diff Add dx fx gy l r = pure (dx, dx) + diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx) + diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx) + diff Div dx fx gy l r = do + let k = getId (Bin Div l r) + v <- fst <$> getVal k + pure (M.computeAs S $ M.zipWith (/) dx gy + , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy)) + diff Power dx fx gy l r = do + let k = getId (Bin Power l r) + v <- fst <$> getVal k + pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v) + + diff PowerAbs dx fx gy l r = do + let k = getId (Bin PowerAbs l r) + v <- fst <$> getVal k + let v2 = M.map abs fx + v3 = M.computeAs S $ M.zipWith (*) fx gy + pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2 + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v) + + diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx + dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl + in pure (M.computeAs S $ dxl, M.computeAs S $ dxy) + + fixNaN x = if isNaN x then 0 else x + + +reverseModeGraph :: SRMatrix -> PVector -> Maybe PVector -> VS.Vector Double -> Fix SRTree -> (Array D Ix1 Double, VS.Vector Double) +reverseModeGraph xss ys mYErr theta tree = (delay $ cachedVal' IntMap.! root + , VS.fromList [M.sum $ cachedGrad Map.! (Param ix) | ix <- [0..p-1]]) + where + yErr = fromJust mYErr + --ys = delay ys' + m = M.size ys + p = VS.length theta + comp = M.getComp xss + one :: Array S Ix1 Double + one = M.replicate comp m 1 + (key2int, int2key, cachedVal, (subtract 1) -> root) = cataM leftToRight alg tree `execState` (Map.empty, IntMap.empty, IntMap.empty, 0) + (key2int', int2key', cachedVal', cachedGrad) = calcGrad root one `execState` (key2int, int2key, cachedVal, Map.empty) + + calcGrad :: Int -> Array S Ix1 Double -> State (Map.Map (SRTree Int) Int, IntMap.IntMap (SRTree Int), IntMap.IntMap (Array S Ix1 Double), Map.Map (SRTree Int) (Array S Ix1 Double)) () + calcGrad key v = do node <- gets ((IntMap.! key) . _int2key) + case node of + Bin op l r -> do xl <- gets (getVal l) + xr <- gets (getVal r) + (dl, dr) <- diff op v xl xr l r + calcGrad l dl + calcGrad r dr + Uni f t -> do x <- gets (getVal t) + calcGrad t (M.computeAs S $ M.zipWith (*) v (M.map (derivative f) x)) + Param ix -> modify' (insertGrad v (Param ix)) + _ -> pure () + where + _int2key (_, b, _, _) = b + insertGrad v k (a, b, c, g) = (a, b, c, Map.insertWith (\v1 v2 -> M.computeAs S $ M.zipWith (+) v1 v2) k v g) + + graph (a, _, _, _) = a + insKey key ev (a, b, c, d) = (Map.insert key d a, IntMap.insert d key b, IntMap.insert d ev c, d+1) + -- get the values from the cache + getVal key (a, b, c, d) = c IntMap.! key + -- maps the the struct to an integer key + getKey key (a, b, c, d) = a Map.! key + + -- this tells the order in which we traverse the tree + leftToRight (Uni f mt) = Uni f <$> mt; + leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr + leftToRight (Var ix) = pure (Var ix) + leftToRight (Param ix) = pure (Param ix) + leftToRight (Const c) = pure (Const c) + + evalKey (Var ix) = pure $ if ix == -1 + then ys + else if ix == -2 + then yErr + else M.computeAs S $ xss gets (getVal t) + evalKey (Bin op l r) = M.computeAs S <$> (M.zipWith (evalOp op) <$> gets (getVal l) <*> gets (getVal r)) + + alg (Var ix) = insertKey (Var ix) + alg (Param ix) = insertKey (Param ix) + alg (Const v) = insertKey (Const v) + alg (Uni f t) = insertKey (Uni f t) + alg (Bin op l r) = insertKey (Bin op l r) + + --diff :: Op -> Array S Ix1 Double -> Array S Ix1 Double -> Array S Ix1 Double -> (Array S Ix1 Double, Array S Ix1 Double) + diff Add dx fx gy l r = pure (dx, dx) + diff Sub dx fx gy l r = pure (dx, M.computeAs S $ M.map negate dx) + diff Mul dx fx gy l r = pure (M.computeAs S $ M.zipWith (*) dx gy, M.computeAs S $ M.zipWith (*) dx fx) + diff Div dx fx gy l r = do + k <- gets (getKey (Bin Div l r)) + v <- gets (getVal k) + pure (M.computeAs S $ M.zipWith (/) dx gy + , M.computeAs S $ M.zipWith (*) dx (M.zipWith (\l r -> negate l/r) v gy)) + diff Power dx fx gy l r = do + k <- gets (getKey (Bin Power l r)) + v <- gets (getVal k) + pure ( M.computeAs S $ M.zipWith4 (\d f g vi -> fixNaN $ d * g * vi / f) dx fx gy v + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx fx v) + + diff PowerAbs dx fx gy l r = do + k <- gets (getKey (Bin PowerAbs l r)) + v <- gets (getVal k) + let v2 = M.map abs fx + v3 = M.computeAs S $ M.zipWith (*) fx gy + pure ( M.computeAs S $ M.zipWith4 (\d v3i vi v2i -> fixNaN $ d * v3i * vi / (v2i^2)) dx v3 v v2 + , M.computeAs S $ M.zipWith3 (\d f vi -> fixNaN $ d * vi * log f) dx v2 v) + + diff AQ dx fx gy l r = let dxl = M.zipWith (\g d -> d * (recip . sqrt . (+1) . (^2)) g) gy dx + dxy = M.zipWith3 (\f g dl -> f * g * dl^3) fx gy dxl + in pure (M.computeAs S $ dxl, M.computeAs S $ dxy) + + fixNaN x = if isNaN x then 0 else x + + insertKey key = do + isCached <- gets ((key `Map.member`) . graph) + when (not isCached) $ do + ev <- evalKey key + modify' (insKey key ev) + gets (getKey key) + +-- | Same as above, but using reverse mode with the tree encoded as an array, that is even faster. +reverseModeArr :: SRMatrix + -> PVector + -> Maybe PVector + -> VS.Vector Double -- PVector + -> [(Int, (Int, Int, Int, Double))] -- arity, opcode, ix, const val + -> IntMap.IntMap Int + -> (Array D Ix1 Double, Array S Ix1 Double) +reverseModeArr xss ys mYErr theta t j2ix = + unsafePerformIO $ do + fwd <- M.newMArray (Sz2 n m) 0 + partial <- M.newMArray (Sz2 n m) 0 + jacob <- M.newMArray (Sz p) 0 + val <- M.newMArray (Sz m) 0 + let + stps = 2 + --delta = m `div` stps + --rngs = [(i*delta, min m $ (i+1)*delta) | i <- [0..stps] ] + (a, b) = (0, m) + + forward (a, b) fwd + calculateYHat (a, b) fwd val + reverseMode (a, b) fwd partial + combine (a, b) partial jacob + j <- UMA.unsafeFreeze (getComp xss) jacob + v <- UMA.unsafeFreeze (getComp xss) val + pure (delay v, j) + + where + (Sz2 m _) = M.size xss + p = VS.length theta + n = length t + toLin i j = i*m + j + yErr = fromJust mYErr + eps = 1e-8 + + myForM_ [] _ = pure () + myForM_ (!x:xs) f = do f x + myForM_ xs f + {-# INLINE myForM_ #-} + + calculateYHat :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO () + calculateYHat (a, b) fwd yhat = myForM_ [a..b-1] $ \i -> do + vi <- UMA.unsafeRead fwd (0 :. i) + UMA.unsafeWrite yhat i vi + {-# INLINE calculateYHat #-} + + forward :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> IO () + forward (a, b) fwd = do + let t' = Prelude.reverse t + myForM_ t' makeFwd + where + makeFwd (j, (0, 0, ix, _)) = + do let j' = j2ix IntMap.! j + myForM_ [a..b-1] $ \i -> do + --let val = xss M.! (i :. ix) + UMA.unsafeWrite fwd (j' :. i) $ case ix of + (-1) -> ys M.! i + (-2) -> yErr M.! i + _ -> xss M.! (i :. ix) + makeFwd (j, (0, 1, ix, _)) = do let j' = j2ix IntMap.! j + v = theta VS.! ix + myForM_ [a..b-1] $ \i -> do + UMA.unsafeWrite fwd (j' :. i) v + makeFwd (j, (0, 2, _, x)) = do let j' = j2ix IntMap.! j + myForM_ [a..b-1] $ \i -> do + UMA.unsafeWrite fwd (j' :. i) x + makeFwd (j, (1, f, _, _)) = do let j' = j2ix IntMap.! j + j2 = j2ix IntMap.! (2*j + 1) + myForM_ [a..b-1] $ \i -> do + v <- UMA.unsafeRead fwd (j2 :. i) + UMA.unsafeWrite fwd (j' :. i) (evalFun (toEnum f) v) + makeFwd (j, (2, op, _, _)) = do let j' = j2ix IntMap.! j + j2 = j2ix IntMap.! (2*j + 1) + j3 = j2ix IntMap.! (2*j + 2) + myForM_ [a..b-1] $ \i -> do + l <- UMA.unsafeRead fwd (j2 :. i) + r <- UMA.unsafeRead fwd (j3 :. i) + UMA.unsafeWrite fwd (j' :. i) (evalOp (toEnum op) l r) + makeFwd _ = pure () + {-# INLINE makeFwd #-} + {-# INLINE forward #-} + + reverseMode :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix2 Double -> IO () + reverseMode (a, b) fwd partial = + do myForM_ [a..b-1] $ \i -> UMA.unsafeWrite partial (0 :. i) 1 + myForM_ t makeRev + where + makeRev (j, (1, f, _, _)) = do let dxj = j2ix IntMap.! j + vj = j2ix IntMap.! (2*j + 1) + myForM_ [a..b-1] $ \i -> do + v <- UMA.unsafeRead fwd (vj :. i) + dx <- UMA.unsafeRead partial (dxj :. i) + --let val = dx * derivative (toEnum f) v + UMA.unsafeWrite partial (vj :. i) (dx * derivative (toEnum f) v) + makeRev (j, (2, op, _, _)) = do let dxj = j2ix IntMap.! j + lj = j2ix IntMap.! (2*j + 1) + rj = j2ix IntMap.! (2*j + 2) + myForM_ [a..b-1] $ \i -> do + l <- UMA.unsafeRead fwd (lj :. i) + r <- UMA.unsafeRead fwd (rj :. i) + dx <- UMA.unsafeRead partial (dxj :. i) + let (dxl, dxr) = diff (toEnum op) dx l r + UMA.unsafeWrite partial (lj :. i) dxl + UMA.unsafeWrite partial (rj :. i) dxr + makeRev _ = pure () + {-# INLINE makeRev #-} + {-# INLINE reverseMode #-} + + --f(x)^g(x) + --d f(x)^g(x) / d f(x) = f(x)^(g(x)-1) + -- f(x) + g(x) = 1, 1 + -- f(x) - g(x) = 1, -1 + -- f(x) * g(x) = g(x), f(x) + -- f(x) / g(x) = 1/g(x), -f(x)/g(x)^2 + -- f(x) ^ g(x) = g(x) * f(x) ^ (g(x) - 1), f(x) ^ g(x) * log f(x) + -- |f(x)| ^ g(x) = g(x) * |f(x)| ^ (g(x) - 2) * f(x), |f(x)| ^ g(x) * log |f(x)| + + -- |f(x)| ^ g(x) = exp (log |f(x)| * g(x)) + -- => |f(x)| ^ (g(x) - 1) * g(x) + -- => |f(x)| ^ g(x) * log |f(x)| * 1 + + fixNaN x | isNaN x = 0 + | otherwise = x + + diff :: Op -> Double -> Double -> Double -> (Double, Double) + diff Add dx fx gy = (dx, dx) + diff Sub dx fx gy = (dx, negate dx) + diff Mul dx fx gy = (dx * gy, dx * fx) + diff Div dx fx gy = (dx / gy, dx * (negate fx / (gy * gy))) + --diff Power dx fx gy = (fixNaN $ dx * ((fx+eps)**gy - fx**gy)/eps, fixNaN $ dx * (fx**(gy+eps) - fx**gy)/eps) + --diff PowerAbs dx fx gy = (fixNaN $ dx * (abs (fx+eps)**gy - abs fx**gy)/eps, fixNaN $ dx * (abs fx**(gy+eps) - abs fx**gy)/eps) + {--} + diff Power 0 _ _ = (0, 0) + diff Power dx 0 0 = (0, 0) + diff Power dx fx 0 = (0, fixNaN $ dx * log fx) + diff Power dx 0 gy = (fixNaN $ dx * gy * if gy < 1 then eps ** (gy - 1) else 0 + , 0) --dx * fx ** gy * log fx) + diff Power dx fx gy = (fixNaN $ dx * gy * fx ** (gy - 1), fixNaN $ dx * fx ** gy * log fx) + + diff PowerAbs 0 fx gy = (0, 0) + diff PowerAbs 0 0 0 = (0, 0) + diff PowerAbs dx fx 0 = (0, fixNaN $ dx * log (abs fx)) + diff PowerAbs dx 0 gy = (0, fixNaN $ dx * if gy < 0 then eps ** gy else 0) + diff PowerAbs dx fx gy = (fixNaN $ dx * gy * fx * abs fx ** (gy - 2), fixNaN $ dx * abs fx ** gy * log (abs fx)) + {--} + diff AQ dx fx gy = let dxl = recip ((sqrt . (+1)) (gy * gy)) + dxy = fx * gy * (dxl^3) -- / (sqrt (gy*gy + 1)) + in (dxl * dx, dxy * dx) + + {-# INLINE diff #-} + + combine :: (Int, Int) -> MArray (PrimState IO) S Ix2 Double -> MArray (PrimState IO) S Ix1 Double -> IO () + combine (lo, hi) partial jacob = myForM_ t makeJacob + where + makeJacob (j, (0, 1, ix, _)) = do val <- UMA.unsafeRead jacob ix + let j' = j2ix IntMap.! j + addI a b acc = do v2 <- UMA.unsafeRead partial (b :. a) + pure (v2 + acc) + acc <- foldM (\a i -> addI i j' a) val [lo..hi-1] + UMA.unsafeWrite jacob ix acc + makeJacob _ = pure () + {-# INLINE combine #-} + +-- | The function `forwardModeUnique` calculates the numerical gradient of the tree and evaluates the tree at the same time. It assumes that each parameter has a unique occurrence in the expression. This should be significantly faster than `forwardMode`. +forwardModeUniqueJac :: SRMatrix -> PVector -> Fix SRTree -> [PVector] +forwardModeUniqueJac xss theta = snd . second (map (M.computeAs M.S) . DL.toList) . cata alg + where + (Sz n) = M.size theta + one = replicateAs xss 1 + + alg (Var ix) = (xss Double + , _theta2tau :: Double -> Double + } + +-- shows the CI with n places +showCI :: Int -> CI -> String +showCI n (CI x l h) = show (rnd l) <> " <= " <> show (rnd x) <> " <= " <> show (rnd h) + where + rnd = (/10^n) . (fromIntegral . round) . (*10^n) +printCI :: Int -> CI -> IO () +printCI n = putStrLn . showCI n + +-- | Calculates the confidence interval of the parameters using +-- Laplace approximation or Profile likelihood +paramCI :: CIType -> Int -> PVector -> Double -> [CI] +paramCI (Laplace stats) nSamples theta alpha = zipWith3 CI (A.toList theta) lows highs + where + -- the Laplace approximation is theta +/- t(1-alpha/2) * standard error + (A.Sz k) = A.size theta + t = quantile (studentT . fromIntegral $ nSamples - k) (1 - alpha / 2.0) + stdErr = _stdErr stats + lows = A.toList $ A.zipWith (-) theta $ A.map (*t) stdErr + highs = A.toList $ A.zipWith (+) theta $ A.map (*t) stdErr + +paramCI (Profile stats profiles) nSamples _ alpha = zipWith3 CI theta lows highs + where + -- for the profile likelihood we use the square root of the F-distribution with (1-alpha) + k = length theta + t = sqrt $ quantile (fDistribution k (fromIntegral $ nSamples - k)) (1 - alpha) + stdErr = _stdErr stats + lows = map (`_tau2theta` (-t)) profiles + highs = map (`_tau2theta` t) profiles + theta = map _opt profiles + +-- | calculates the prediction confidence interval using Laplace approximation or profile likelihood. +-- +predictionCI :: CIType -> Distribution -> (SRMatrix -> PVector) -> (SRMatrix -> [PVector]) -> (CI -> PVector -> Fix SRTree -> (Double -> Double, Double)) -> SRMatrix -> Fix SRTree -> PVector -> Double -> [CI] -> [CI] +predictionCI (Laplace stats) _ predFun jacFun _ xss tree theta alpha _ = zipWith3 CI yhat lows highs + where + yhat = A.toList $ predFun xss + jac' :: A.Matrix A.S Double + jac' = A.fromLists' compMode $ map A.toList $ jacFun xss + jac :: [PVector] + jac = A.toList $ A.outerSlices $ A.computeAs A.S $ A.transpose jac' + n = length yhat + (A.Sz k) = A.size theta + t = quantile (studentT . fromIntegral $ n - k) (1 - alpha / 2.0) + covs = A.toList $ A.outerSlices $ _cov stats + lows = zipWith (-) yhat $ map (*t) resStdErr + highs = zipWith (+) yhat $ map (*t) resStdErr + + getResStdError row = sqrt $ (A.!.!) row $ A.fromList compMode $ map (row A.!.!) covs + resStdErr = map getResStdError jac + +predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat 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

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 diff --git a/src/Algorithm/SRTree/Likelihoods.hs b/src/Algorithm/SRTree/Likelihoods.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Likelihoods.hs @@ -0,0 +1,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 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 #-} diff --git a/src/Algorithm/SRTree/ModelSelection.hs b/src/Algorithm/SRTree/ModelSelection.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/ModelSelection.hs @@ -0,0 +1,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 #-} diff --git a/src/Algorithm/SRTree/NonlinearOpt.hs b/src/Algorithm/SRTree/NonlinearOpt.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/NonlinearOpt.hs @@ -0,0 +1,976 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : Numeric.NLOPT +Copyright : (c) Matthew Peddie 2017 +License : BSD3 +Maintainer : Matthew Peddie +Stability : provisional +Portability : GHC + +This module provides a high-level, @hmatrix@-compatible interface to +the library by +Steven G. Johnson. + +NOTE: This is an adaptation from https://hackage.haskell.org/package/hmatrix-nlopt-0.2.0.0 +that removes the dependency to hmatrix and support any Vector Storage. + += Documentation + +Most non-numerical details are documented, but for specific +information on what the optimization methods do, how constraints are +handled, etc., you should consult: + + * The + + * The + + * The + += Example program + +The following interactive session example uses the Nelder-Mead simplex +algorithm, a derivative-free local optimizer, to minimize a trivial +function with a minimum of 22.0 at @(0, 0)@. + +>>> import Numeric.LinearAlgebra ( dot, fromList ) +>>> let objf x = x `dot` x + 22 -- define objective +>>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +>>> let algorithm = NELDERMEAD objf [] Nothing -- specify algorithm +>>> let problem = LocalProblem 2 stop algorithm -- specify problem +>>> let x0 = fromList [5, 10] -- specify initial guess +>>> minimizeLocal problem x0 +Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) + +-} + +module Algorithm.SRTree.NonlinearOpt ( + -- * Specifying the objective function + Objective + , ObjectiveD + , Preconditioner + -- * Specifying the constraints + -- ** Bound constraints + , Bounds(..) + -- ** Nonlinear constraints + -- + -- $nonlinearconstraints + + -- *** Constraint functions + , ScalarConstraint + , ScalarConstraintD + , VectorConstraint + , VectorConstraintD + -- *** Constraint types + , Constraint(..) + , EqualityConstraint(..) + , InequalityConstraint(..) + -- *** Collections of constraints + , EqualityConstraints + , EqualityConstraintsD + , InequalityConstraints + , InequalityConstraintsD + -- * Stopping conditions + -- + -- $nonempty + , StoppingCondition(..) + , NonEmpty(..) + -- * Additional configuration + , RandomSeed(..) + , Population(..) + , VectorStorage(..) + , InitialStep(..) + -- * Minimization problems + -- ** Local minimization + , LocalAlgorithm(..) + , LocalProblem(..) + , minimizeLocal + -- ** Global minimization + , GlobalAlgorithm(..) + , GlobalProblem(..) + , minimizeGlobal + -- ** Minimization by augmented Lagrangian + , AugLagAlgorithm(..) + , AugLagProblem(..) + , minimizeAugLag + -- ** Results + , Solution(..) + , N.Result(..) + ) where + +import qualified Numeric.Optimization.NLOPT.Bindings as N + +import Data.List.NonEmpty (NonEmpty(..)) + +import qualified Data.Vector.Storable as V +import Data.Vector.Storable ( Vector ) + +import Control.Exception ( Exception ) +import qualified Control.Exception as Ex +import Data.Typeable ( Typeable ) +import Data.Foldable ( traverse_ ) + +import System.IO.Unsafe ( unsafePerformIO ) + +-- each element i contains a row vec +type Matrix a = [Vector a] + +flatten :: V.Storable a => Matrix a -> Vector a +flatten = V.concat +{-# INLINE flatten #-} + +{- Function wrapping for the immutable HMatrix interface -} +wrapScalarFunction :: (Vector Double -> Double) -> N.ScalarFunction () +wrapScalarFunction f params _ _ = return $ f params + +wrapScalarFunctionD :: (Vector Double -> (Double, Vector Double)) + -> N.ScalarFunction () +wrapScalarFunctionD f params grad _ = do + case grad of + Nothing -> return () + Just g -> V.copy g usergrad + return result + where + (result, usergrad) = f params + +wrapVectorFunction :: (Vector Double -> Word -> Vector Double) + -> Word -> N.VectorFunction () +wrapVectorFunction f n params vout _ _ = V.copy vout $ f params n + +wrapVectorFunctionD :: (Vector Double -> Word -> (Vector Double, Matrix Double)) + -> Word -> N.VectorFunction () +wrapVectorFunctionD f n params vout jac _ = do + V.copy vout result + case jac of + Nothing -> return () + Just j -> V.copy j (flatten userjac) + where + (result, userjac) = f params n + +wrapPreconditionerFunction :: (Vector Double -> Vector Double -> Vector Double) + -> N.PreconditionerFunction () +wrapPreconditionerFunction f params v vpre _ = V.copy vpre (f params v) + +{- Objective functions -} +-- | An objective function that calculates the objective value at the +-- given parameter vector. +type Objective + = Vector Double -- ^ Parameter vector + -> Double -- ^ Objective function value + +-- | An objective function that calculates both the objective value +-- and the gradient of the objective with respect to the input +-- parameter vector, at the given parameter vector. +type ObjectiveD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Objective function value, gradient) + +-- | A preconditioner function, which computes @vpre = H(x) v@, where +-- @H@ is the Hessian matrix: the positive semi-definite second +-- derivative at the given parameter vector @x@, or an approximation +-- thereof. +type Preconditioner + = Vector Double -- ^ Parameter vector @x@ + -> Vector Double -- ^ Vector @v@ to precondition at @x@ + -> Vector Double -- ^ Preconditioned vector @vpre@ + +data ObjectiveFunction f + = MinimumObjective f + | PreconditionedMinimumObjective Preconditioner f + +applyObjective :: N.Opt -> ObjectiveFunction Objective -> IO N.Result +applyObjective opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunction f) () +applyObjective opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunction f) + (wrapPreconditionerFunction p) () + +applyObjectiveD :: N.Opt -> ObjectiveFunction ObjectiveD -> IO N.Result +applyObjectiveD opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunctionD f) () +applyObjectiveD opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunctionD f) + (wrapPreconditionerFunction p) () + +{- Constraint functions -} +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraint + = Vector Double -- ^ Parameter vector @x@ + -> Double -- ^ Constraint violation (deviation from 0) + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the gradient of @c(x)@ with respect to @x@ at +-- that point. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraintD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Constraint violation, constraint gradient) + +-- | A constraint function which returns a vector @c(x)@ given the +-- parameter vector @x@. The constraint will enforce that @c(x) == 0@ +-- (equality constraint) or @c(x) <= 0@ (inequality constraint). +type VectorConstraint + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> Vector Double -- ^ Constraint violation vector + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the Jacobian (first derivative) matrix of +-- @c(x)@ with respect to @x@ at that point. The constraint will +-- enforce that @c(x) == 0@ (equality constraint) or @c(x) <= 0@ +-- (inequality constraint). +type VectorConstraintD + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> (Vector Double, Matrix Double) -- ^ (Constraint violation vector, + -- constraint Jacobian) + +-- $nonlinearconstraints +-- +-- Note that most NLOPT algorithms do not support nonlinear +-- constraints natively; if you need to enforce nonlinear constraints, +-- you may want to use the 'AugLagAlgorithm' family of solvers, which +-- can add nonlinear constraints to some algorithm that does not +-- support them by a principled modification of the objective +-- function. +-- +-- == Example program +-- +-- The following interactive session example enforces a scalar +-- constraint on the problem given in the beginning of the module: the +-- parameters must always sum to 1. The minimizer finds a constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = COBYLA objf [] [] [constraint] Nothing +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.500000000013028, solutionParams = [0.5000025521533521,0.49999744784664796], solutionResult = FTOL_REACHED}) + + +data Constraint s v + -- | A scalar constraint. + = Scalar s + -- | A vector constraint. + | Vector Word v + -- | A scalar constraint with an attached preconditioning function. + | Preconditioned Preconditioner s + +-- | An equality constraint, comprised of both the constraint function +-- (or functions, if a preconditioner is used) along with the desired +-- tolerance. +data EqualityConstraint s v = EqualityConstraint + { eqConstraintFunctions :: Constraint s v + , eqConstraintTolerance :: Double + } + +-- | An inequality constraint, comprised of both the constraint +-- function (or functions, if a preconditioner is used) along with the +-- desired tolerance. +data InequalityConstraint s v = InequalityConstraint + { ineqConstraintFunctions :: Constraint s v + , ineqConstraintTolerance :: Double + } + +-- | A collection of equality constraints that do not supply +-- constraint derivatives. +type EqualityConstraints = + [EqualityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of inequality constraints that do not supply +-- constraint derivatives. +type InequalityConstraints = + [InequalityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of equality constraints that supply constraint +-- derivatives. +type EqualityConstraintsD = [EqualityConstraint ScalarConstraintD VectorConstraintD] + +-- | A collection of inequality constraints that supply constraint +-- derivatives. +type InequalityConstraintsD = [InequalityConstraint ScalarConstraintD VectorConstraintD] + +class ApplyConstraint constraint where + applyConstraint :: N.Opt -> constraint -> IO N.Result + +instance ApplyConstraint (EqualityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (EqualityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +{- Bounds -} + +-- | Bound constraints are specified by vectors of the same dimension +-- as the parameter space. +-- +-- == Example program +-- +-- The following interactive session example enforces lower bounds on +-- the example from the beginning of the module. This prevents the +-- optimizer from locating the true minimum at @(0, 0)@; a slightly +-- higher constrained minimum at @(1, 1)@ is found. Note that the +-- optimizer returns 'N.XTOL_REACHED' rather than 'N.FTOL_REACHED', +-- because the bound constraint is active at the final minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +-- >>> let lowerbound = LowerBounds $ fromList [1, 1] -- specify bounds +-- >>> let algorithm = NELDERMEAD objf [lowerbound] Nothing -- specify algorithm +-- >>> let problem = LocalProblem 2 stop algorithm -- specify problem +-- >>> let x0 = fromList [5, 10] -- specify initial guess +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 24.0, solutionParams = [1.0,1.0], solutionResult = XTOL_REACHED}) +data Bounds + -- | Lower bound vector @v@ means we want @x >= v@. + = LowerBounds (Vector Double) + -- | Upper bound vector @u@ means we want @x <= u@. + | UpperBounds (Vector Double) + deriving (Eq, Show, Read) + +applyBounds :: N.Opt -> Bounds -> IO N.Result +applyBounds opt (LowerBounds lbvec) = N.set_lower_bounds opt lbvec +applyBounds opt (UpperBounds ubvec) = N.set_upper_bounds opt ubvec + +{- Stopping conditions -} + +-- | A 'StoppingCondition' tells NLOPT when to stop working on a +-- minimization problem. When multiple 'StoppingCondition's are +-- provided, the problem will stop when any one condition is met. +data StoppingCondition + -- | Stop minimizing when an objective value @J@ less than or equal + -- to the provided value is found. + = MinimumValue Double + -- | Stop minimizing when an optimization step changes the objective + -- value @J@ by less than the provided tolerance multiplied by @|J|@. + | ObjectiveRelativeTolerance Double + -- | Stop minimizing when an optimization step changes the objective + -- value by less than the provided tolerance. + | ObjectiveAbsoluteTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than @x@ scaled by the provided + -- tolerance. + | ParameterRelativeTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than the corresponding element in + -- the provided vector of tolerances values. + | ParameterAbsoluteTolerance (Vector Double) + -- | Stop when the number of evaluations of the objective function + -- exceeds the provided count. + | MaximumEvaluations Word + -- | Stop when the optimization time exceeds the provided time (in + -- seconds). This is not a precise limit. + | MaximumTime Double + deriving (Eq, Show, Read) + +-- $nonempty +-- +-- The 'NonEmpty' data type from 'Data.List.NonEmpty' is re-exported +-- here, because it is used to ensure that you always specify at least +-- one stopping condition. + +applyStoppingCondition :: N.Opt -> StoppingCondition -> IO N.Result +applyStoppingCondition opt (MinimumValue x) = N.set_stopval opt x +applyStoppingCondition opt (ObjectiveRelativeTolerance x) = N.set_ftol_rel opt x +applyStoppingCondition opt (ObjectiveAbsoluteTolerance x) = N.set_ftol_abs opt x +applyStoppingCondition opt (ParameterRelativeTolerance x) = N.set_xtol_rel opt x +applyStoppingCondition opt (ParameterAbsoluteTolerance v) = N.set_xtol_abs opt v +applyStoppingCondition opt (MaximumEvaluations n) = N.set_maxeval opt n +applyStoppingCondition opt (MaximumTime deltat) = N.set_maxtime opt deltat + +{- Random seed control -} + +-- | This specifies how to initialize the random number generator for +-- stochastic algorithms. +data RandomSeed + -- | Seed the RNG with the provided value. + = SeedValue Word + -- | Seed the RNG using the system clock. + | SeedFromTime + -- | Don't perform any explicit initialization of the RNG. + | Don'tSeed + deriving (Eq, Show, Read) + +applyRandomSeed :: RandomSeed -> IO () +applyRandomSeed Don'tSeed = return () +applyRandomSeed (SeedValue n) = N.srand n +applyRandomSeed SeedFromTime = N.srand_time + +{- Random stuff -} + +-- | This specifies the population size for algorithms that use a pool +-- of solutions. +newtype Population = Population Word deriving (Eq, Show, Read) + +applyPopulation :: N.Opt -> Population -> IO N.Result +applyPopulation opt (Population n) = N.set_population opt n + +-- | This specifies the memory size to be used by algorithms like +-- 'LBFGS' which store approximate Hessian or Jacobian matrices. +newtype VectorStorage = VectorStorage Word deriving (Eq, Show, Read) + +applyVectorStorage :: N.Opt -> VectorStorage -> IO N.Result +applyVectorStorage opt (VectorStorage n) = N.set_vector_storage opt n + +-- | This vector with the same dimension as the parameter vector @x@ +-- specifies the initial step for the optimizer to take. (This +-- applies to local gradient-free algorithms, which cannot use +-- gradients to estimate how big a step to take.) +newtype InitialStep = InitialStep (Vector Double) deriving (Eq, Show, Read) + +applyInitialStep :: N.Opt -> InitialStep -> IO N.Result +applyInitialStep opt (InitialStep v) = N.set_initial_step opt v + +{- Algorithms -} + +data GlobalProblem = GlobalProblem + { lowerBounds :: Vector Double -- ^ Lower bounds for @x@ + , upperBounds :: Vector Double -- ^ Upper bounds for @x@ + , gstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , galgorithm :: GlobalAlgorithm -- ^ Algorithm specification + } + +-- | These are the global minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'Population', you can simply specify 'Nothing' to use +-- the default behavior. +data GlobalAlgorithm + -- | DIviding RECTangles + = DIRECT Objective + -- | DIviding RECTangles, locally-biased variant + | DIRECT_L Objective + -- | DIviding RECTangles, "slightly randomized" + | DIRECT_L_RAND Objective RandomSeed + -- | DIviding RECTangles, unscaled version + | DIRECT_NOSCAL Objective + -- | DIviding RECTangles, locally-biased and unscaled + | DIRECT_L_NOSCAL Objective + -- | DIviding RECTangles, locally-biased, unscaled and "slightly + -- randomized" + | DIRECT_L_RAND_NOSCAL Objective RandomSeed + -- | DIviding RECTangles, original FORTRAN implementation + | ORIG_DIRECT Objective InequalityConstraints + -- | DIviding RECTangles, locally-biased, original FORTRAN + -- implementation + | ORIG_DIRECT_L Objective InequalityConstraints + -- | Stochastic Global Optimization. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO ObjectiveD + -- | Stochastic Global Optimization, randomized variant. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO_RAND ObjectiveD RandomSeed + -- | Controlled Random Search with Local Mutation + | CRS2_LM Objective RandomSeed (Maybe Population) + -- | Improved Stochastic Ranking Evolution Strategy + | ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population) + -- | Evolutionary Algorithm + | ESCH Objective + -- | Original Multi-Level Single-Linkage + | MLSL Objective LocalProblem (Maybe Population) + -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy + -- Sequence for starting points + | MLSL_LDS Objective LocalProblem (Maybe Population) + +algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm +algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT +algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L +algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND +algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL +algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT +algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L +algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO +algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND +algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM +algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES +algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH +algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL +algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS + +applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + + go (DIRECT o) = obj o + go (DIRECT_L o) = obj o + go (DIRECT_NOSCAL o) = obj o + go (DIRECT_L_NOSCAL o) = obj o + go (ESCH o) = obj o + go (STOGO o) = objD o + go (DIRECT_L_RAND o _) = obj o + go (DIRECT_L_RAND_NOSCAL o _) = obj o + go (ORIG_DIRECT o _) = obj o + go (ORIG_DIRECT_L o _) = obj o + go (STOGO_RAND o _) = objD o + go (CRS2_LM o _ _) = obj o + go (ISRES o _ _ _ _) = obj o + go (MLSL o _ _) = obj o + go (MLSL_LDS o _ _) = obj o + +applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalAlgorithm opt alg = do + applyGlobalObjective opt alg + go alg + where + seed = applyRandomSeed + pop = maybe (return ()) (tryTo . applyPopulation opt) + ic = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + + local lp = setupLocalProblem lp >>= N.set_local_optimizer opt + + go (DIRECT_L_RAND _ s) = seed s + go (DIRECT_L_RAND_NOSCAL _ s) = seed s + go (ORIG_DIRECT _ ineq) = ic ineq + go (ORIG_DIRECT_L _ ineq) = ic ineq + go (STOGO_RAND _ s) = seed s + go (CRS2_LM _ s p) = seed s *> pop p + go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p + go (MLSL _ lp p) = local lp *> pop p + go (MLSL_LDS _ lp p) = local lp *> pop p + go _ = return () + +tryTo :: IO N.Result -> IO () +tryTo act = do + result <- act + if (N.isSuccess result) + then return () + else Ex.throw $ NloptException result + +data NloptException = NloptException N.Result deriving (Show, Typeable) +instance Exception NloptException + +-- | Solve the specified global optimization problem. +-- +-- = Example program +-- +-- The following interactive session example uses the 'ISRES' +-- algorithm, a stochastic, derivative-free global optimizer, to +-- minimize a trivial function with a minimum of 22.0 at @(0, 0)@. +-- The search is conducted within a box from -10 to 10 in each +-- dimension. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion +-- >>> let algorithm = ISRES objf [] [] (SeedValue 22) Nothing -- specify algorithm +-- >>> let lowerbounds = fromList [-10, -10] -- specify bounds +-- >>> let upperbounds = fromList [10, 10] -- specify bounds +-- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm +-- >>> let x0 = fromList [5, 8] -- specify initial guess +-- >>> minimizeGlobal problem x0 +-- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED}) +minimizeGlobal :: GlobalProblem -- ^ Problem specification + -> Vector Double -- ^ Initial parameter guess + -> Either N.Result Solution -- ^ Optimization results +minimizeGlobal prob x0 = + unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +applyGlobalProblem :: N.Opt -> GlobalProblem -> IO () +applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do + tryTo $ applyBounds opt (LowerBounds lb) + tryTo $ applyBounds opt (UpperBounds ub) + traverse_ (tryTo . applyStoppingCondition opt) stop + applyGlobalAlgorithm opt alg + +newOpt :: N.Algorithm -> Word -> IO N.Opt +newOpt alg sz = do + opt' <- N.create alg sz + case opt' of + Nothing -> Ex.throw $ NloptException N.FAILURE + Just opt -> return opt + +setupGlobalProblem :: GlobalProblem -> IO N.Opt +setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do + opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp) + applyGlobalProblem opt gp + return opt + +solveProblem :: N.Opt -> Vector Double -> IO Solution +solveProblem opt x0 = do + (N.Output outret outcost outx nevals) <- N.optimize opt x0 + if (N.isSuccess outret) + then return $ Solution outcost outx outret nevals + else Ex.throw $ NloptException outret + +minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution +minimizeGlobal' gp x0 = do + opt <- setupGlobalProblem gp + solveProblem opt x0 + +data LocalProblem = LocalProblem + { lsize :: Word -- ^ The dimension of the + -- parameter vector. + , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification + } + +-- | These are the local minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one +-- another. Note that some local methods require you provide +-- derivatives (gradients or Jacobians) for your objective function +-- and constraint functions. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to +-- use the default behavior. +data LocalAlgorithm + -- | Limited-memory BFGS + = LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage) + -- | Limited-memory BFGS + | LBFGS ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-2 + | VAR2 ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-1 + | VAR1 ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method + | TNEWTON ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method with automatic restarting + | TNEWTON_RESTART ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method + | TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method with automatic + -- restarting + | TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage) + -- | Method of moving averages + | MMA ObjectiveD InequalityConstraintsD + -- | Sequential Least-Squares Quadratic Programming + | SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD + -- | Conservative Convex Separable Approximation + | CCSAQ ObjectiveD Preconditioner + -- | PRincipal AXIS gradient-free local optimization + | PRAXIS Objective [Bounds] (Maybe InitialStep) + -- | Constrained Optimization BY Linear Approximations + | COBYLA Objective [Bounds] InequalityConstraints EqualityConstraints + (Maybe InitialStep) + -- | Powell's NEWUOA algorithm + | NEWUOA Objective (Maybe InitialStep) + -- | Powell's NEWUOA algorithm with bounds by SGJ + | NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep) + -- | Nelder-Mead Simplex gradient-free method + | NELDERMEAD Objective [Bounds] (Maybe InitialStep) + -- | NLOPT implementation of Rowan's Subplex algorithm + | SBPLX Objective [Bounds] (Maybe InitialStep) + -- | Bounded Optimization BY Quadratic Approximations + | BOBYQA Objective [Bounds] (Maybe InitialStep) + +algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm +algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL +algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS +algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2 +algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1 +algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON +algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART +algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND +algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART +algorithmEnumOfLocal (MMA _ _) = N.LD_MMA +algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP +algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ +algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS +algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA +algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD +algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX +algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA + +applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO () +applyLocalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p + + go (LBFGS_NOCEDAL o _) = objD o + go (LBFGS o _) = objD o + go (VAR2 o _) = objD o + go (VAR1 o _) = objD o + go (TNEWTON o _) = objD o + go (TNEWTON_RESTART o _) = objD o + go (TNEWTON_PRECOND o _) = objD o + go (TNEWTON_PRECOND_RESTART o _) = objD o + go (MMA o _) = objD o + go (SLSQP o _ _ _) = objD o + go (CCSAQ o prec) = precond prec o + go (PRAXIS o _ _) = obj o + go (COBYLA o _ _ _ _) = obj o + go (NEWUOA o _) = obj o + go (NEWUOA_BOUND o _ _) = obj o + go (NELDERMEAD o _ _) = obj o + go (SBPLX o _ _) = obj o + go (BOBYQA o _ _) = obj o + +applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO () +applyLocalAlgorithm opt alg = do + applyLocalObjective opt alg + go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + ecd = traverse_ (tryTo . applyConstraint opt) + store = maybe (return ()) (tryTo . applyVectorStorage opt) + bound = traverse_ (tryTo . applyBounds opt) + step0 = maybe (return ()) (tryTo . applyInitialStep opt) + + go (LBFGS_NOCEDAL _ vs) = store vs + go (LBFGS _ vs) = store vs + go (VAR2 _ vs) = store vs + go (VAR1 _ vs) = store vs + go (TNEWTON _ vs) = store vs + go (TNEWTON_RESTART _ vs) = store vs + go (TNEWTON_PRECOND _ vs) = store vs + go (TNEWTON_PRECOND_RESTART _ vs) = store vs + go (MMA _ ineqd) = icd ineqd + go (SLSQP _ b ineqd eqd) = + bound b *> icd ineqd *> ecd eqd + go (CCSAQ _ _ ) = return () + go (PRAXIS _ b s) = bound b *> step0 s + go (COBYLA _ b ineq eq s) = + bound b *> ic ineq *> ec eq *> step0 s + go (NEWUOA _ s) = step0 s + go (NEWUOA_BOUND _ b s) = bound b *> step0 s + go (NELDERMEAD _ b s) = bound b *> step0 s + go (SBPLX _ b s) = bound b *> step0 s + go (BOBYQA _ b s) = bound b *> step0 s + +applyLocalProblem :: N.Opt -> LocalProblem -> IO () +applyLocalProblem opt (LocalProblem _ stop alg) = do + traverse_ (tryTo . applyStoppingCondition opt) stop + applyLocalAlgorithm opt alg + +setupLocalProblem :: LocalProblem -> IO N.Opt +setupLocalProblem lp@(LocalProblem sz _ alg) = do + opt <- newOpt (algorithmEnumOfLocal alg) sz + applyLocalProblem opt lp + return opt + +minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution +minimizeLocal' lp x0 = do + opt <- setupLocalProblem lp + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the SLSQP solver to find the minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale ) +-- >>> let objf x = (x `dot` x + 22, 2 `scale` x) +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1]) +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = SLSQP objf [] [] [constraint] +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED}) +minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution +minimizeLocal prob x0 = + unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +class ProblemSize c where + problemSize :: c -> Word + +instance ProblemSize LocalProblem where + problemSize = lsize + +instance ProblemSize GlobalProblem where + problemSize = fromIntegral . V.length . lowerBounds + +instance ProblemSize AugLagProblem where + problemSize (AugLagProblem _ _ alg) = case alg of + AUGLAG_LOCAL lp _ _ -> problemSize lp + AUGLAG_EQ_LOCAL lp -> problemSize lp + AUGLAG_GLOBAL gp _ _ -> problemSize gp + AUGLAG_EQ_GLOBAL gp -> problemSize gp + + +-- | __IMPORTANT NOTE__ +-- +-- For augmented lagrangian problems, you, the user, are responsible +-- for providing the appropriate type of constraint. If the +-- subsidiary problem requires an `ObjectiveD`, then you should +-- provide constraint functions with derivatives. If the subsidiary +-- problem requires an `Objective`, you should provide constraint +-- functions without derivatives. If you don't do this, you may get a +-- runtime error. +data AugLagProblem = AugLagProblem + { alEquality :: EqualityConstraints -- ^ Possibly empty set of + -- equality constraints + , alEqualityD :: EqualityConstraintsD -- ^ Possibly empty set of + -- equality constraints with + -- derivatives + , alalgorithm :: AugLagAlgorithm -- ^ Algorithm specification. + } + +-- | The Augmented Lagrangian solvers allow you to enforce nonlinear +-- constraints while using local or global algorithms that don't +-- natively support them. The subsidiary problem is used to do the +-- minimization, but the @AUGLAG@ methods modify the objective to +-- enforce the constraints. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- See the documentation for 'AugLagProblem' for an important note +-- about the constraint functions. +data AugLagAlgorithm + -- | AUGmented LAGrangian with a local subsidiary method + = AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a local subsidiary method and with + -- penalty functions only for equality constraints + | AUGLAG_EQ_LOCAL LocalProblem + -- | AUGmented LAGrangian with a global subsidiary method + | AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a global subsidiary method and with + -- penalty functions only for equality constraints. + | AUGLAG_EQ_GLOBAL GlobalProblem + +algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm +algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ +algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ + +-- | This structure is returned in the event of a successful +-- optimization. +data Solution = Solution + { solutionCost :: Double -- ^ The objective function value + -- at the minimum + , solutionParams :: Vector Double -- ^ The parameter vector which + -- minimizes the objective + , solutionResult :: N.Result -- ^ Why the optimizer stopped + + , nEvals :: Int -- ^ Number of evaluations until stop + } deriving (Eq, Show, Read) + +applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO () +applyAugLagAlgorithm opt alg = go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + -- AUGLAG won't work at all if you don't pass it the same + -- objective as the subproblem -- here we pull out the subproblem + -- objectives from the algorithm spec and set the same objective + -- function so the user can't mess it up. + local lp = tryTo $ do + localopt <- setupLocalProblem lp + applyLocalObjective opt (lalgorithm lp) + N.set_local_optimizer opt localopt + global gp = do + tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt + applyGlobalObjective opt (galgorithm gp) + + go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_LOCAL lp) = local lp + go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_GLOBAL gp) = global gp + +applyAugLagProblem :: N.Opt -> AugLagProblem -> IO () +applyAugLagProblem opt (AugLagProblem eq eqd alg) = do + traverse_ (tryTo . applyConstraint opt) eq + traverse_ (tryTo . applyConstraint opt) eqd + applyAugLagAlgorithm opt alg + +minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution +minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do + opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap) + applyAugLagProblem opt ap + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the augmented Lagrangian method to enforce the constraint and +-- the 'SBPLX' algorithm, which does not support nonlinear constraints +-- itself, to perform the minimization. As before, the parameters +-- must always sum to 1, and the minimizer finds the same constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let algorithm = SBPLX objf [] Nothing +-- >>> let subproblem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal subproblem x0 +-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem) +-- >>> minimizeAugLag problem x0 +-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED}) + +minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution +minimizeAugLag prob x0 = + unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode diff --git a/src/Algorithm/SRTree/Opt.hs b/src/Algorithm/SRTree/Opt.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Opt.hs @@ -0,0 +1,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 diff --git a/src/Data/SRTree.hs b/src/Data/SRTree.hs --- a/src/Data/SRTree.hs +++ b/src/Data/SRTree.hs @@ -1,7 +1,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -16,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 (..) ) diff --git a/src/Data/SRTree/Datasets.hs b/src/Data/SRTree/Datasets.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Datasets.hs @@ -0,0 +1,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), + ( 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 [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 diff --git a/src/Data/SRTree/Derivative.hs b/src/Data/SRTree/Derivative.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Derivative.hs @@ -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 #-} diff --git a/src/Data/SRTree/Eval.hs b/src/Data/SRTree/Eval.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Eval.hs @@ -0,0 +1,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 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 #-} diff --git a/src/Data/SRTree/Internal.hs b/src/Data/SRTree/Internal.hs --- a/src/Data/SRTree/Internal.hs +++ b/src/Data/SRTree/Internal.hs @@ -1,11 +1,12 @@ {-# language FlexibleInstances, DeriveFunctor #-} {-# language ScopedTypeVariables #-} {-# language RankNTypes #-} -{-# language ViewPatterns #-} +{-# language OverloadedStrings #-} +{-# language LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Internal --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -21,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 diff --git a/src/Data/SRTree/Print.hs b/src/Data/SRTree/Print.hs --- a/src/Data/SRTree/Print.hs +++ b/src/Data/SRTree/Print.hs @@ -1,7 +1,9 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Print --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -12,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 diff --git a/src/Data/SRTree/Random.hs b/src/Data/SRTree/Random.hs --- a/src/Data/SRTree/Random.hs +++ b/src/Data/SRTree/Random.hs @@ -2,7 +2,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Random --- Copyright : (c) Fabricio Olivetti 2021 - 2021 +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental @@ -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 #-} diff --git a/src/Data/SRTree/Recursion.hs b/src/Data/SRTree/Recursion.hs --- a/src/Data/SRTree/Recursion.hs +++ b/src/Data/SRTree/Recursion.hs @@ -1,5 +1,17 @@ {-# language RankNTypes #-} {-# language DeriveFunctor #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.Recursion +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Recursion schemes +-- +----------------------------------------------------------------------------- module Data.SRTree.Recursion where import Control.Monad ( (>=>) ) diff --git a/src/Numeric/Optimization/NLOPT/Bindings.hs b/src/Numeric/Optimization/NLOPT/Bindings.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Optimization/NLOPT/Bindings.hs @@ -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 +Stability : provisional +Portability : GHC + +Low-level interface to the NLOPT library. Please see + +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) diff --git a/src/Text/ParseSR.hs b/src/Text/ParseSR.hs new file mode 100644 --- /dev/null +++ b/src/Text/ParseSR.hs @@ -0,0 +1,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 diff --git a/src/Text/ParseSR/IO.hs b/src/Text/ParseSR/IO.hs new file mode 100644 --- /dev/null +++ b/src/Text/ParseSR/IO.hs @@ -0,0 +1,73 @@ +{-# language LambdaCase #-} +----------------------------------------------------------------------------- +-- | +-- Module : Text.ParseSR.IO +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to parse multiple expressions from stdin or a text file. +-- +----------------------------------------------------------------------------- +module Text.ParseSR.IO ( withInput, withOutput, withOutputDebug ) + where + +-- import Data.SRTree.EqSat1 +import Algorithm.EqSat.Simplify ( simplifyEqSatDefault ) +import Control.Monad (forM_, unless) +import qualified Data.ByteString.Char8 as B +import Data.SRTree +import Data.SRTree.Recursion (Fix (..)) +import System.IO +import Text.ParseSR (Output, SRAlgs, parseSR, showOutput) + +-- | given a filename, the symbolic regression algorithm, a string of variables name, +-- and two booleans indicating whether to convert float values to parameters and +-- whether to simplify the expression or not, it will read the file and parse everything +-- returning a list of either an error message or a tree. +-- +-- empty filename defaults to stdin +withInput :: String -> SRAlgs -> String -> Bool -> Bool -> IO [Either String (Fix SRTree)] +withInput fname sr hd param simpl = do + h <- if null fname then pure stdin else openFile fname ReadMode + contents <- hGetLines h + let myParserFun = parseSR sr (B.pack hd) param . B.pack + -- myParser = if simpl then fmap simplifyEqSat . myParserFun else myParserFun + myParser = if simpl then fmap simplifyEqSatDefault . myParserFun else myParserFun + es = map myParser $ filter (not . null) contents + unless (null fname) $ hClose h + pure es + +-- | outputs a list of either error or trees to a file using the Output format. +-- +-- empty filename defaults to stdout +withOutput :: String -> Output -> [Either String (Fix SRTree)] -> IO () +withOutput fname output exprs = do + h <- if null fname then pure stdout else openFile fname WriteMode + forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err + Right ex -> hPutStrLn h (showOutput output ex) + unless (null fname) $ hClose h + +-- | debug version of output function to check the invalid parsers +withOutputDebug :: String -> Output -> [Either String (Fix SRTree, Fix SRTree)] -> IO () +withOutputDebug fname output exprs = do + h <- if null fname then pure stdout else openFile fname WriteMode + forM_ exprs $ \case + Left err -> hPutStrLn h $ "invalid expression: " <> err + Right (t1, t2) -> do + hPutStrLn h ("First: " <> showOutput output t1) + hPutStrLn h ("Second: " <> showOutput output t2) + hFlush h + unless (null fname) $ hClose h + +hGetLines :: Handle -> IO [String] +hGetLines h = do + done <- hIsEOF h + if done + then return [] + else do + line <- hGetLine h + (line :) <$> hGetLines h diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -1,12 +1,12 @@ cabal-version: 1.12 --- This file has been generated from package.yaml by hpack version 0.35.2. +-- This file has been generated from package.yaml by hpack version 0.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 diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -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 ()