diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,97 @@ # Changelog for srtree +## 3.0.0.2 + +- Added parser for NeoGP.jl + +## 3.0.0.1 + +- Fixed some wrong bounds + +## 3.0.0.0 + +- **BREAKING**: Removed the Accelerate AD backend (`Algorithm.SRTree.AD.Accelerate`). + The `ADBackEnd` type now only has `SingleThread` and `MultiThread` constructors. + This removes the `accelerate` and `accelerate-llvm-native` dependencies. +- Out-of-core equality saturation with paged e-graph store (SQLite/PostgreSQL) +- Frontier re-saturation: mark changed classes and re-saturate only the frontier +- Streaming matcher for n-ary and cached genericJoin paths (O(1) memory on paged graphs) +- Cycle-safe and size-budgeted `getBestExpr` extraction +- Bounded cost/best fixpoints so recalc terminates on cyclic graphs +- Bounded node-to-class and canonical maps on paged graphs (LRU caches) +- Fast ByteString double parser for dataset loading +- Thread `Loss` (not `Distribution`) through fitness functions; add `readLoss` +- Multiset e-graph improvements + +## 2.0.1.7 + +- Added log10 MSE fitness function + +## 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/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,5 @@ -Copyright Author name here (c) 2021 +Copyright (c) 2026, folivetti -All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: @@ -13,7 +12,7 @@ disclaimer in the documentation and/or other materials provided with the distribution. - * Neither the name of Author name here nor the names of other + * Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. @@ -21,7 +20,7 @@ "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY diff --git a/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/Bench/Main.hs b/apps/Bench/Main.hs new file mode 100644 --- /dev/null +++ b/apps/Bench/Main.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE BangPatterns #-} + +import Criterion.Main +import Control.DeepSeq (force, NFData) +import Control.Exception (evaluate) +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector as VB +import qualified Data.Vector.Generic as G +import qualified Data.Vector.Storable as VS + +import Data.SRTree +import Data.SRTree.Print +import Data.SRTree.Datasets +import Data.SRTree.Eval +import Data.SRTree.Random +import System.Random +import Control.Monad.State.Strict +import Algorithm.SRTree.NonlinearOpt +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.AD + +-- Assuming these are exported by your project modules: +-- import SRTree +-- import Compiler +-- import DatasetLoader + +-- Mock signatures based on your provided functions +-- randomTree :: Int -> Int -> Int -> IO Term -> IO NonTerm -> Bool -> IO Tree +-- loadDataset :: FilePath -> Bool -> IO [V.Vector Double] +-- evalTree :: Tree -> [V.Vector Double] -> V.Vector Double +-- compile :: [V.Vector Double] -> Tree -> (Theta -> V.Vector Double) + +genTerm = do coin <- tossBiased 0.4 + if coin then randomFrom [Fix $ Var ix | ix <- [0..8]] else randomFrom [Fix $ Param ix | ix <- [0..9]] +genNonTerm = randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Uni LogAbs (), Uni SqrtAbs ()] + +genMultipleTrees 0 = pure [] +genMultipleTrees n = do + t <- randomTree 5 10 150 genTerm genNonTerm False + ts <- genMultipleTrees (n-1) + pure (t:ts) + +getF (_, x, _) = x +{-# INLINE getF #-} +getT (t, _, _) = t +{-# INLINE getT #-} + +main :: IO () +main = do + -- 1. Initialization: Load the dataset + putStrLn "Loading dataset..." + ((dataset, y, _, _), _, _, _) <- loadDataset "data.tsv" True + + -- 2. Initialization: Generate the random expression tree + putStrLn "Generating random tree..." + -- Replace 'genTerm' and 'genNonTerm' with your actual generators + --g <- getStdGen + let g = mkStdGen 42 + -- tree <- evalStateT (randomTree 7 10 150 genTerm genNonTerm True) g + trees' <- evalStateT (genMultipleTrees 5) g + -- let trees' = [Fix (Uni LogAbs (Fix (Bin PowerAbs (param 0) (param 1 * var 0))))] :: [Fix SRTree] + + -- IMPORTANT: Force deep evaluation of the tree and dataset. + -- If we do not do this, GHC's lazy evaluation will cause the benchmark + -- to measure the time it takes to parse the CSV and build the tree in memory! + -- _ <- evaluate (force tree) + _ <- evaluate (force dataset) + + + -- 3. Initialization: Pre-compile the tree + -- We evaluate this strictly (!) so the one-time compilation cost + -- is not included in the runtime benchmark. + putStrLn "Compiling tree..." + let !compiledFn = [compile dataset tree | tree <- trees] + evalTree x th t = compile x t th + -- Mock theta (parameter vector) to pass into the closures + !theta = V.fromList [1.0, 0.5, 0.2, 0.3, 0.1, 0.5, 0.9, 0.3, 0.2, 0.4] + !theta1 = V.fromList [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + trees = map relabelParamsOrder $ filter (\t -> let v = V.sum (evalTree dataset theta t) in not (isInfinite v || isNaN v)) trees' + naiveEval = evalTree dataset theta + dataset' = map G.convert dataset + y' = G.convert y + theta1' = G.convert theta1 + + _ <- evaluate (force theta) + _ <- evaluate (force theta1) + print $ sum $ map (\t -> V.sum $ naiveEval t) trees + print $ sum $ map (\t -> V.sum $ t theta) compiledFn + print $ sum $ map (\t -> getF $ minimizeNLL MultiThread MSE Nothing 0 dataset y t theta1) trees + --print $ sum $ map (\t -> getF $ minimizeNLLCompiled MSE Nothing 0 dataset y t theta1) trees + + --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraph MSE dataset' y' Nothing t theta1') trees + --print $ sum $ map (\t -> VS.sum . snd $ gradNLLGraphO MSE dataset' y' Nothing t theta1') trees + --print $ sum $ map (\t -> VS.sum . snd $ compileGrad dataset' y' Nothing t 100 theta1') trees + --print $ sum $ map (\ct -> V.sum $ ct theta) compiledFn + --print $ sum $ map (\ct -> V.sum $ executeVM ct rowDataset theta) bytecodes + -- print $ V.sum $ evalTree dataset theta tree + -- print $ V.sum $ compiledFn theta + + putStrLn "Running benchmarks..." + + -- 4. The Benchmarks + defaultMain [ + bgroup "Tree Evaluation (Fixed Dataset)" [ + + -- The slow version: dynamically traversing the AST at runtime + bench "evalTree (Naive AST Traversal)" $ + nf (\ts -> sum [V.sum $ evalTree dataset theta1 t | t <- ts]) trees, + + + -- The fast version: executing the pre-compiled, stream-fused closure + bench "compile (Compiled Closure)" $ + nf (\t -> sum [V.sum (ct t) | ct <- compiledFn]) theta1, + + -- The fast version: executing the pre-compiled, stream-fused closure + bench "minimizeNLLCompiled (Compiled Closure)" $ + nf (\ts -> sum [V.sum . getT $ minimizeNLL MultiThread MSE Nothing 100 dataset' y' t theta1' | t <- ts]) trees + + --bench "minimizeNLLO (Naive optimized AST Traversal)" $ + -- nf (\ts -> sum [V.sum . getT $ minimizeNLLO MSE Nothing 100 dataset y t theta1 | t <- ts]) trees + + -- The slow version: dynamically traversing the AST at runtime + --bench "minimizeNLL (Naive AST Traversal)" $ + -- nf (\ts -> sum [V.sum . getT $ minimizeNLL (NLL MSE) Nothing 100 dataset y t theta1 | t <- ts]) trees + + ] + ] diff --git a/apps/BenchEqSat/Main.hs b/apps/BenchEqSat/Main.hs new file mode 100644 --- /dev/null +++ b/apps/BenchEqSat/Main.hs @@ -0,0 +1,259 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +import Criterion.Main +import qualified Data.Vector.Unboxed as VU +import qualified Data.IntMap as IntMap +import qualified Data.HashMap.Strict as HashMap +import qualified Data.HashSet as Set + +import Data.SRTree +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Control.Monad.State.Strict +import Control.Monad (replicateM, zipWithM_) +import Control.Monad.Identity + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin _ l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +evalEG :: EGraphST Identity a -> (a, EGraph) +evalEG m = runIdentity $ runStateT m emptyGraph + +runInEG :: EGraph -> EGraphST Identity a -> (a, EGraph) +runInEG eg m = runIdentity $ runStateT m eg + +-- Expression generators for benchmarking +chainAdd :: Int -> Fix SRTree +chainAdd 0 = var 0 +chainAdd n = chainAdd (n-1) + var n + +deepBinTree :: Int -> Fix SRTree +deepBinTree 0 = var 0 +deepBinTree n = deepBinTree (n-1) + constv (fromIntegral n) + +complexTree :: Int -> Fix SRTree +complexTree n = go n + where + go 0 = var 0 + go i = (var i + constv (fromIntegral i)) * (go (i-1) + constv (fromIntegral i)) + +simplifyRules :: [Rule] +simplifyRules = + [ "a" + 0 :=> "a" + , "a" * 1 :=> "a" + , "a" + "a" :=> 2 * "a" + , "a" * 0 :=> 0 + , 0 + "a" :=> "a" + , 1 * "a" :=> "a" + ] + +-- More rules including commutativity (triggers more merges) +moreRules :: [Rule] +moreRules = + [ "a" + 0 :=> "a" + , "a" * 1 :=> "a" + , "a" + "a" :=> 2 * "a" + , "a" * 0 :=> 0 + , 0 + "a" :=> "a" + , 1 * "a" :=> "a" + , "a" + "b" :=> "b" + "a" + , "a" * "b" :=> "b" * "a" + ] + +addZero :: Fix SRTree -> Fix SRTree -> Fix SRTree +addZero l r = Fix (Bin Add l r) + +main :: IO () +main = do + putStrLn "Generating benchmark expressions..." + let smallExpr = chainAdd 5 + mediumExpr = chainAdd 20 + largeExpr = chainAdd 100 + complex = complexTree 8 + + putStrLn "Running benchmarks..." + defaultMain [ + bgroup "E-graph Construction" [ + bench "fromTree (5-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) smallExpr, + bench "fromTree (20-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) mediumExpr, + bench "fromTree (100-leaf chain)" $ + whnf (\e -> evalEG $ fromTree myCost e) largeExpr, + bench "fromTree (complex-ternary tree)" $ + whnf (\e -> evalEG $ fromTree myCost e) complex + ], + + bgroup "E-graph Add" [ + bench "add single e-node (Var)" $ + whnf (\eg -> runInEG eg $ add myCost (EVar 999)) (snd $ evalEG $ fromTree myCost smallExpr), + bench "add single e-node (Const)" $ + whnf (\eg -> runInEG eg $ add myCost (EConst 42.0)) (snd $ evalEG $ fromTree myCost smallExpr), + bench "add single e-node (Bin Add)" $ + whnf (\eg -> runInEG eg $ add myCost (ENAry EAdd (imFromList [0, 1]))) (snd $ evalEG $ fromTree myCost mediumExpr) + ], + + bgroup "Merge" [ + bench "merge two distinct eclasses (size 1)" $ + whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 1), + bench "merge two distinct eclasses (size 3)" $ + whnf (\(e1,e2,eg) -> runInEG eg $ merge myCost e1 e2) (makeMergePair 3) + ], + + bgroup "Pattern Matching" [ + bench "match simple pattern (a+0)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMatchableEG), + bench "match commutative pattern (a+b)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMatchableEG), + bench "match triple pattern (a+b+c)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" + "c" :: Pattern)) (makeMatchableEG) + ], + + bgroup "Match After Merge" [ + bench "match (a+0) after merge (stale trie keys)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + 0 :: Pattern)) (makeMergedEG), + bench "match (a+b) after merge (stale trie keys)" $ + whnf (\(eg,_) -> runInEG eg $ match ("a" + "b" :: Pattern)) (makeMergedEG) + ], + + bgroup "Rebuild" [ + bench "rebuild after 5 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 5), + bench "rebuild after 20 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 20), + bench "rebuild after 100 adds" $ + whnf (\(eg,_) -> runInEG eg $ rebuild myCost) (makeDirtyEG 100) + ], + + bgroup "Cost Propagation" [ + bench "recalculateBest (10 eclasses)" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 10), + bench "recalculateBest (100 eclasses)" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (recalculateBest myCost) eids) (makeNEclasses 100) + ], + + bgroup "DB Operations" [ + bench "addToDB single enode" $ + whnf (\(en,eid,eg) -> runInEG eg $ addToDB en eid) (makeDBEntry), + bench "addToDB 10 enodes" $ + whnf (\(ens,eg) -> runInEG eg $ mapM_ (uncurry addToDB) ens) (makeDBEntries 10) + ], + + bgroup "Equality Saturation" [ + bench "eqSat small expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, simplifyRules), + bench "eqSat medium expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (mediumExpr, simplifyRules), + bench "eqSat small expr (8 rules, commutative)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (smallExpr, moreRules), + bench "eqSat large expr (5 rules)" $ + whnf (\(e,r) -> evalEG $ eqSat e r myCost 10) (largeExpr, simplifyRules) + ], + + bgroup "Extraction" [ + bench "getBestExpr (5-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 5), + bench "getBestExpr (20-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 20), + bench "getBestExpr (100-leaf)" $ + whnf (\(eid,eg) -> runInEG eg $ getBestExpr eid) (makeExtractable 100) + ], + + bgroup "Fitness Operations" [ + bench "insertFitness single" $ + whnf (\(eid,eg) -> runInEG eg $ insertFitness eid 0.5 []) (makeExtractable 1), + bench "insertFitness 10 eclasses" $ + whnf (\(eids,eg) -> runInEG eg $ mapM_ (\eid -> insertFitness eid 0.5 []) eids) (makeNEclasses 10), + bench "getTopFitEClassWithSize" $ + whnf (\(eids,eg) -> runInEG eg $ getTopFitEClassWithSize 1 3) (makeFitnessEG) + ] + ] + where + addZeroTree = addZero (var 0) (constv 0.0) + + makeMergePair :: Int -> (EClassId, EClassId, EGraph) + makeMergePair n = + let tree = deepBinTree n + (eid1, eg1) = evalEG $ fromTree myCost tree + (eid2, eg2) = runInEG eg1 $ fromTree myCost tree + in (eid1, eid2, eg2) + + makeMatchableEG :: (EGraph, EClassId) + makeMatchableEG = + let tree = complexTree 4 + (eid, eg) = evalEG $ do + eid' <- fromTree myCost tree + _ <- fromTree myCost (var 0 + constv 1.0) + _ <- fromTree myCost (var 1 * constv 2.0) + _ <- fromTree myCost (var 0 + constv 0.0) + _ <- fromTree myCost (var 1 * constv 1.0) + rebuild myCost + pure eid' + in (eg, eid) + + -- E-graph with merges applied, creating stale trie keys + makeMergedEG :: (EGraph, EClassId) + makeMergedEG = + let (_, eg) = evalEG $ do + eid1 <- fromTree myCost (var 0) + eid2 <- fromTree myCost (constv 0.0) + eid3 <- fromTree myCost (var 0 + constv 1.0) + _ <- fromTree myCost (var 1) + rebuild myCost + -- merge to create stale trie entries + merge myCost eid1 eid2 + merge myCost eid2 eid3 + rebuild myCost + pure eid1 + in (eg, 0) + + makeDirtyEG :: Int -> (EGraph, EClassId) + makeDirtyEG n = + let tree = deepBinTree n + (eid, eg) = evalEG $ do + eid' <- fromTree myCost tree + _ <- fromTree myCost (tree + var 999) + rebuild myCost + _ <- fromTree myCost (tree * var 998) + pure eid' + in (eg, eid) + + makeExtractable :: Int -> (EClassId, EGraph) + makeExtractable n = + let tree = deepBinTree n + in evalEG $ fromTree myCost tree + + makeNEclasses :: Int -> ([EClassId], EGraph) + makeNEclasses n = + evalEG $ replicateM n (fromTree myCost (constv (fromIntegral n))) + + makeFitnessEG :: ([EClassId], EGraph) + makeFitnessEG = evalEG $ do + eids <- mapM (fromTree myCost . constv . fromIntegral) [1..10] + zipWithM_ (\eid i -> insertFitness eid (fromIntegral i) []) eids [1..] + pure eids + + makeDBEntry :: (ENode, EClassId, EGraph) + makeDBEntry = + let (eid, eg) = evalEG $ do + eid <- fromTree myCost (var 999) + rebuild myCost + pure eid + in (EVar 777, eid, eg) + + makeDBEntries :: Int -> ([(ENode, EClassId)], EGraph) + makeDBEntries n = + let (eids, eg) = evalEG $ do + eids <- mapM (fromTree myCost . var) [999..(999 + n - 1)] + rebuild myCost + pure eids + in (zip (map EVar [1000..]) eids, eg) diff --git a/apps/Report/Main.hs b/apps/Report/Main.hs new file mode 100644 --- /dev/null +++ b/apps/Report/Main.hs @@ -0,0 +1,247 @@ +module Main (main) where + +import Options.Applicative +import qualified Data.ByteString.Char8 as B +import qualified Data.Vector.Unboxed as U +import Data.SRTree +import Data.SRTree.Eval (Target, Columns, compileLoss) +import Data.SRTree.Datasets (loadTrainingOnly) +import Data.SRTree.Print (showExpr) +import Text.ParseSR (parseSR, SRAlgs(..)) +import Algorithm.SRTree.Compile (compileTree, EvalTree(..), logParameters, logParametersLatt) +import Algorithm.SRTree.Likelihoods (Distribution(..), Loss(..), buildLoss, fisherNLL, hessianNLL) +import Algorithm.SRTree.ConfidenceIntervals + ( getStatsFromModel, paramCI, CIType(..), CI(..), BasicStats(..) + , ProfileT(..), PType(..), getAllProfiles, getCol + ) +import Algorithm.SRTree.ModelSelection (ModelEval(..), logFunctional, logFunctionalFreq) +import Statistics.Distribution (ContDistr(quantile)) +import Statistics.Distribution.FDistribution (fDistribution) +import Control.Exception (try, SomeException) +import Data.List.Split (splitOn) +import Text.Printf (printf) +import Control.Monad (forM_, when) + +---------------------------------------------------------------------- +-- CLI argument types +---------------------------------------------------------------------- +data CIMethod = LaplaceCI | ProfileCI deriving (Show) + +data ProfileTypeArg = BatesArg | ODEArg | ConstrainedArg deriving (Read) +instance Show ProfileTypeArg where + show BatesArg = "Bates" + show ODEArg = "ODE" + show ConstrainedArg = "Constrained" + +data ReportArgs = ReportArgs + { raExprs :: !FilePath + , raFormat :: !SRAlgs + , raData :: !FilePath + , raHeader :: !Bool + , raDist :: !Distribution + , raCriteria :: ![ModelEval] + , raCI :: !CIMethod + , raAlpha :: !Double + , raCIType :: !ProfileTypeArg + , raDbg :: !Bool + } + +---------------------------------------------------------------------- +-- Argument parser +---------------------------------------------------------------------- +argParser :: Parser ReportArgs +argParser = ReportArgs + <$> strOption ( long "exprs" <> short 'e' <> help "File with expressions, one per line" <> metavar "FILE" ) + <*> option auto ( long "format" <> short 'f' <> help "Expression format: TIR, HL, OPERON, BINGO, GOMEA, PYSR, SBP, EPLEX" <> metavar "FMT" ) + <*> strOption ( long "data" <> short 'd' <> help "Dataset file (optionally with :start:end:target:features:y_err)" <> metavar "FILE" ) + <*> switch ( long "header" <> help "Dataset has a header row" ) + <*> option auto ( long "dist" <> value Gaussian <> help "Distribution: Gaussian, Bernoulli, Poisson, LeastSquares" <> metavar "DIST" <> showDefault ) + <*> option parseCriteria ( long "criteria" <> short 'c' <> value [RMSE, R2, AIC, BIC] <> help "Comma-separated criteria" <> metavar "CRITERIA" <> showDefault ) + <*> option parseCI ( long "ci" <> value LaplaceCI <> help "CI method: Laplace, Profile" <> metavar "METHOD" <> showDefault ) + <*> option auto ( long "alpha" <> value 0.05 <> help "Significance level" <> metavar "ALPHA" <> showDefault ) + <*> option parseProfileType ( long "ci-type" <> value BatesArg <> help "Profile CI type: Bates, ODE, Constrained" <> metavar "TYPE" <> showDefault ) + <*> switch ( long "dbg" <> help "Debug: dump profile tau/theta spline points" ) + +parseCriteria :: ReadM [ModelEval] +parseCriteria = eitherReader $ \s -> + case traverse parseOne (splitOn "," s) of + Right es -> Right es + Left e -> Left e + where + parseOne "RMSE" = Right RMSE + parseOne "R2" = Right R2 + parseOne "AIC" = Right AIC + parseOne "BIC" = Right BIC + parseOne "Evidence" = Right Evidence + parseOne "FBF" = Right FBF + parseOne "MDL" = Right MDL + parseOne "MDLLatt" = Right MDLLatt + parseOne "MDLFreq" = Right MDLFreq + parseOne "NLL" = Right (EvalLoss (NLL Gaussian)) + parseOne s = Left ("unknown criterion: " ++ s) + +parseCI :: ReadM CIMethod +parseCI = eitherReader $ \s -> case s of + "Laplace" -> Right LaplaceCI + "Profile" -> Right ProfileCI + _ -> Left ("unknown CI method: " ++ s ++ " (use Laplace or Profile)") + +parseProfileType :: ReadM ProfileTypeArg +parseProfileType = eitherReader $ \s -> case s of + "Bates" -> Right BatesArg + "ODE" -> Right ODEArg + "Constrained" -> Right ConstrainedArg + _ -> Left ("unknown profile type: " ++ s ++ " (use Bates, ODE, or Constrained)") + +---------------------------------------------------------------------- +-- Report data +---------------------------------------------------------------------- +data ReportData = ReportData + { rdTree :: Fix SRTree + , rdTheta :: Target + , rdStdErr :: Target + , rdCriteria :: [(ModelEval, Double)] + , rdCIs :: [CI] + } + +---------------------------------------------------------------------- +-- Main +---------------------------------------------------------------------- +main :: IO () +main = do + args <- execParser (info (argParser <**> helper) fullDesc) + (xss, ys, mYerr) <- loadTrainingOnly (raData args) (raHeader args) + content <- B.readFile (raExprs args) + let exprs = filter (not . B.null) $ B.lines content + mapM_ (processOne args xss ys mYerr) (zip [(1 :: Int) ..] exprs) + +---------------------------------------------------------------------- +-- Process a single expression +---------------------------------------------------------------------- +processOne :: ReportArgs -> Columns -> Target -> Maybe Target -> (Int, B.ByteString) -> IO () +processOne args xss ys mYerr (idx, src) = do + result <- try $ do + tree <- case parseSR (raFormat args) B.empty True src of + Left e -> fail ("parse error: " ++ e) + Right t -> return $! relabelParams t + let dist = raDist args + nRows = U.length ys + nModelParams = countParamsUniq tree + nParams = nModelParams + + case dist of + Gaussian -> 1 + ROXY -> 3 + _ -> 0 + + let et = compileTree dist xss ys mYerr tree + theta0 = U.replicate nParams 1.0 + thetaOpt = ctOptimizer et theta0 + + when (any isNaN (U.toList thetaOpt)) $ + fail "optimisation returned NaN" + + let mseTree = buildLoss MSE (fromIntegral nRows) tree + mseLoss = compileLoss xss mseTree ys mYerr thetaOpt + nllLoss = ctNLL et thetaOpt + tss = ctVar et + + let fisherDiag = fisherNLL dist mYerr xss ys tree thetaOpt + hessCols = hessianNLL dist mYerr xss ys tree thetaOpt + hessLists = map U.toList hessCols + logP = logParameters fisherDiag thetaOpt + logPLatt = logParametersLatt hessLists fisherDiag thetaOpt + logF = logFunctional tree + logFFreq = logFunctionalFreq tree + nF = fromIntegral nRows + kF = fromIntegral nParams + crits = map (\c -> (c, evalOne c mseLoss nllLoss tss nF kF logP logPLatt logF logFFreq)) + (raCriteria args) + + let stats = getStatsFromModel dist mYerr xss ys tree thetaOpt + laplaceCIs = paramCI (Laplace stats) nRows thetaOpt (raAlpha args) + let ptype = case raCIType args of + BatesArg -> Bates + ODEArg -> ODE + ConstrainedArg -> Constrained + let kInt = U.length thetaOpt + nInt = U.length ys + profT = sqrt $ quantile (fDistribution (fromIntegral kInt) (fromIntegral $ nInt - kInt)) (1 - raAlpha args) + cis <- case raCI args of + LaplaceCI -> return laplaceCIs + ProfileCI -> do + let profiles = getAllProfiles ptype et thetaOpt (_stdErr stats) laplaceCIs (raAlpha args) + when (raDbg args) $ forM_ (zip [0..] profiles) $ \(i, ProfileT taus thetas _ tau2theta _) -> do + putStrLn $ "DEBUG Profile " ++ show i ++ " (opt=" ++ show (thetaOpt U.! i) ++ "):" + putStrLn $ " tau range: [" ++ show (if U.null taus then 0 else U.head taus) + ++ ", " ++ show (if U.null taus then 0 else U.last taus) ++ "]" + putStrLn $ " t=" ++ show profT + putStrLn $ " tau2theta(-t)=" ++ show (tau2theta (-profT)) + ++ " tau2theta(+t)=" ++ show (tau2theta profT) + putStrLn $ " profile points:" + let tausL = U.toList taus + thetasL = U.toList (getCol i thetas) + forM_ (zip tausL thetasL) $ \(tau, th) -> + putStrLn $ " tau=" ++ show tau ++ " theta=" ++ show th + return $ paramCI (Profile stats profiles) nRows thetaOpt (raAlpha args) + + return $! ReportData + { rdTree = tree + , rdTheta = thetaOpt + , rdStdErr = _stdErr stats + , rdCriteria = crits + , rdCIs = cis + } + + case result of + Right rd -> printReport idx src rd + Left e -> printFailure idx src (show (e :: SomeException)) + +---------------------------------------------------------------------- +-- Evaluate a single ModelEval from base quantities +---------------------------------------------------------------------- +evalOne :: ModelEval -> Double -> Double -> Double -> Double -> Double + -> Double -> Double -> Double -> Double -> Double +evalOne RMSE mse _ _ _ _ _ _ _ _ = sqrt mse +evalOne R2 mse _ tss n _ _ _ _ _ = 1 - n * mse / tss +evalOne AIC _ nll _ _ k _ _ _ _ = 2*k + 2*nll +evalOne BIC _ nll _ n k _ _ _ _ = k * log n + 2*nll +evalOne Evidence _ nll _ n k _ _ _ _ = (1 - b) * nll - k/2 * log b + where b = 1 / sqrt n +evalOne FBF _ nll _ n k _ _ _ _ = res + where b = 1 / sqrt n; nup = exp (1 - log 3) + res = (1 - b) * nll - k/2 * log b + k/2 * log (2*pi*nup) +evalOne MDL _ nll _ _ _ logP _ logF _ = nll + logF + logP +evalOne MDLLatt _ nll _ _ _ _ logPL logF _ = nll + logF + logPL +evalOne MDLFreq _ nll _ _ _ logP _ _ logFF = nll + logFF + logP +evalOne (EvalLoss (NLL Gaussian)) _ nll _ _ _ _ _ _ _ = nll +evalOne _ _ _ _ _ _ _ _ _ _ = 0 -- unreachable + +---------------------------------------------------------------------- +-- Output +---------------------------------------------------------------------- +printReport :: Int -> B.ByteString -> ReportData -> IO () +printReport idx src rd = do + putStrLn $ "=== Expression " ++ show idx ++ " ===" + putStrLn $ "Tree: " ++ showExpr (rdTree rd) + putStrLn "Parameters:" + let thetaList = U.toList (rdTheta rd) + ciList = rdCIs rd + forM_ (zip3 [0..] thetaList ciList) $ \(i, th, ci) -> + putStrLn $ " theta" ++ show i ++ ": " ++ fmt th + ++ " [" ++ fmt (lower_ ci) ++ ", " ++ fmt (upper_ ci) ++ "]" + putStrLn "Model Selection:" + forM_ (rdCriteria rd) $ \(c, v) -> + putStrLn $ " " ++ padRight 12 (show c) ++ ": " ++ fmt v + putStrLn "" + where + fmt x | abs x < 1e-10 = "0.0000" + | abs x >= 1e4 = printf "%.4e" x + | otherwise = printf "%.6f" x + padRight n s = s ++ replicate (max 0 (n - length s)) ' ' + +printFailure :: Int -> B.ByteString -> String -> IO () +printFailure idx src msg = do + putStrLn $ "=== Expression " ++ show idx ++ " ===" + putStrLn $ "Tree: " ++ B.unpack src + putStrLn $ "Error: " ++ msg + putStrLn "" diff --git a/src/Algorithm/EqSat.hs b/src/Algorithm/EqSat.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat.hs @@ -0,0 +1,380 @@ +{-# 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 qualified Data.IntSet as IntSet +import Data.List (intercalate) +import Data.Map (Map) +import qualified Data.Map as Map +import Data.Maybe (mapMaybe) +import Data.SRTree +import Data.HashSet (HashSet) +import qualified Data.HashSet as Set +import Control.Monad ( zipWithM, forM_ ) + +-- | The `Scheduler` stores a map with the banned iterations of a certain rule . +-- TODO: make it more customizable. +type Scheduler a = State (IntMap Int) a + +-- to avoid importing +-- | runs equality saturation from an expression tree, +-- a given set of rules, and a cost function. +-- Returns the tree with the smallest cost. +eqSat :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) +eqSat expr rules costFun maxIt = + do root <- fromTree costFun expr + _ <- runEqSat costFun rules maxIt + recalculateBest costFun root + +type CostMap = Map EClassId (Int, Fix SRTree) + +-- | recalculates the costs with a new cost function +recalculateBest :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree) +recalculateBest costFun eid = + do ecls <- allClasses + let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls] + costs = fillUpCosts classes Map.empty + eid' <- canonical eid + case Map.lookup eid' costs of + Just (_, t) -> pure t + Nothing -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid' + <> " nClasses=" <> show (IntMap.size classes) + <> " costSize=" <> show (Map.size costs) + where + nodeCost :: CostMap -> ENode -> (Int, Fix SRTree) + nodeCost costMap enode = + -- A child that has not been costed yet (a cycle, or a class whose + -- cost is computed later in this iteration) contributes a large + -- sentinel instead of 0: a 0 placeholder is cheaper than the real + -- cost, so the fixpoint below would keep the stale placeholder tree + -- (e.g. `x * 0.0` for `x * (y + z)`). Real costs always beat it. + let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (costMap Map.!? cid) | cid <- eChildren enode ] + c = case enode of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + _ -> costFun (replaceChildren cc (fromENode enode)) + in (c + sum cc, Fix $ case enode of + ENAry op _ -> unfix (naryTree op nc) + _ -> replaceChildren nc (fromENode enode)) -- | missing children (cyclic classes) get cost 0 so every class is costed + costSentinel :: Int + costSentinel = 1000000 + + fillUpCosts :: IntMap EClass -> CostMap -> CostMap + fillUpCosts classes = go (IntMap.size classes + 1) (IntMap.keysSet classes) + where + go 0 _ m = m + go n dirty m + | IntSet.null dirty = m + | otherwise = go (n - 1) dirty' m' + where + (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty + step (d, cm) eid = case IntMap.lookup eid classes of + Nothing -> (d, cm) + Just ecl -> + let currentCost = Map.lookup eid cm + minCost = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c') + ) Nothing (_eNodes ecl) + (changed, cm') = case (currentCost, minCost) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, Map.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, Map.insert eid new cm) + d' = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl) + else d + in d' `seq` cm' `seq` (d', cm') + +-- | Recompute every e-class's cost-minimal @_best@/_cost@ bottom-up and write +-- it back into the graph. Needed after loading a graph whose best/cost were +-- not persisted (e.g. via srtree-db), where @_best@ may otherwise hold an +-- arbitrary (potentially large) e-node. +recalculateBestAll :: ClassStore m => CostFun -> EGraphST m () +recalculateBestAll costFun = do + ecls <- allClasses + let classes = IntMap.fromList [(_eClassId ec, ec) | ec <- ecls] + bests = fixpoint classes IntMap.empty + forM_ (IntMap.toList bests) $ \(eid, (c, en)) -> + case IntMap.lookup eid classes of + Nothing -> pure () + Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } } + where + nodeCost :: IntMap (Int, ENode) -> ENode -> (Int, ENode) + nodeCost cm en = + let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc + _ -> costFun (replaceChildren cc (fromENode en)) + sum cc + in (c, en) + costSentinel :: Int + costSentinel = 1000000 + + fixpoint :: IntMap EClass -> IntMap (Int, ENode) -> IntMap (Int, ENode) + fixpoint classes0 = go (IntMap.size classes0 + 1) (IntMap.keysSet classes0) + where + go 0 _ m = m + go n dirty m + | IntSet.null dirty = m + | otherwise = go (n - 1) dirty' m' + where + (dirty', m') = IntSet.foldl' step (IntSet.empty, m) dirty + step (d, cm) eid = case IntMap.lookup eid classes0 of + Nothing -> (d, cm) + Just ecl -> + let current = IntMap.lookup eid cm + minNode = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minNode) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, IntMap.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, IntMap.insert eid new cm) + d' = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) d (_parents ecl) + else d + in d' `seq` cm' `seq` (d', cm') + +-- | Like 'recalculateBestAll' but streamed: each e-class body is fetched on +-- demand through 'ClassStore' (so a paged graph never materializes every body +-- at once) and only the small @(cost, best e-node)@ map is kept resident. The +-- structural worklist fixpoint is identical. +recalculateBestAllStream :: ClassStore m => CostFun -> EGraphST m () +recalculateBestAllStream costFun = do + ids <- allKeys + let idSet = IntSet.fromList ids + costSentinel = 1000000 + nodeCost cm en = + let cc = [ maybe costSentinel fst (IntMap.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + sum cc + _ -> costFun (replaceChildren cc (fromENode en)) + sum cc + in (c, en) + stepEid cm eid = do + mec <- readDirect eid + case mec of + Nothing -> pure (IntSet.empty, cm) + Just ecl -> do + let current = IntMap.lookup eid cm + minNode = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minNode) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, IntMap.insert eid new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, IntMap.insert eid new cm) + dirty = if changed + then Set.foldl' (\acc (pid, _) -> IntSet.insert pid acc) IntSet.empty (_parents ecl) + else IntSet.empty + pure (dirty, cm') + fixpoint n dirty cm + | n <= 0 || IntSet.null dirty = pure cm + | otherwise = go (IntSet.toList dirty) IntSet.empty cm + where + go [] d acc = fixpoint (n - 1) d acc + go (e : es) d acc = do + (d', m') <- stepEid acc e + go es (IntSet.union d d') m' + cm <- fixpoint (IntSet.size idSet + 1) idSet IntMap.empty + forM_ (IntMap.toList cm) $ \(eid, (c, en)) -> do + mec <- readDirect eid + case mec of + Nothing -> pure () + Just ec -> writeDirect ec { _info = (_info ec) { _cost = c, _best = en } } + +-- | Streaming variant of 'recalculateBest': computes the cost-minimal tree for a +-- single root without materializing every e-class body at once. +recalculateBestStream :: ClassStore m => CostFun -> EClassId -> EGraphST m (Fix SRTree) +recalculateBestStream costFun eid = do + ids <- allKeys + let idSet = IntSet.fromList ids + costSentinel = 1000000 + nodeCost cm en = + let (cc, nc) = unzip [ maybe (costSentinel, Fix (Const 0)) id (Map.lookup cid cm) | cid <- eChildren en ] + c = case en of + ENAry op _ -> costFun (Bin (toOp op) 0 0) + _ -> costFun (replaceChildren cc (fromENode en)) + in (c + sum cc, Fix $ case en of + ENAry op _ -> unfix (naryTree op nc) + _ -> replaceChildren nc (fromENode en)) + stepEid cm eid' = do + mec <- lookupClass eid' + case mec of + Nothing -> pure (IntSet.empty, cm) + Just ecl -> do + let current = Map.lookup eid' cm + minCost = Set.foldl' (\acc en -> let c = nodeCost cm en + in case acc of + Nothing -> Just c + Just c' -> Just (if fst c <= fst c' then c else c')) + Nothing (_eNodes ecl) + (changed, cm') = case (current, minCost) of + (_, Nothing) -> (False, cm) + (Nothing, Just new) -> (True, Map.insert eid' new cm) + (Just old, Just new) + | fst old <= fst new -> (False, cm) + | otherwise -> (True, Map.insert eid' new cm) + dirty = if changed + then Set.foldl' (\acc (pid,_) -> IntSet.insert pid acc) IntSet.empty (_parents ecl) + else IntSet.empty + pure (dirty, cm') + fixpoint n dirty cm + | n <= 0 || IntSet.null dirty = pure cm + | otherwise = go (IntSet.toList dirty) IntSet.empty cm + where + go [] d acc = fixpoint (n - 1) d acc + go (e : es) d acc = do + (d', m') <- stepEid acc e + go es (IntSet.union d d') m' + cm <- fixpoint (IntSet.size idSet + 1) idSet Map.empty + eid' <- canonical eid + case Map.lookup eid' cm of + Just (_, t) -> pure t + Nothing -> error $ "EQSAT_RECALC_MISSING eid=" <> show eid' + <> " costSize=" <> show (Map.size cm) + +-- | Run equality saturation and stream the final extraction (see +-- 'recalculateBestStream'), so a paged graph is never fully materialized. +eqSatStream :: ClassStore m => Fix SRTree -> [Rule] -> CostFun -> Int -> EGraphST m (Fix SRTree) +eqSatStream expr rules costFun maxIt = do + root <- fromTree costFun expr + _ <- runEqSat costFun rules maxIt + recalculateBestAllStream costFun + recalculateBestStream costFun root + +-- | replaces the equality rules with two one-way rules +replaceEqRules :: Rule -> [Rule] +replaceEqRules (p1 :=> p2) = [p1 :=> p2] +replaceEqRules (p1 :==: p2) = [p1 :=> p2, p2 :=> p1] +replaceEqRules (r :| cond) = map (:| cond) $ replaceEqRules r + +-- | Compile a rule source into a query, or `Nothing` for n-ary patterns that +-- use the direct multiset matcher instead. +compileSource :: Rule -> Maybe (Query, [ClassOrVar], ClassOrVar) +compileSource r = if hasNAry (source r) + then Nothing + else Just (compileToQuery (source r)) + +-- | Cap on the total number of rule matches applied in a single eqsat +-- iteration. Combined with the per-rule caps ('ruleBudget'/'ruleRootVisit' for +-- n-ary, 'ruleMatchBudget' for the cached path) and the persistent +-- mark-on-attempt seen-set (which makes each rule's budget advance to new +-- matches), this bounds a single iteration's apply/rebuild work regardless of +-- graph size. +iterMatchBudget :: Int +iterMatchBudget = 2000 + +-- | run equality saturation for a number of iterations +runEqSat :: ClassStore m => CostFun -> [Rule] -> Int -> EGraphST m (Bool, Int) +runEqSat costFun rules maxIter = go maxIter IntMap.empty compiledRules + where + rules' = concatMap replaceEqRules rules + compiledRules = map (\r -> (r, compileSource r)) rules' + + go it sch compiled = + do -- reset dirty flag before processing this iteration + modify' $ over (eDB . changed) (const False) + + -- step 1: match the rules using cached compiled queries + let matchSch = matchWithScheduler it + adapted i (r, cq) = map (,cq) <$> matchSch i r + matchAll = zipWithM adapted [0..] + (filtered, sch') = runState (matchAll compiled) sch + + -- step 2: apply matches and rebuild + matches <- mapM (\(rule, cq) -> map (rule,) <$> case cq of + Just q -> do paged <- isPagedGraph + if paged + then matchStreamCached (Just (show (source rule))) (source rule) + else matchCachedWith (Just (show (source rule))) q + Nothing -> matchSaturated (source rule)) $ concat filtered + -- bound the total number of matches applied per iteration so a + -- single iteration's apply/rebuild work stays bounded on huge + -- graphs (genuine matches; we just process them over more iters). + mapM_ (uncurry (applyMatch costFun)) (take iterMatchBudget (concat matches)) + rebuild costFun + + -- check dirty flag: if no modifications occurred, we've saturated + changed <- gets (_changed . _eDB) + if it == 1 || not changed + then pure (True, it) + else + do eClasses <- gets _eClass + if IntMap.size eClasses > 1500 + then throttle it sch' compiled + else go (it-1) sch' compiled + + throttle it sch compiled = do + cleanMaps + eClasses <- gets _eClass + if IntMap.size eClasses <= 1500 + then go (it-1) sch compiled + else do applySingleMergeOnlyEqSat costFun rules + changed <- gets (_changed . _eDB) + if it <= 1 || not changed + then pure (False, it) -- give up and return early stop + else throttle (it-1) sch compiled + +-- | apply a single step of merge-only equality saturation +applySingleMergeOnlyEqSat :: ClassStore m => CostFun -> [Rule] -> EGraphST m () +applySingleMergeOnlyEqSat costFun rules = + do let matchSch = matchWithScheduler 10 + matchAll = zipWithM matchSch [0..] + (rls, _) = runState (matchAll rules') IntMap.empty + matches <- getNMatches 500 rls + rebuild costFun + where + rules' = concatMap replaceEqRules rules + + getNMatches n [] = pure [] + getNMatches 0 _ = pure [] + getNMatches n ([]:rss) = getNMatches n rss + getNMatches n ((r:rs):rss) = do matches <- map (r,) <$> matchSaturated (source r) + let (x, _) = splitAt n matches + m = length x + if m == n + then pure matches + 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 maybe False (<= it) mbBan -- check if the rule is banned + then pure [] + else do -- let matches = match db (source rule) + modify (IntMap.insert ruleNumber (it+5)) + 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,743 @@ +{-# 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 +import Data.SRTree +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.DB +import qualified Data.IntMap.Strict as IntMap +import Data.IntMap.Strict (IntMap) +import Data.Map.Strict ( Map ) +import qualified Data.Map.Strict as Map +import qualified Data.HashMap.Strict as HashMap +import qualified Data.HashSet as Set +import Control.Monad.State.Strict +import Control.Monad.Identity +import GHC.Stack (HasCallStack) + +import Data.SRTree.Recursion (cataM) +import Data.List (sort) +import Algorithm.EqSat.Info +import qualified Data.IntSet as IntSet + +import qualified Data.Set as RangeSet + + +-- | adds a new or existing e-node (merging if necessary) +add :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m EClassId +add costFun enode = do + enode'' <- canonize enode + enode''' <- foldConsts costFun enode'' + + maybeEid <- lookupNode enode''' + case maybeEid of + Just eid -> pure eid + Nothing -> do + curId <- gets (_nextId . _eDB) -- get the next available e-class id + insertCanonical curId curId -- register the class as its own representative + insertNode enode''' curId -- associate new e-node with id (bounded on paged graphs) + modify' $ over (eDB . nextId) (+1) -- update next id + . over (eDB . worklist) (Set.insert (curId, enode''')) -- add e-node and id into worklist + forM_ (eChildren enode''') (addParents curId enode''') -- update the children's parent list + info <- makeAnalysis costFun enode''' + h <- getChildrenMinHeight enode''' + let newClass = createEClass curId enode''' info h -- create e-class + -- insert via 'insertClass' so a paged (DB-backed) class store also + -- persists the new class's page; for a pure graph this is identical + -- to inserting into @_eClass@ directly. + insertClass newClass + --modifyEClass costFun curId -- simplify eclass if it evaluates to a number + + -- update database + addToDB enode''' curId -- add new node to db + tracking <- gets (_trackDBs . _eDB) + when tracking $ + modify' $ over (eDB . sizeDB) + $ IntMap.insertWith (IntSet.union) (_size info) (IntSet.singleton curId) + modify' $ over (eDB . unevaluated) (IntSet.insert curId) + . over (eDB . changed) (const True) + pure curId + where + addParents :: ClassStore m => EClassId -> ENode -> EClassId -> EGraphST m () + addParents cId node c = + do ec <- getEClass c + let ec' = ec{ _parents = Set.insert (cId, node) (_parents ec) } + -- write through 'insertClass' so a paged store keeps the updated parents + insertClass ec' + +-- | Add a binary (SRTree-based) node, converting it to a flattened ENode. +-- Sub and Div are canonicalized away at insertion: `x - y` becomes +-- `x + (-1)*y` and `x / y` becomes `x * recip y`, so no Sub/Div e-node ever +-- enters the e-graph and the Sub/Div-aware rules become redundant. +addTree :: (ClassStore m, HasCallStack) => CostFun -> SRTree EClassId -> EGraphST m EClassId +addTree costFun (Bin Sub l r) = do + neg <- addNegate costFun r + add costFun =<< mkENary EAdd [l, neg] +addTree costFun (Bin Div l r) = do + rec <- add costFun (EUni Recip r) + add costFun =<< mkENary EMul [l, rec] +addTree costFun t = toENode t >>= add costFun +{-# INLINE addTree #-} + +-- | builds the e-class for the negation of the e-class `t`, represented as +-- `(-1) * t` (matching the pattern-level `negate` encoding in Algorithm.EqSat.DB). +addNegate :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId +addNegate costFun t = do + negOne <- add costFun (EConst (-1)) + add costFun =<< mkENary EMul [negOne, t] + +-- | Fused 'calculateConsts' + 'foldConstants': fetches each child's constant +-- info a single time, detects fully-constant nodes (replaced by EConst/EParam) +-- and folds together all-but-one constant children of an ENAry +-- (e.g. 2+3+x becomes 5+x). Constants that are already folded single subtrees +-- are handled by the same child-constant walk. +foldConsts :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode +foldConsts _ en@(ENAry _ m) | IntMap.null m = pure en +foldConsts costFun en@(ENAry op m) = do + let xs = expandedList m + infos <- mapM (fmap (_consts . _info) . getEClass) xs + case foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) infos of + ConstVal x -> pure (EConst x) + ParamIx x -> pure (EParam x) + _ -> foldENary costFun op m infos +foldConsts _ en = do + infos <- mapM (fmap (_consts . _info) . getEClass) (eChildren en) + case combineConsts (replaceChildren infos (fromENode en)) of + ConstVal x -> pure (EConst x) + ParamIx x -> pure (EParam x) + _ -> pure en +{-# INLINE foldConsts #-} + +-- | Fold together all-but-one constant children of an ENAry multiset. +foldENary :: (ClassStore m, HasCallStack) => CostFun -> NOp -> IntMap Int -> [Consts] -> EGraphST m ENode +foldENary costFun op m infos = do + let xs = expandedList m + (consts, rest) = foldr step ([], []) (zip xs infos) + step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs) + step (x, _) (cs, rs) = (cs, x:rs) + if length consts >= 2 + then do + let folded = case op of + EAdd -> sum consts + EMul -> product consts + if isNaN folded || isInfinite folded + then pure (ENAry op m) + else do + cid <- add costFun (EConst folded) + pure (ENAry op (imFromList (cid : rest))) + else pure (ENAry op m) +{-# INLINE foldENary #-} + +-- | Fold together all-but-one constant children of an ENAry at insertion +-- time (e.g. 2+3+x becomes 5+x). Constants that are already folded +-- single subtrees are handled by 'calculateConsts' above; this handles the +-- flattened case where several constant terms land in one multiset. +foldConstants :: (ClassStore m, HasCallStack) => CostFun -> ENode -> EGraphST m ENode +foldConstants _ en@(ENAry _ m) | IntMap.size m < 2 = pure en +foldConstants costFun en@(ENAry op m) = do + let xs = expandedList m + infos <- mapM (fmap (_consts . _info) . getEClass) xs + let (consts, rest) = foldr step ([], []) (zip xs infos) + step (_, ConstVal v) (cs, rs) | not (isNaN v) && not (isInfinite v) = (v:cs, rs) + step (x, _) (cs, rs) = (cs, x:rs) + if length consts >= 2 + then do + let folded = case op of + EAdd -> sum consts + EMul -> product consts + if isNaN folded || isInfinite folded + then pure en + else do + cid <- add costFun (EConst folded) + pure (ENAry op (imFromList (cid : rest))) + else pure en +foldConstants _ en = pure en + +-- | rebuilds the e-graph after inserting or merging +-- e-classes +rebuild :: (ClassStore m, HasCallStack) => 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 :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> ENode -> EGraphST m () +repair costFun ecId enode = + do modify' $ over eNodeToEClass (HashMap.delete enode) + enode' <- canonize enode + ecId' <- canonical ecId + doExist <- lookupNode enode' + case doExist of + Just ecIdCanon -> do mergedId <- merge costFun ecIdCanon ecId' + insertNode enode' mergedId + addToDB enode' mergedId + Nothing -> do insertNode enode' ecId' + addToDB enode' ecId' +{-# INLINE repair #-} + +-- | repair the analysis of the e-class +-- considering the new added e-node +repairAnalysis :: (ClassStore m, HasCallStack) => 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 let bestChanged = _best (_info eclass) /= _best newData + modify' $ over (eDB . analysis) (_parents eclass <>) + . (if bestChanged && isJust (_fitness (_info eclass)) then over (eDB . refits) (IntSet.insert ecId') else id) + -- write through 'insertClass' so a paged store keeps the updated body + insertClass eclass' + _ <- modifyEClass costFun ecId' + pure () +{-# INLINE repairAnalysis #-} + +-- | merge to equivalent e-classes +merge :: (ClassStore m, HasCallStack) => 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 :: (ClassStore m, HasCallStack) => EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m EClassId + mergeClasses led ledC ledO sub subC subO = + do insertCanonical sub led -- persist/register the canonical merges + insertCanonical subO led + let newC = EClass led + (_eNodes ledC `Set.union` _eNodes subC) + (_parents ledC <> _parents subC) + (min (_height ledC) (_height subC)) + (joinData (_info ledC) (_info subC)) + forM_ (Set.toList (_eNodes subC)) $ \en -> insertNode en led + -- write the merged body through the class store (a paged store keeps the + -- authoritative page) and drop the absorbed class + insertClass newC + deleteClass sub + modify' $ over (eDB . worklist) (_parents subC <>) + when (_info newC /= _info ledC) + $ do let bestChanged = _best (_info newC) /= _best (_info ledC) + modify' $ over (eDB . analysis) (_parents ledC <>) + . (if bestChanged && isJust (_fitness (_info ledC)) then over (eDB . refits) (IntSet.insert led) else id) + when (_info newC /= _info subC) + $ modify' $ over (eDB . analysis) (_parents subC <>) + tracking <- gets (_trackDBs . _eDB) + when tracking $ updateDBs newC led ledC ledO sub subC subO + modifyEClass costFun led + modify' $ over (eDB . changed) (const True) + pure led + + getLeaderSub c1 c1O c2 c2O = + do ec1 <- getEClass c1 + ec2 <- getEClass c2 + let n1 = Set.size (_parents ec1) + n2 = Set.size (_parents ec2) + pure $ if n1 >= n2 + then (c1, ec1, c1O, c2, ec2, c2O) + else (c2, ec2, c2O, c1, ec1, c1O) + + updateDBs :: (ClassStore m, HasCallStack) => EClass -> EClassId -> EClass -> EClassId -> EClassId -> EClass -> EClassId -> EGraphST m () + updateDBs newC led ledC ledO sub subC subO = do + updateFitnessDB newC led ledC ledO sub subC subO + updateSizeDB newC led ledC ledO sub subC subO + + 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 = + case fitNew of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.insert led . IntSet.delete ledO . IntSet.delete sub . IntSet.delete subO) + Just fn -> do + when (fitNew /= fitLed) $ do + modify' $ case fitLed of + Nothing -> over (eDB . unevaluated) (IntSet.delete led . IntSet.delete ledO) + Just fl -> over (eDB . fitRangeDB) (removeRange led fl . removeRange ledO fl) + . over (eDB . sizeFitDB) (IntMap.adjust (removeRange ledO fl . removeRange led fl) szLed) + modify' $ over (eDB . fitRangeDB) (insertRange led fn) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange led fn) szNew . IntMap.insertWith RangeSet.union szNew RangeSet.empty) + modify' $ case fitSub of + Nothing -> over (eDB . unevaluated) (IntSet.delete sub . IntSet.delete subO) + Just fs -> over (eDB . fitRangeDB) (removeRange sub fs . removeRange subO fs) + . over (eDB . sizeFitDB) (IntMap.adjust (removeRange subO fs . removeRange sub fs) szSub) + where + fitNew = (_fitness . _info) newC + fitLed = (_fitness . _info) ledC + 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 :: (ClassStore m, HasCallStack) => CostFun -> EClassId -> EGraphST m EClassId +modifyEClass costFun ecId = + do ec <- getEClass ecId + case (_consts . _info) ec of + ConstVal x -> + do let en = EConst x + c <- calculateCost costFun en + let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } + maybeEid <- lookupNode en + -- write through 'insertClass' (a paged store keeps the authoritative page) + insertClass ec{ _eNodes = Set.singleton en, _info = infoEc } + when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.insert ecId) + case maybeEid of + Nothing -> pure ecId + Just eid' -> merge costFun eid' ecId + + ParamIx x -> + do let en = EParam x + c <- calculateCost costFun en + let infoEc = (_info ec){ _cost = c, _best = en, _consts = toConst en } + maybeEid <- lookupNode en + insertClass ec{ _eNodes = Set.insert en (_eNodes ec), _info = infoEc } + when (isJust $ _fitness $ _info ec) $ modify' $ over (eDB . refits) (IntSet.insert ecId) + case maybeEid of + Nothing -> pure ecId + Just eid' -> merge costFun eid' ecId + + _ -> pure ecId + + where + isTerm (EVar _) = True + isTerm (EConst _) = True + isTerm (EParam _) = True + isTerm _ = False + + toConst (EParam ix) = ParamIx ix + toConst (EConst x) = ConstVal x + toConst _ = NotConst + +-- * DB + +-- | `addToDB` adds an e-node and e-class id to the database +addToDB :: (ClassStore m, HasCallStack) => ENode -> EClassId -> EGraphST m () -- State DB () +addToDB enode' eid = do + eid' <- canonical eid + ec <- getEClass eid' + let isConst = _consts . _info $ ec + let enode = case isConst of + ConstVal x -> EConst x + ParamIx x -> EParam x + _ -> enode' + let ids = eid : eChildren enode -- we will add the e-class id and the children ids + op = eOpKey enode -- changes Bin op l r to Bin op () () so `op` as a single entry in the DB + trie <- gets (Map.lookup op . _patDB . _eDB) + case populate trie ids of -- populates the trie + Nothing -> pure () + Just t -> modify' $ over (eDB . patDB) (Map.insert op t) -- if something was created, insert back into the DB + recordNode enode eid -- register the node for the streaming matcher's source +{-# INLINE addToDB #-} + +-- | Populates an IntTrie with a sequence of e-class ids +populate :: Maybe IntTrie -> [EClassId] -> Maybe IntTrie +populate _ [] = Nothing +populate Nothing eids = foldr f Nothing eids + where + f :: EClassId -> Maybe IntTrie -> Maybe IntTrie + f eid (Just t) = Just $ IntTrie (IntMap.singleton eid t) + f eid Nothing = Just $ IntTrie (IntMap.singleton eid (IntTrie IntMap.empty)) +populate (Just tId) (eid:eids) = let nextTrie = IntMap.lookup eid (_trie tId) + val = fromMaybe (IntTrie IntMap.empty) $ populate nextTrie eids + in Just $ IntTrie (IntMap.insert eid val (_trie tId)) +{-# INLINE populate #-} + +canonizeMap :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m (Subst, ClassOrVar) +canonizeMap (subst, cv) = (,cv) <$> traverse g subst + where + g :: ClassStore m => SubVal -> EGraphST m SubVal + g (SVOne e2) = SVOne <$> canonOne e2 + g (SVMap m) = SVMap . IntMap.fromListWith (+) <$> mapM (\(e2, n) -> do + e2' <- canonOne (Left e2) + pure (getInt e2', n)) (IntMap.toList m) + canonOne :: ClassStore m => ClassOrVar -> EGraphST m ClassOrVar + canonOne (Left e2) = Left <$> canonical e2 + canonOne e2 = pure e2 +{-# INLINE canonizeMap #-} + +applyMatch :: (ClassStore m, HasCallStack) => CostFun -> Rule -> (Subst, 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 #-} + +-- | gets the e-node of the target of the rule +-- TODO: add consts and modify +classOfENode :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m (Maybe EClassId) +classOfENode costFun subst (VarPat c) = do let maybeEid = case Map.lookup (Right (fromEnum c)) subst of + Just (SVOne v) -> Just v + _ -> Nothing + case maybeEid of + Nothing -> pure Nothing + Just eid -> Just <$> canonical (getInt eid) +classOfENode costFun subst (Fixed (Const x)) = Just <$> add costFun (EConst x) +classOfENode costFun subst (Fixed target) = do newChildren <- mapM (classOfENode costFun subst) (getElems target) + case sequence newChildren of + Nothing -> pure Nothing + Just cs -> do let new_enode = replaceChildren cs target + cs' <- mapM canonical cs + areConsts <- mapM isConst cs' + if and areConsts + then do eid <- addTree costFun new_enode + rebuild costFun -- eid new_enode + pure (Just eid) + else do en <- toENode new_enode + en' <- canonize en + gets (HashMap.lookup en' . _eNodeToEClass) +classOfENode _ _ (NAry _ _) = error "classOfENode: n-ary pattern unsupported" +classOfENode _ _ Hole = error "classOfENode: Hole is only valid in MapP targets" +{-# INLINE classOfENode #-} + +-- | adds the target of the rule into the e-graph +reprPrat :: (ClassStore m, HasCallStack) => CostFun -> Subst -> Pattern -> EGraphST m EClassId +reprPrat costFun subst (VarPat c) = do + let k = Right (fromEnum c) + v <- case Map.lookup k subst of + Nothing -> error $ "REPRPRAT_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst) + Just (SVOne x) -> pure x + Just (SVMap _) -> error $ "REPRPRAT_REST_AS_SINGLE var=" <> show (fromEnum c) + canonical $ getInt v +reprPrat costFun subst (Fixed target) = do newChildren <- mapM (reprPrat costFun subst) (getElems target) + addTree costFun (replaceChildren newChildren target) +reprPrat costFun subst Hole = error "REPRPRAT_HOLE: Hole must be filled by MapP" +reprPrat costFun subst (NAry op ncs) = do + m <- IntMap.unionsWith (+) <$> mapM (childEidM costFun subst) ncs + case IntMap.toList m of + [] -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1))) + [(c, 1)] -> canonical c + _ -> do en <- mkENaryM op m + add costFun en +{-# INLINE reprPrat #-} + +-- | Adds a single child of an n-ary target pattern to the e-graph, returning +-- its contribution as a canonical multiset (so 'Rest' children carry their +-- 'IntMap' straight through without expansion). +childEidM :: (ClassStore m, HasCallStack) => CostFun -> Subst -> NChild -> EGraphST m (IntMap Int) +childEidM costFun subst (Ch p) = (`IntMap.singleton` 1) <$> reprPrat costFun subst p +childEidM costFun subst (Rest c) = restEidsM subst c +childEidM costFun subst (MapP p c) = do + es <- restEids subst c + ms <- forM es $ \e -> reprMapP costFun subst e p + pure (imFromList ms) +{-# INLINE childEidM #-} + +-- | The e-class ids bound to a rest variable, as a canonical multiset. +restEidsM :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m (IntMap Int) +restEidsM subst c = do + let k = Right (fromEnum c) + case Map.lookup k subst of + Just (SVMap m) -> pure m + Just (SVOne _) -> error $ "REPRPRAT_SINGLE_AS_REST var=" <> show (fromEnum c) + Nothing -> error $ "REPRPRAT_MISSING_REST var=" <> show (fromEnum c) +{-# INLINE restEidsM #-} + +-- | The e-class ids bound to a rest variable, expanded one entry per +-- occurrence (used by 'MapP', which needs to instantiate per child). +restEids :: (Monad m, HasCallStack) => Subst -> Char -> EGraphST m [EClassId] +restEids subst c = expandedList <$> restEidsM subst c +{-# INLINE restEids #-} + +-- | Build the target of a pattern where every `Hole` is filled with the +-- e-class `e` (used by 'MapP'). +reprMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> Pattern -> EGraphST m EClassId +reprMapP costFun subst e Hole = canonical e +reprMapP costFun subst e (VarPat c) = reprPrat costFun subst (VarPat c) +reprMapP costFun subst e (Fixed target) = do + newChildren <- mapM (reprMapP costFun subst e) (getElems target) + addTree costFun (replaceChildren newChildren target) +reprMapP costFun subst e (NAry op ncs) = do + m <- IntMap.unionsWith (+) <$> mapM (childMapP costFun subst e) ncs + case IntMap.toList m of + [] -> reprPrat costFun subst (Fixed (Const (if op == EAdd then 0 else 1))) + [(c, 1)] -> canonical c + _ -> do en <- mkENaryM op m + add costFun en +{-# INLINE reprMapP #-} + +-- | A single child of an n-ary pattern inside a 'MapP' function. +childMapP :: (ClassStore m, HasCallStack) => CostFun -> Subst -> EClassId -> NChild -> EGraphST m (IntMap Int) +childMapP costFun subst e (Ch p) = (`IntMap.singleton` 1) <$> reprMapP costFun subst e p +childMapP costFun subst e (Rest c) = restEidsM subst c +childMapP costFun subst e (MapP _ _) = error "nested MapP unsupported" +{-# INLINE childMapP #-} + +isValidHeight :: (ClassStore m, HasCallStack) => (Subst, ClassOrVar) -> EGraphST m Bool +isValidHeight match = do + h <- case snd match of + Left ec -> _height <$> getEClass ec + Right _ -> pure 0 + pure $ h < 15 +{-# INLINE isValidHeight #-} + +-- | returns `True` if the condition of a rule is valid for that match +isValidConditions :: ClassStore m => Condition -> (Subst, ClassOrVar) -> EGraphST m Bool +isValidConditions (Condition f) match = f (fst match) +{-# INLINE isValidConditions #-} + +-- * Tree to e-graph conversion and utility functions + +-- | Creates an e-graph from an expression tree +fromTree :: (ClassStore m, HasCallStack) => CostFun -> Fix SRTree -> EGraphST m EClassId +fromTree costFun = cataM sequence (addTree costFun) +{-# INLINE fromTree #-} + +-- | Builds an e-graph from multiple independent trees +fromTrees :: ClassStore 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 + + +getBestENode eid = (_best . _info) <$> getEClass eid +{-# INLINE getBestENode #-} + +-- | returns one expression rooted at e-class `eId` +-- TODO: avoid loopings +getExpressionFrom :: ClassStore m => EClassId -> EGraphST m (Fix SRTree) +getExpressionFrom eId' = do + nodes <- _eNodes <$> getEClass eId' + case Set.toList nodes of + (n:_) -> case n of + EVar ix -> pure $ Fix $ Var ix + EParam ix -> pure $ Fix $ Param ix + EConst x -> pure $ Fix $ Const x + EUni f t -> Fix . Uni f <$> getExpressionFrom t + EBin op l r -> Fix <$> (Bin op <$> getExpressionFrom l <*> getExpressionFrom r) + ENAry op xs -> naryTree op <$> mapM getExpressionFrom (expandedList xs) + [] -> error "getExpressionFrom: empty eclass" +{-# INLINE getExpressionFrom #-} + +-- | returns all expressions rooted at e-class `eId` +-- TODO: check for infinite list +getAllExpressionsFrom :: ClassStore m => EClassId -> EGraphST m [Fix SRTree] +getAllExpressionsFrom eId' = do + nodes <- Set.toList . _eNodes <$> getEClass eId' + go nodes + where + go [] = pure [] + go (n:ns) = do + t <- case n of + EVar ix -> pure [Fix $ Var ix] + EParam ix -> pure [Fix $ Param ix] + EConst x -> pure [Fix $ Const x] + EUni f t -> Prelude.map (Fix . Uni f) <$> getAllExpressionsFrom t + EBin op l r -> do l' <- getAllExpressionsFrom l + r' <- getAllExpressionsFrom r + pure $ [Fix $ Bin op li ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM getAllExpressionsFrom (expandedList xs) + pure [ naryTree op comb | comb <- sequence ts ] + ts <- go ns + pure (t ++ ts) +{-# INLINE getAllExpressionsFrom #-} + +getNExpressionsFrom :: ClassStore m => Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom n eId' = getNExpressionsFrom' n 15 eId' + +getNExpressionsFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [Fix SRTree] +getNExpressionsFrom' _ 0 _ = pure [] +getNExpressionsFrom' n d eId' = do + nodes <- Set.toList . _eNodes <$> getEClass eId' + (concat <$> go n d nodes) + where + isTerm (EVar _) = True + isTerm (EConst _) = True + isTerm (EParam _) = True + isTerm _ = False + toTree (EVar ix) = Fix $ Var ix + toTree (EConst x) = Fix $ Const x + toTree (EParam ix) = Fix $ Param ix + toTree _ = undefined + + go n' _ [] = pure [] + go n' 0 ts = pure [] + go n' d (node:ns) = do + tt <- case node of + EVar ix -> pure [Fix $ Var ix] + EParam ix -> pure [Fix $ Param ix] + EConst x -> pure [Fix $ Const x] + EUni f t -> Prelude.map (Fix . Uni f) <$> getNExpressionsFrom' n' (d-1) t + EBin op l r -> do l' <- getNExpressionsFrom' n' (d-1) l + r' <- getNExpressionsFrom' n' (d-1) r + pure $ Prelude.take n [Fix $ Bin op li ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM (getNExpressionsFrom' n' (d-1)) (expandedList xs) + pure $ Prelude.take n [ naryTree op comb | comb <- sequence ts ] + let n'' = n' - length tt + if n'' <= 0 + then pure [tt] + else do ts <- go n'' (d-1) ns + pure (tt:ts) + +getNEclassFrom :: ClassStore m => Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom n eid = getNEclassFrom' n 15 eid + +getNEclassFrom' :: ClassStore m => Int -> Int -> EClassId -> EGraphST m [[EClassId]] +getNEclassFrom' _ 0 _ = pure [] +getNEclassFrom' n d eId' = do + eId <- canonical eId' + nodes <- Set.toList . _eNodes <$> getEClass eId' + (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 + EBin op l r -> do l' <- getNEclassFrom' n' (d-1) l + r' <- getNEclassFrom' n' (d-1) r + pure $ Prelude.take n [li <> ri | li <- l', ri <- r'] + ENAry op xs -> do ts <- mapM (getNEclassFrom' n' (d-1)) xs + pure $ Prelude.take n [ concat comb | comb <- sequence ts ] + EUni f t -> getNEclassFrom' n' (d-1) t -- [[eid2:eid1]] + EVar ix -> pure [[]] + EConst x -> pure [[]] + EParam ix -> pure [[]] + pure tt + --let n'' = n' - length tt + --if n'' <= 0 + -- then pure [tt] + -- else do ts <- go n'' (d-1) ns + -- pure (tt:ts) + +getAllChildEClasses :: ClassStore m => EClassId -> EGraphST m [EClassId] +getAllChildEClasses eId' = do + eId <- canonical eId' + IntSet.toList <$> go [eId] IntSet.empty + + where + hasNoTerminal :: [ENode] -> Bool + hasNoTerminal = all (not . null . eChildren) + getNodes :: ClassStore m => EClassId -> EGraphST m [ENode] + getNodes n = Set.toList . _eNodes <$> getEClass n + + go :: ClassStore m => [Int] -> IntSet.IntSet -> EGraphST m IntSet.IntSet + go [] visited = pure visited + go queue visited = do + nodes <- concatMap eChildren . concat . filter hasNoTerminal <$> mapM getNodes queue + eids <- filter (\e -> e `IntSet.notMember` visited) <$> (mapM canonical nodes) + go eids (visited `IntSet.union` IntSet.fromList queue) + {- + 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 :: ClassStore m => EClassId -> EGraphST m [EClassId] +getAllChildBestEClasses eId' = do + IntSet.toList <$> go IntSet.empty eId' + where + go :: ClassStore m => IntSet.IntSet -> EClassId -> EGraphST m IntSet.IntSet + go acc n + | IntSet.member n acc = pure acc + | otherwise = do + let acc' = IntSet.insert n acc + node <- (_best . _info) <$> getEClass n + eids <- mapM canonical $ eChildren node + foldM go acc' eids + +getAllChildBestEClassesRep :: ClassStore m => EClassId -> EGraphST m [EClassId] +getAllChildBestEClassesRep eId' = do + go eId' + where + go :: ClassStore m => EClassId -> EGraphST m [EClassId] + go n = do node <- (_best . _info) <$> getEClass n + let hasTerminal = (null . eChildren) node + eids <- mapM canonical $ eChildren node + if hasTerminal + then pure [n] + else do eids' <- mapM go eids + pure (n : concat eids') + +-- | returns a random expression rooted at e-class `eId` +getRndExpressionFrom :: EClassId -> EGraphST (State StdGen) (Fix SRTree) +getRndExpressionFrom eId' = do + nodes <- Set.toList . _eNodes <$> getEClass eId' + n <- lift $ randomFrom nodes + case n of + EUni f t -> Fix . Uni f <$> getRndExpressionFrom t + EBin op l r -> Fix <$> (Bin op <$> getRndExpressionFrom l <*> getRndExpressionFrom r) + ENAry op xs -> naryTree op <$> mapM getRndExpressionFrom (expandedList xs) + EVar ix -> pure $ Fix $ Var ix + EConst x -> pure $ Fix $ Const x + EParam ix -> pure $ Fix $ Param ix + where + randomRange rng = state (randomR rng) + randomFrom xs = do n <- randomRange (0, length xs - 1) + pure $ xs !! n +{-# INLINE getRndExpressionFrom #-} + +cleanMaps :: ClassStore m => EGraphST m () +cleanMaps = do + hasStore <- gets (isJust . _classStore) + if hasStore + -- the paged store is authoritative for both node->class and canonical + -- lookups, so the bounded resident caches are simply reset (an O(n) rebuild + -- of an unbounded map would defeat the out-of-core goal). + then modify' $ \eg -> eg { _eNodeToEClass = HashMap.empty + , _canonicalMap = IntMap.empty + , _eClass = IntMap.empty } + else do + enode2eclass <- gets _eNodeToEClass + entries <- forM (HashMap.toList enode2eclass) $ \(k,v) -> do + k' <- canonize k + v' <- canonical v + pure (k',v') + let enode2eclass' = HashMap.fromList entries + eclassMap <- gets _eClass + entries' <- forM (IntMap.toList eclassMap) $ \(k,v) -> do + k' <- canonical k + pure $ if k==k' then (Just (k,v)) else Nothing + let eclassMap' = IntMap.fromList (catMaybes entries') + modify' $ \eg -> eg { _eNodeToEClass = enode2eclass' + , _eClass = eclassMap' } +{-# INLINE cleanMaps #-} 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,695 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +----------------------------------------------------------------------------- +-- | +-- 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 GHC.Stack (HasCallStack) +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IntMap +import Data.Map (Map) +import qualified Data.Map as Map +import Data.List (sortBy) +import Data.Maybe (fromMaybe) +import Data.Ord (comparing) +import Data.SRTree +import Data.HashSet (HashSet) +import qualified Data.HashSet as Set +import qualified Data.Set as RangeSet +import Data.String (IsString (..)) +import Data.SRTree.Recursion (cata) +import Text.Read (readMaybe) + + +-- A Pattern is either a fixed-point of a tree, an index to a pattern variable +-- (which matches anything), a hole (only used inside a 'MapP' target function), +-- or an n-ary Add/Mul pattern whose children are matched as a multiset. +data Pattern = Fixed (SRTree Pattern) | VarPat Char | Hole | NAry NOp [NChild] + deriving (Show, Eq, Ord) + +-- | A child of an n-ary pattern: a single child pattern ('Ch'), a rest +-- variable binding every remaining child of the node ('Rest'), or a +-- target-side map that splices one instantiation of a pattern (with its 'Hole' +-- filled) per child bound to a rest variable ('MapP'). +data NChild = Ch Pattern | Rest Char | MapP Pattern Char + deriving (Show, Eq, Ord) + +-- The instance for `IsString` for a `Pattern` is +-- valid only for a single letter char from a-zA-Z. +-- The patterns can be written as "x" + "y", for example, +-- 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 Add l r) = NAry EAdd [Ch l, Ch r] + alg (Bin Mul l r) = NAry EMul [Ch l, Ch r] + alg (Bin op l r) = Fixed $ Bin op l r + alg (Uni f t) = Fixed $ Uni f t +-- A rule is either a directional rule where pat1 can be replaced by pat2, a bidirectional rule +-- 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 predicate over a match's substitution that runs inside +-- the e-graph monad so it can fetch e-class data through 'ClassStore' (which +-- streams from a paged store when the graph is out-of-core). The quantification +-- over the monad is intentional: the same condition works for any 'ClassStore' +-- instance, including the IO-backed paged store. +newtype Condition = Condition (forall m. ClassStore m => Subst -> EGraphST m Bool) + +-- An Atom is composed of either an e-class id or pattern variable id +-- and the tree that generated that pattern. Left is e-class id and Right is a VarPat. +type ClassOrVar = Either EClassId Int +data Atom = Atom ClassOrVar (SRTree ClassOrVar) deriving Show + +-- | A substitution value: a single e-class (a matched pattern variable) or the +-- canonical multiset of e-class ids (a matched rest variable). +data SubVal = SVOne ClassOrVar | SVMap (IntMap Int) deriving Show + +-- | Substitution map produced by matching a pattern. +type Subst = Map ClassOrVar SubVal + +unFixPat :: Pattern -> SRTree Pattern +unFixPat (Fixed p) = p +unFixPat (VarPat _) = error "unFixPat: VarPat is not a fixed pattern" +unFixPat Hole = error "unFixPat: Hole is not a fixed pattern" +unFixPat (NAry _ _) = error "unFixPat: NAry is not a fixed pattern" +{-# INLINE unFixPat #-} + + +instance Num Pattern where + l + r = NAry EAdd [Ch l, Ch r] + {-# INLINE (+) #-} + l - r = NAry EAdd [Ch l, Ch (negate r)] + {-# INLINE (-) #-} + l * r = NAry EMul [Ch l, Ch r] + {-# INLINE (*) #-} + + abs = Fixed . Uni Abs + {-# INLINE abs #-} + + negate t = NAry EMul [Ch (Fixed (Const (-1))), Ch 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 = NAry EMul [Ch l, Ch (Fixed (Uni Recip 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. This is the pure +-- matcher (no seen-set) used by user pattern queries; saturation uses +-- 'matchSaturated'. +match :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)] +match src = if hasNAry src + then matchNAryWith Nothing src + else do + paged <- isPagedGraph + if paged + then matchStreamCached Nothing src + else matchCachedWith Nothing (compileToQuery src) +{-# INLINE match #-} + +-- | Non-n-ary matching. The match's root e-class anchors it the same way the +-- n-ary matcher anchors one match per trie root, so it shares the same cheap +-- persistent mark-on-attempt seen-set ('_seenMatches', keyed by rule source -> +-- root class id): already-processed roots are skipped so the per-rule budget +-- advances to new matches across the scheduler's ban/unban cycles. Keying by +-- the root (an @O(1)@ class id) avoids serializing every substitution, which +-- would dominate on rules whose @genericJoin@ yields many matches. 'Nothing' +-- disables the seen-set (pure queries). +matchCachedWith :: ClassStore m => Maybe String -> (Query, [ClassOrVar], ClassOrVar) -> EGraphST m [(Subst, ClassOrVar)] +matchCachedWith mSk (q, vars, root) = do + ss <- genericJoin q vars root + seenSk <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + let rootOf s = case Map.lookup root s of + Just (SVOne (Left eid)) -> eid + _ -> 0 + fresh = [ s | s <- ss + , Map.size s > 0 + , maybe True (\_ -> not (RangeSet.member (show (rootOf s)) seenSk)) mSk ] + taken = take ruleMatchBudget fresh + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.fromList (map (show . rootOf) taken))) + Nothing -> pure () + pure [ (s, case Map.lookup root s of + Nothing -> error $ "MATCHCACHED_MISSING root=" <> show (getInt root) <> " substSize=" <> show (Map.size s) + Just v -> fromSVOne v) + | s <- taken ] +{-# INLINE matchCachedWith #-} + +-- | Saturation matching: consults/marks the persistent seen-set so each rule's +-- per-iteration budget advances to genuinely new matches across ban/unban. +matchSaturated :: ClassStore m => Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchSaturated src = if hasNAry src + then matchNAryWith (Just (show src)) src + else do + paged <- isPagedGraph + if paged + then matchStreamCached (Just (show src)) src + else matchCachedWith (Just (show src)) (compileToQuery src) +{-# INLINE matchSaturated #-} + +-- | True if the pattern (or a nested child) is an n-ary Add/Mul pattern. +hasNAry :: Pattern -> Bool +hasNAry (NAry _ _) = True +hasNAry (Fixed t) = any hasNAry (getElems t) +hasNAry _ = False +{-# INLINE hasNAry #-} + +-- | The operator trie key of the top-level pattern. +opOf :: Pattern -> SRTree () +opOf (NAry EAdd _) = Bin Add () () +opOf (NAry EMul _) = Bin Mul () () +opOf (Fixed t) = getOperator t +opOf _ = error "opOf: pattern has no operator" +{-# INLINE opOf #-} + +-- | Matches an n-ary pattern against every root e-node of the operator trie. +-- A per-rule result budget ('ruleBudget') bounds the total number of matches +-- returned for one rule against one individual's nodes, and only the first +-- match per root e-class is kept, taming the O(k^2*m^2) backtracking of +-- Rest/Ch rules (e.g. factoring a common term out of a sum of products). +-- Keeping one match per root is sound: every returned match is genuine, and +-- the egraph merges the equivalent rewrites that further matches would apply, +-- so the rest of the root's matches are redundant work. +ruleBudget :: Int +ruleBudget = 64 + +-- | Cap on how many operator-trie root e-classes a single rule may visit per +-- match. 'ruleBudget' bounds the number of *results* returned, but a rule whose +-- matches are rare would otherwise still scan every root e-class in the trie +-- (every @+@/@*@ class in the graph), doing an expensive 'recursiveMatch' per +-- root -- which blows up on large graphs even though few matches result. +-- Capping root visits bounds the *search work* independently of the result +-- count. Sound: we only stop enumerating (fewer) genuine matches early. +ruleRootVisit :: Int +ruleRootVisit = 512 + +-- | Cap on how many matches a non-n-ary rule (the cached @genericJoin@ path) +-- may return per match. The n-ary matcher has 'ruleBudget'; give the cached +-- path a separate (larger) budget so a single rule cannot flood the iteration. +ruleMatchBudget :: Int +ruleMatchBudget = 1024 + +-- | Cap on how many operator-root e-classes the streaming cached matcher visits +-- per match, bounding the search work (and the page reads) independently of the +-- result count, exactly as 'ruleRootVisit' does for the n-ary matcher. +ruleMatchRootVisit :: Int +ruleMatchRootVisit = 2048 + +-- | Match an n-ary pattern against every root e-class of its operator trie. +-- +-- A persistent per-source set of already-attempted roots ('_seenMatches') lets +-- the matcher skip roots it has already tried, so the per-rule result/search +-- budgets keep advancing to *new* roots across the scheduler's ban/unban cycles +-- instead of re-enumerating the same head of the trie (which starves the tail). +-- Roots are marked as attempted on the first try ('mark-on-attempt'), whether or +-- not they yielded a match, so a match that fails 'applyMatch' conditions is not +-- re-attempted every cycle. +matchNAryWith :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchNAryWith mSk src = do + seen <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + -- skip already-attempted roots so the per-rule budget advances to new roots + -- across the scheduler's ban/unban cycles (matches the trie path's semantics). + let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ] + roots <- streamRoots (opOf src) ruleRootVisit exclude + go roots 0 0 [] + where + go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)] + go [] _ _ acc = pure (reverse acc) + go _ n _ acc | n >= ruleBudget = pure (reverse acc) + go (_ : _) _ r acc | r >= ruleRootVisit = pure (reverse acc) + go (eid : eids) n r acc = do + -- mark-on-attempt: remember this root as tried for this rule source + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid))) + Nothing -> pure () + substs <- recursiveMatch src eid Map.empty + let newMs = take 1 [ (s, Left eid) | s <- substs ] + go eids (n + length newMs) (r + 1) (foldr (:) acc newMs) +{-# INLINE matchNAryWith #-} + +-- | Streaming matcher for the cached (non-n-ary @genericJoin@) path on a paged +-- graph. Instead of enumerating candidates from the in-RAM @_patDB@ trie, it +-- streams the candidate root e-classes of the pattern's operator through +-- 'streamRoots' (bounded, skipping the already-attempted seen-set) and matches +-- each root incrementally with 'recursiveMatch' (which reads e-classes through +-- the paged store). This is the out-of-core analogue of 'matchCachedWith': the +-- resident/pure path keeps the optimized trie 'genericJoin', and only a paged +-- graph takes this route, so the matcher never builds an O(nodes) structure. +-- +-- 'ruleMatchBudget' bounds the results and 'ruleMatchRootVisit' bounds the root +-- visits; the persistent mark-on-attempt seen-set makes each rule's budgets +-- advance to new roots across the scheduler's ban/unban cycles. +matchStreamCached :: ClassStore m => Maybe String -> Pattern -> EGraphST m [(Subst, ClassOrVar)] +matchStreamCached mSk src = do + seen <- case mSk of + Nothing -> pure RangeSet.empty + Just sk -> gets (Map.findWithDefault RangeSet.empty sk . _seenMatches . _eDB) + let exclude = [ i | s <- RangeSet.toList seen, Just i <- [readMaybe s :: Maybe EClassId] ] + roots <- case opOfMay src of + Just op -> streamRoots op ruleMatchRootVisit exclude + Nothing -> pure [] + go roots 0 0 [] + where + go :: ClassStore m => [EClassId] -> Int -> Int -> [(Subst, ClassOrVar)] -> EGraphST m [(Subst, ClassOrVar)] + go [] _ _ acc = pure (reverse acc) + go _ n _ acc | n >= ruleMatchBudget = pure (reverse acc) + go (_ : _) _ r acc | r >= ruleMatchRootVisit = pure (reverse acc) + go (eid : eids) n r acc = do + case mSk of + Just sk -> modify' $ over (eDB . seenMatches) + (Map.insertWith RangeSet.union sk (RangeSet.singleton (show eid))) + Nothing -> pure () + substs <- recursiveMatch src eid Map.empty + let newMs = take (ruleMatchBudget - n) [ (s, Left eid) | s <- substs ] + go eids (n + length newMs) (r + 1) (foldr (:) acc newMs) +{-# INLINE matchStreamCached #-} + +-- | The operator trie key of the top-level pattern, or @Nothing@ for a pattern +-- with no operator (e.g. a bare variable), which the streaming matcher treats +-- as matching nothing. +opOfMay :: Pattern -> Maybe (SRTree ()) +opOfMay (NAry EAdd _) = Just (Bin Add () ()) +opOfMay (NAry EMul _) = Just (Bin Mul () ()) +opOfMay (Fixed t) = Just (getOperator t) +opOfMay _ = Nothing +{-# INLINE opOfMay #-} + +-- | Recursively match a pattern against the e-class `eid`, threading a +-- substitution map, returning every substitution that completes the match. +recursiveMatch :: ClassStore m => Pattern -> EClassId -> Subst -> EGraphST m [Subst] +recursiveMatch (VarPat c) eid subst = + pure (bindVar subst (Right (fromEnum c)) eid) +recursiveMatch Hole _ subst = pure [subst] +recursiveMatch (Fixed t) eid subst = matchFixed t eid subst +recursiveMatch (NAry op ncs) eid subst = matchNAryNode op ncs eid subst +{-# INLINE recursiveMatch #-} + +-- | Bind `v` to the e-class `eid`, enforcing that re-occurrences of `v` are +-- consistent. +bindVar :: Subst -> ClassOrVar -> EClassId -> [Subst] +bindVar subst v eid = + case Map.lookup v subst of + Just (SVOne e) | e == Left eid -> [subst] + Just _ -> [] + Nothing -> [Map.insert v (SVOne (Left eid)) subst] +{-# INLINE bindVar #-} + +-- | Match a fixed tree pattern against the e-nodes of the e-class `eid`, +-- returning every substitution that completes the match across all candidate +-- e-nodes. +matchFixed :: ClassStore m => SRTree Pattern -> EClassId -> Subst -> EGraphST m [Subst] +matchFixed t eid subst = do + ec <- getEClass eid + let cands = [n | n <- Set.toList (_eNodes ec), eOpKey n == getOperator t] + fmap concat $ forM cands $ \n -> matchChildren t subst n + where + matchChildren t s n = go (zip (getElems t) (enodeChildren n)) [s] + go [] ss = pure ss + go ((p, c) : ps) ss = do + ms <- concat <$> mapM (\s -> recursiveMatch p c s) ss + go ps ms +{-# INLINE matchFixed #-} + +-- | The child e-class ids of an e-node, in canonical (sorted for ENAry) order. +enodeChildren :: ENode -> [EClassId] +enodeChildren (EUni _ t) = [t] +enodeChildren (EBin _ l r) = [l, r] +enodeChildren (ENAry _ m) = expandedList m +enodeChildren _ = [] +{-# INLINE enodeChildren #-} + +-- | Match an n-ary pattern node against the e-class `eid`: it must contain an +-- ENAry node of the given op, whose children are matched as a multiset. Every +-- ENAry node in the class is tried. +matchNAryNode :: ClassStore m => NOp -> [NChild] -> EClassId -> Subst -> EGraphST m [Subst] +matchNAryNode op ncs eid subst = do + ec <- getEClass eid + let nodes = [m | ENAry op' m <- Set.toList (_eNodes ec), op' == op] + fmap concat $ forM nodes $ \m -> + matchNChildren ncs m subst +{-# INLINE matchNAryNode #-} + +-- | Match a sequence of n-ary children against a multiset of e-class ids. +-- Each 'Ch' consumes one matched child; a 'Rest' child consumes all remaining +-- children. Every multiset assignment is returned. Iterating over the distinct +-- child ids (the multiset's keys) is sound (duplicate copies only differ by +-- position, which 'decChild' already resolves) and avoids duplicate result +-- sets. +-- +-- A per-call result budget ('matchCap') caps the number of substitutions +-- returned, bounding the O(k^2*m^2) backtracking of Rest/Ch rules such as +-- factoring a common term out of a sum of products. Sound: each result is a +-- genuine match; we merely stop enumerating once the budget is exhausted. +matchCap :: Int +matchCap = 64 + +matchNChildren :: ClassStore m => [NChild] -> IntMap Int -> Subst -> EGraphST m [Subst] +matchNChildren ncs children subst = reverse <$> goB ncs children subst matchCap + where + goB :: ClassStore m => [NChild] -> IntMap Int -> Subst -> Int -> EGraphST m [Subst] + goB [] m s _ + | IntMap.null m = pure [s] + | otherwise = pure [] + goB (Rest c : ps) m s b = do + let v = Right (fromEnum c) + case Map.lookup v s of + Just _ -> pure [] -- rest variable already bound + Nothing -> goB ps IntMap.empty (Map.insert v (SVMap m) s) b + goB (Ch p : ps) m s b + | multiplicity m <= nCh ps = pure [] -- not enough children left + | otherwise = goC (IntMap.keys m) 0 [] + where + goC :: ClassStore m => [EClassId] -> Int -> [Subst] -> EGraphST m [Subst] + goC [] _ acc = pure acc + goC _ n acc | n >= b = pure acc + goC (c : cs) n acc = do + ms <- recursiveMatch p c s + goMs c ms cs n acc + goMs :: ClassStore m => EClassId -> [Subst] -> [EClassId] -> Int -> [Subst] -> EGraphST m [Subst] + goMs c [] cs n acc = goC cs n acc + goMs c (s' : ms) cs n acc + | n >= b = pure acc + | otherwise = do + r <- goB ps (decChild c m) s' (b - n) + let r' = take (b - n) r + n' = n + length r' + goMs c ms cs n' (foldr (:) acc r') + goB (MapP _ _ : _) _ _ _ = error "matchNChildren: MapP is only valid in targets" +{-# INLINE matchNChildren #-} + +-- | Total number of children (counting multiplicities) in a multiset. +multiplicity :: IntMap Int -> Int +multiplicity = IntMap.foldr' (+) 0 +{-# INLINE multiplicity #-} + +-- | Remove one occurrence of `c` from the multiset (decrementing its +-- multiplicity, or dropping the key entirely when it reaches zero). +decChild :: Int -> IntMap Int -> IntMap Int +decChild c = IntMap.update (\n -> if n > 1 then Just (n - 1) else Nothing) c +{-# INLINE decChild #-} + +-- | Number of 'Ch' patterns in a child pattern sequence (each consumes one +-- child, so at least this many children must remain). +nCh :: [NChild] -> Int +nCh = length . filter isCh + where + isCh (Ch _) = True + isCh _ = False +{-# INLINE nCh #-} + +-- | Unwrap a single-e-class substitution value. +fromSVOne :: SubVal -> ClassOrVar +fromSVOne (SVOne v) = v +fromSVOne (SVMap _) = error "fromSVOne: expected a single e-class" +{-# INLINE fromSVOne #-} + +-- | Returns a Query (list of atoms) of a pattern with pre-computed ordered vars +compileToQuery :: Pattern -> (Query, [ClassOrVar], ClassOrVar) +compileToQuery pat = (atoms, orderedVars atoms, root) + where (atoms, root) = evalState (processPat pat) 256 + -- creates the atoms of a pattern + processPat :: Pattern -> State Int (Query, ClassOrVar) + processPat (VarPat x) = pure ([], Right $ fromEnum x) + processPat (NAry _ _) = error "compileToQuery: n-ary pattern (use matchNAry instead)" + processPat Hole = error "compileToQuery: Hole is only valid in MapP targets" + processPat (Fixed pat) = do + -- get the next available var id and add as root + v <- get + let root = Right v + -- updates the next available id + modify (+1) + -- recursivelly process the children of the pattern + patChilds <- mapM processPat (getElems pat) + -- create an atom composed of the + -- root and the tree with the children + -- replaced by the childs roots + -- add the child atoms to the list + let atoms = concatMap fst patChilds + roots = map snd patChilds + atom = Atom root (replaceChildren roots pat) + atoms' = atom:atoms + pure (atoms', root) +{-# INLINE compileToQuery #-} + +-- get the value from the Either Int Int +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 :: (ClassStore m, HasCallStack) => Query -> [ClassOrVar] -> ClassOrVar -> EGraphST m [Subst] +genericJoin atoms vars root = go atoms vars + where + -- for each variable + -- for each possible e-class id for that variable + -- replace the var id with this e-class id, and + -- recurse to find the possible matches for the next atom + go :: ClassStore m => Query -> [ClassOrVar] -> EGraphST m [Subst] + go atoms [] = pure [Map.empty] -- | _ <- atoms] + go atoms (x:vars) = do cIds1 <- domainX x atoms root + maps <- forM cIds1 $ \classId -> do + map (Map.insert x (SVOne classId)) <$> go (updateVar x classId atoms) vars + pure (concat maps) +{-# INLINE genericJoin #-} + + + +-- | returns the e-class id for a certain variable that +-- matches the pattern described by the atoms +domainX :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [ClassOrVar] +domainX var atoms root = do + let atoms' = filter (elemOfAtom var) atoms -- :: [ClassOrVar] -- look only in the atoms with this var + map Left <$> intersectAtoms var atoms' root -- find the intersection of possible keys by each atom +{-# INLINE domainX #-} + +-- | returns all e-class id that can matches this sequence of atoms +intersectAtoms :: (ClassStore m, HasCallStack) => ClassOrVar -> Query -> ClassOrVar -> EGraphST m [EClassId] +intersectAtoms _ [] root = pure [] +intersectAtoms var (a:atoms) root = do + a0 <- toCanon =<< go a + Set.toList <$> (foldM (\acc atom -> do + res <- go atom + Set.intersection acc <$> toCanon res) a0 atoms) + where + 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 IntMap.empty trie (r:getElems t)) + Nothing -> 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 -> IntMap EClassId -> IntTrie -> [ClassOrVar] -> Maybe (HashSet EClassId) +intersectTries var xs trie [] = Just Set.empty +intersectTries var xs trie (i:ids) = + case i of + Left x -> case IntMap.lookup x (_trie trie) of + Just subtrie -> intersectTries var xs subtrie ids + Nothing -> Nothing + Right x -> if IntMap.member x xs + then case IntMap.lookup (xs IntMap.! x) (_trie trie) of + Just subtrie -> intersectTries var xs subtrie ids + Nothing -> Nothing + else if Right x == var + then if all (isDiffFrom x) ids + then Just $ Set.fromList (IntMap.keys (_trie trie)) + else Just $ IntMap.foldrWithKey (\k v acc -> + case intersectTries var (IntMap.insert x k xs) v ids of + Nothing -> acc + _ -> Set.insert k acc) Set.empty (_trie trie) + else Just $ IntMap.foldrWithKey (\k v acc -> + case intersectTries var (IntMap.insert x 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 +-- Ties are broken by putting an atom ROOT first. The root indexes the +-- operator trie directly, so matching it first replaces repeated whole-trie +-- folds (O(candidates x nodes)) with direct per-node trie descents. The old +-- tie-break (by id) put low-id pattern leaves before the high-id fresh root, +-- which made the root's domain include every operator node regardless of the +-- already-bound children (over-enumeration and O(n^2) folds). +-- Measured on the user config: 33s -> 19s (MT -N8), best loss unchanged. +orderedVars :: Query -> [ClassOrVar] +orderedVars atoms = sortBy (comparing key) $ RangeSet.toList $ RangeSet.fromList [a | atom <- atoms, a <- getIdsFrom atom, isRight a] + where + getIdsFrom (Atom r t) = r : getElems t + isRight (Right _) = True + isRight _ = False + + -- is the variable the ROOT of some atom (an index into the operator trie)? + isHeader v = any (\a -> case a of Atom r _ -> r == v) atoms + + varCost :: ClassOrVar -> Int + varCost var = foldr (\a acc -> if elemOfAtom var a then acc - 100 + atomLen a else acc) 0 atoms + + key :: ClassOrVar -> (Int, Int) + key v = (varCost v, if isHeader v then 0 else 1) + + atomLen (Atom _ t) = 1 + length (getElems t) +{-# INLINE orderedVars #-} diff --git a/src/Algorithm/EqSat/Egraph.hs b/src/Algorithm/EqSat/Egraph.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Egraph.hs @@ -0,0 +1,925 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE StrictData #-} +{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} +----------------------------------------------------------------------------- +-- | +-- 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_, when, foldM, void) +import Data.List ( intercalate, foldl' ) +import Control.Monad (forM) +import Control.Monad.State.Strict hiding ( get, put ) +import Control.Monad.IO.Class (MonadIO(..)) +import Data.Functor.Identity (Identity) +import GHC.Stack (HasCallStack) +import System.Random (StdGen) +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IntMap +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap +import Data.HashSet (HashSet) +import qualified Data.HashSet as Set +import Data.IntSet (IntSet) +import qualified Data.IntSet as IntSet +import qualified Data.Set as RangeSet +import Data.SRTree +import Data.SRTree.Eval +import Data.SRTree.Recursion (cata) +import Data.Hashable +import Data.Binary +import qualified Data.Binary as Bin +import qualified Data.Vector.Unboxed as VU +import Control.DeepSeq (NFData) + +import GHC.Generics + + +type EClassId = Int -- NOTE: DO NOT CHANGE THIS, this will break the use of IntMap and IntSet +type ClassIdMap = IntMap + +-- | N-ary operators represented as flattened multisets inside the e-graph. +-- Only Add and Mul are associative-commutative in this library; the remaining +-- ops (Sub, Div, Power, PowerAbs, AQ) stay binary and live in 'EBin'. +data NOp = EAdd | EMul deriving (Show, Eq, Ord, Enum, Generic, NFData) + +-- | The e-graph's node language. +-- +-- 'ENAry' stores Add/Mul as a canonical multiset of e-class ids: children are +-- path-compressed, keys sorted by canonical 'EClassId' (commutativity), and +-- nested same-op ENAry children are absorbed at insertion time +-- (associativity), so no commutativity/associativity rewrite rules are needed +-- for Add/Mul. The children are an 'IntMap' of e-class id to multiplicity. +data ENode + = EVar {-# UNPACK #-} !Int + | EParam {-# UNPACK #-} !Int + | EConst {-# UNPACK #-} !Double + | EUni Function EClassId + | EBin Op EClassId EClassId -- Sub | Div | Power | PowerAbs | AQ + | ENAry NOp (IntMap Int) -- canonical multiset: eclass -> multiplicity + deriving (Show, Eq, Generic, NFData) + +type EGraphST m a = StateT EGraph m a +type Cost = Int +type CostFun = SRTree Cost -> Cost +type ECache = IntMap.IntMap Target + +instance Hashable NOp where + hashWithSalt n EAdd = n `hashWithSalt` (0 :: Int) + hashWithSalt n EMul = n `hashWithSalt` (1 :: Int) + +instance Hashable ENode where + hashWithSalt n (EVar ix) = n `hashWithSalt` (0 :: Int) `hashWithSalt` ix + hashWithSalt n (EParam ix) = n `hashWithSalt` (1 :: Int) `hashWithSalt` ix + hashWithSalt n (EConst x) = n `hashWithSalt` (2 :: Int) `hashWithSalt` x + hashWithSalt n (EUni f t) = n `hashWithSalt` (3 :: Int) `hashWithSalt` (fromEnum f) `hashWithSalt` t + hashWithSalt n (EBin op l r) = n `hashWithSalt` (4 :: Int) `hashWithSalt` (fromEnum op) `hashWithSalt` l `hashWithSalt` r + hashWithSalt n (ENAry op m) = n `hashWithSalt` (5 :: Int) `hashWithSalt` op `hashWithSalt` m + +type RangeTree a = RangeSet.Set (a, EClassId) + +-- | Expand a canonical multiset back to the equivalent (multi-)set of child +-- e-class ids, one entry per occurrence. +expandedList :: IntMap Int -> [EClassId] +expandedList = concatMap (\(k, n) -> replicate n k) . IntMap.toAscList +{-# INLINE expandedList #-} + +-- | Build a canonical multiset from a list of child ids (duplicates allowed). +imFromList :: [EClassId] -> IntMap Int +imFromList = IntMap.fromListWith (+) . map (, 1) +{-# INLINE imFromList #-} + + + +insertRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a +insertRange eid x = RangeSet.insert (x, eid) +{-# INLINE insertRange #-} + +removeRange :: (Ord a, Show a) => EClassId -> a -> RangeTree a -> RangeTree a +removeRange eid x = RangeSet.delete (x, eid) +{-# INLINE removeRange #-} + + + + + +-- TODO: check this \/ +getWithinRange :: Ord a => a -> a -> RangeTree a -> [EClassId] +getWithinRange lb ub rt = + let (_, ge) = RangeSet.split (lb, minBound) rt + (inR, _) = RangeSet.split (ub, maxBound) ge + in map snd (RangeSet.toList inR) + +getSmallest :: Ord a => RangeTree a -> Maybe (a, EClassId) +getSmallest = RangeSet.lookupMin +{-# INLINE getSmallest #-} + +getGreatest :: Ord a => RangeTree a -> Maybe (a, EClassId) +getGreatest = RangeSet.lookupMax +{-# INLINE getGreatest #-} + +-- | Handle to an external, lazily paged e-class store (provided by the +-- storage layer, e.g. srtree-db's 'PageStore'). An 'EGraph' carries one when +-- e-classes are backed by a database; the IO actions fetch / persist / +-- evict a single e-class page. 'Nothing' keeps the classic fully-resident +-- behaviour. +data EClassPageStore = EClassPageStore + { cpsLookup :: EClassId -> IO (Maybe EClass) + , cpsInsert :: EClass -> IO () + , cpsDelete :: EClassId -> IO () + , cpsFlush :: IO () -- ^ write back all pending dirty pages + , cpsAll :: IO [EClass] -- ^ all e-classes currently in the store + , cpsKeys :: IO [EClassId] -- ^ all e-class ids currently in the store + , cpsStreamRoots :: SRTree () -> Int -> [EClassId] -> IO [EClassId] -- ^ bounded candidate roots for an operator, skipping an attempted set + , cpsRecordNode :: ENode -> EClassId -> IO () -- ^ register a newly-created node for write-back + , cpsNodeToClass :: ENode -> IO (Maybe EClassId) -- ^ content-address node -> class lookup (live) + , cpsCanonicalOf :: EClassId -> IO (Maybe EClassId) -- ^ e-class -> canonical representative (live) + , cpsRecordCanonical :: EClassId -> EClassId -> IO () -- ^ persist a canonical mapping (write-back) + , cpsBeginFrontier :: IO () -- ^ start a frontier re-saturation (restrict matcher to changed classes) + , cpsEndFrontier :: IO () -- ^ end it: clear the frontier (a pass re-saturated everything) + } + +data EGraph = EGraph { _canonicalMap :: ClassIdMap EClassId -- maps an e-class id to its canonical form + , _eNodeToEClass :: HashMap ENode EClassId -- maps an e-node to its e-class id + , _eClass :: ClassIdMap EClass -- maps an e-class id to its e-class data (resident cache) + , _eDB :: EGraphDB + , _classStore :: Maybe EClassPageStore -- optional lazily paged store for _eClass + } + +data EGraphDB = EDB { _worklist :: HashSet (EClassId, ENode) -- e-nodes and e-class schedule for analysis + , _analysis :: HashSet (EClassId, ENode) -- e-nodes and e-class that changed data + , _refits :: IntSet + , _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 + , _changed :: !Bool -- dirty flag: true if modified since last check + , _trackDBs :: !Bool -- maintain range DBs (False during pure simplify) + , _seenMatches :: Map String (RangeSet.Set String) -- persistent (rule source -> attempted match keys) + } deriving (Show, Generic) + +data EClass = EClass { _eClassId :: {-# UNPACK #-} !Int -- e-class id (maybe we don't need that here) + , _eNodes :: HashSet ENode -- set of e-nodes inside this e-class + , _parents :: HashSet (EClassId, ENode) -- parents (e-class, e-node)'s + , _height :: {-# UNPACK #-} !Int -- height + , _info :: EClassData -- data + } deriving (Show, Eq, Generic) + +data Consts = NotConst | ParamIx {-# UNPACK #-} !Int | ConstVal {-# UNPACK #-} !Double deriving (Show, Eq, Generic) +data Property = Positive | Negative | NonZero | Real deriving (Show, Eq, Generic) -- TODO: incorporate properties + +data EClassData = EData { _cost :: {-# UNPACK #-} !Cost + , _best :: ENode + , _consts :: Consts + , _fitness :: Maybe Double -- NOTE: this cannot be NaN + , _dl :: Maybe Double + , _theta :: [Target] + , _size :: {-# UNPACK #-} !Int + -- , _properties :: Property + -- TODO: include evaluation of expression from this e-class + } deriving (Show, Generic) + +-- * Serialization +instance Generic (EClassId, ENode) + +instance Binary NOp where + put EAdd = put (0 :: Word8) + put EMul = put (1 :: Word8) + + get = do t <- get :: Get Word8 + case t of + 0 -> pure EAdd + 1 -> pure EMul + +instance Binary ENode where + put (EVar ix) = put (0 :: Word8) >> put ix + put (EParam ix) = put (1 :: Word8) >> put ix + put (EConst x) = put (2 :: Word8) >> put x + put (EUni f t) = put (3 :: Word8) >> put (fromEnum f) >> put t + put (EBin op l r) = put (4 :: Word8) >> put (fromEnum op) >> put l >> put r + put (ENAry op m) = put (5 :: Word8) >> put op >> put (expandedList m) + + get = do t <- get :: Get Word8 + case t of + 0 -> EVar <$> get + 1 -> EParam <$> get + 2 -> EConst <$> get + 3 -> EUni <$> (toEnum <$> get) <*> get + 4 -> EBin <$> (toEnum <$> get) <*> get <*> get + 5 -> ENAry <$> get <*> (imFromList <$> get) + +instance Binary (SRTree ()) where + put (Var ix) = put (0 :: Word8) >> put ix + put (Param ix) = put (1 :: Word8) >> put ix + 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 k, Binary v, Hashable k, Eq k) => Binary (HashMap k v) where + put hm = put (HashMap.toList hm) + get = HashMap.fromList <$> get + +instance Binary Target where + put xs = put (VU.toList xs) + get = VU.fromList <$> get + +instance Binary IntTrie +instance Binary EClass +instance Binary Consts +instance Binary Property +instance Binary EClassData +-- Custom: keep `_trackDBs` out of the wire format so on-disk EGraphDB data +-- (written before the flag existed) decodes unchanged; it defaults to True. +instance Binary EGraphDB where + put (EDB w a r p f d s sf sdl u n c _ _) = + put w >> put a >> put r >> put p >> put f >> put d >> put s >> put sf >> put sdl >> put u >> put n >> put c + get = EDB <$> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> get <*> pure True <*> pure Map.empty +-- Custom: the wire format omits `_classStore` (a runtime handle to the paged +-- store, never serialized); it decodes to Nothing. +instance Binary EGraph where + put (EGraph c n e d _) = put c >> put n >> put e >> put d + get = EGraph <$> get <*> get <*> get <*> get <*> pure Nothing + +instance Eq EClassData where + EData c1 b1 cs1 ft1 dl1 _ s1 == EData c2 b2 cs2 ft2 dl2 _ s2 = c1==c2 && b1==b2 && cs1==cs2 && ft1==ft2 && dl1==dl2 && s1==s2 + +-- 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 +newtype IntTrie = IntTrie { _trie :: IntMap IntTrie } deriving (Generic) + +instance Show IntTrie where + show (IntTrie t) = "{" <> intercalate "," (map (\(k,v) -> show k <> " -> " <> show v) $ IntMap.toList t) <> "}" + +makeLenses ''EGraph +makeLenses ''EClass +makeLenses ''EClassData +makeLenses ''EGraphDB + +-- * Paged e-class access + +-- | A monad that can serve e-class data. +-- +-- The pure instances ('Identity', 'State StdGen') serve classes from the +-- resident @_eClass@ map; the 'MonadIO' instance consults the optional +-- 'EClassPageStore' when the graph carries one, falling back to the resident +-- map otherwise. All e-class read/write goes through these accessors, which +-- are the single choke point for a paged (out-of-core) e-graph. +class Monad m => ClassStore m where + lookupClass :: EClassId -> EGraphST m (Maybe EClass) + getClass :: HasCallStack => EClassId -> EGraphST m EClass + insertClass :: EClass -> EGraphST m () + deleteClass :: EClassId -> EGraphST m () + adjustClass :: EClassId -> (EClass -> EClass) -> EGraphST m () + -- | Enumerate every e-class (ids / values) in the graph. Paged graphs stream + -- from the store; resident graphs read the full @_eClass@ map. + allClasses :: EGraphST m [EClass] + allKeys :: EGraphST m [EClassId] + -- | Read/write a class directly from/to the backing store, bypassing the + -- resident LRU cache (and its O(n) 'trimResidentCache'). Bulk single-pass + -- traversals such as 'recalculateBestAllStream' must use these: routing every + -- one of ~n classes through 'lookupClass'/'insertClass' inserts each into the + -- resident map and calls 'trimResidentCache' (a full O(n) rebuild) after each + -- write, degenerating to O(n^2) and never terminating at scale. + readDirect :: EClassId -> EGraphST m (Maybe EClass) + writeDirect :: EClass -> EGraphST m () + allClasses = gets (IntMap.elems . _eClass) + allKeys = gets (IntMap.keys . _eClass) + readDirect = lookupClass + writeDirect = insertClass + -- | Enumerate (bounded) candidate e-class ids that contain a node of the + -- given operator, to drive the streaming matcher, skipping any ids in + -- @exclude@ (the already-attempted seen-set, so the per-rule budget advances + -- to new roots across scheduler cycles). The default reads the resident + -- @_patDB@ trie (the fully-in-RAM path); a paged graph streams the candidates + -- from its backing store instead, so the matcher never builds an O(nodes) + -- structure. + streamRoots :: SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId] + streamRoots = streamRootsFromDB + -- | Record a newly-created e-node (and its e-class) so a streaming matcher's + -- candidate source can see it. The default (fully resident graph) is a no-op: + -- the resident @_patDB@ is already updated by 'addToDB'. + recordNode :: ENode -> EClassId -> EGraphST m () + recordNode _ _ = pure () + -- | Content-address node -> class lookup. The default reads the resident + -- @_eNodeToEClass@ map (complete for a resident graph); a paged graph bounds + -- that map and falls back to the backing store on a miss. + lookupNode :: ENode -> EGraphST m (Maybe EClassId) + lookupNode en = gets (HashMap.lookup en . _eNodeToEClass) + -- | Record a node -> class mapping. The default keeps the resident (full) + -- map; a paged graph bounds it (evicting, since the store is authoritative). + insertNode :: ENode -> EClassId -> EGraphST m () + insertNode en eid = modify' $ over eNodeToEClass (HashMap.insert en eid) + -- | Record a canonical mapping (e-class -> representative), persisting it on a + -- paged graph so the store-backed canonical lookup sees merges/new classes. + insertCanonical :: EClassId -> EClassId -> EGraphST m () + insertCanonical eid canon = modify' $ over canonicalMap (IntMap.insert eid canon) + -- | The canonical representative of an e-class, or @Nothing@ when unknown. The + -- default reads the resident @_canonicalMap@; a paged graph bounds it and + -- falls back to the store. + canonicalOf :: EClassId -> EGraphST m (Maybe EClassId) + canonicalOf eid = gets (IntMap.lookup eid . _canonicalMap) + +-- | Default candidate-root enumeration from the resident @_patDB@ trie, capped +-- at @budget@ after skipping @exclude@ (used by the pure instances and as the +-- no-store fallback for a @MonadIO@ graph). +streamRootsFromDB :: Monad m => SRTree () -> Int -> [EClassId] -> EGraphST m [EClassId] +streamRootsFromDB op budget exclude = do + db <- gets (_patDB . _eDB) + let ex = IntSet.fromList exclude + case Map.lookup op db of + Nothing -> pure [] + Just trie -> pure (take budget [ e | e <- IntMap.keys (_trie trie), not (IntSet.member e ex) ]) +{-# INLINE streamRootsFromDB #-} + +-- | Whether the graph is backed by a lazily paged e-class store. Streaming +-- matchers dispatch on this: a paged graph enumerates candidates from the +-- backing store (bounded memory), a resident graph from @_patDB@. +isPagedGraph :: Monad m => EGraphST m Bool +isPagedGraph = gets (maybe False (const True) . _classStore) +{-# INLINE isPagedGraph #-} + +-- Resident-map implementations (used by every pure monad) ------------------ + +pureLookupClass :: Monad m => EClassId -> EGraphST m (Maybe EClass) +pureLookupClass cid = gets (IntMap.lookup cid . _eClass) +{-# INLINE pureLookupClass #-} + +pureGetClass :: (Monad m, HasCallStack) => EClassId -> EGraphST m EClass +pureGetClass cid = do + m <- pureLookupClass cid + case m of + Just ec -> pure ec + Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid +{-# INLINE pureGetClass #-} + +pureInsertClass :: Monad m => EClass -> EGraphST m () +pureInsertClass ec = modify' $ over eClass (IntMap.insert (_eClassId ec) ec) +{-# INLINE pureInsertClass #-} + +pureDeleteClass :: Monad m => EClassId -> EGraphST m () +pureDeleteClass cid = modify' $ over eClass (IntMap.delete cid) +{-# INLINE pureDeleteClass #-} + +pureAdjustClass :: Monad m => EClassId -> (EClass -> EClass) -> EGraphST m () +pureAdjustClass cid f = modify' $ over eClass (IntMap.adjust f cid) +{-# INLINE pureAdjustClass #-} + +-- | Maximum number of e-classes kept in the resident @_eClass@ cache when the +-- graph is backed by a paged store. When exceeded, the largest-id classes are +-- retained and the rest evicted from the resident map. The store remains +-- authoritative (and Little-data reads fall back to it), so eviction only +-- bounds memory, never correctness. +residentClassCap :: Int +residentClassCap = 50000 + +-- | Trim the resident @_eClass@ cache to at most 'residentClassCap' entries +-- by keeping the largest ids. No-op for graphs without a paged store (their +-- resident map must stay complete for the pure instances). Halving on 2x keeps +-- steady churn from triggering an O(n) rebuild on every insert. +trimResidentCache :: Monad m => EGraphST m () +trimResidentCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _eClass eg + n = IntMap.size m + in if n <= 2 * residentClassCap + then eg + else over eClass (const (IntMap.fromList (Prelude.drop (n - residentClassCap) (IntMap.toAscList m)))) eg + +-- | Bound on the resident @_eNodeToEClass@ cache on a paged graph. Beyond the +-- cap (checked at 2x, halved back to cap) the map is pruned; the backing store +-- is authoritative, so eviction only trades a little dedup accuracy for bounded +-- memory, never correctness. +nodeCacheCap :: Int +nodeCacheCap = 100000 + +-- | Bound on the resident @_canonicalMap@ cache on a paged graph (same +-- halve-on-2x policy; evicted entries are re-read from the store). +canonicalCacheCap :: Int +canonicalCacheCap = 100000 +{-# INLINE nodeCacheCap #-} +{-# INLINE canonicalCacheCap #-} + +trimNodeCache :: Monad m => EGraphST m () +trimNodeCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _eNodeToEClass eg + n = HashMap.size m + in if n <= 2 * nodeCacheCap + then eg + else over eNodeToEClass (const (HashMap.fromList (Prelude.take nodeCacheCap (HashMap.toList m)))) eg +{-# INLINE trimNodeCache #-} + +trimCanonicalCache :: Monad m => EGraphST m () +trimCanonicalCache = modify' $ \eg -> + case _classStore eg of + Nothing -> eg + Just _ -> + let m = _canonicalMap eg + n = IntMap.size m + in if n <= 2 * canonicalCacheCap + then eg + else over canonicalMap (const (IntMap.fromList (Prelude.take canonicalCacheCap (IntMap.toAscList m)))) eg +{-# INLINE trimCanonicalCache #-} + +instance ClassStore Identity where + lookupClass = pureLookupClass + getClass = pureGetClass + insertClass = pureInsertClass + deleteClass = pureDeleteClass + adjustClass = pureAdjustClass + +instance ClassStore (State StdGen) where + lookupClass = pureLookupClass + getClass = pureGetClass + insertClass = pureInsertClass + deleteClass = pureDeleteClass + adjustClass = pureAdjustClass + +-- Any monad that can run IO is potentially paged: the graph's optional +-- store, when present, is authoritative; otherwise classes come from the +-- resident map. +instance {-# OVERLAPPABLE #-} (Monad m, MonadIO m) => ClassStore m where + -- The resident map is kept in sync by 'insertClass'/'deleteClass', so it is + -- consulted first: repeated reads never touch the store, and a class that + -- was evicted from the store's LRU while still dirty is never served stale. + lookupClass cid = do + eg <- gets id + case IntMap.lookup cid (_eClass eg) of + Just ec -> pure (Just ec) + Nothing -> case _classStore eg of + Nothing -> pure Nothing + Just h -> liftIO (cpsLookup h cid) + getClass cid = do + eg <- gets id + case IntMap.lookup cid (_eClass eg) of + Just ec -> pure ec + Nothing -> case _classStore eg of + Nothing -> pureGetClass cid + Just h -> do + m <- liftIO (cpsLookup h cid) + case m of + Just ec -> do + modify' (over eClass (IntMap.insert cid ec)) + trimResidentCache + pure ec + Nothing -> error $ "GETECLASS_MISSING eid=" <> show cid + insertClass ec = do + eg <- gets id + case _classStore eg of + Nothing -> pureInsertClass ec + Just h -> do liftIO (cpsInsert h ec) + pureInsertClass ec + trimResidentCache + deleteClass cid = do + eg <- gets id + case _classStore eg of + Nothing -> pureDeleteClass cid + Just h -> do liftIO (cpsDelete h cid) + pureDeleteClass cid + adjustClass cid f = do + eg <- gets id + case _classStore eg of + Nothing -> pureAdjustClass cid f + Just _ -> do + m <- lookupClass cid + case m of + Nothing -> pure () + Just ec -> insertClass (f ec) + allClasses = do + eg <- gets id + case _classStore eg of + Nothing -> pure (IntMap.elems (_eClass eg)) + Just h -> liftIO (cpsAll h) + allKeys = do + eg <- gets id + case _classStore eg of + Nothing -> pure (IntMap.keys (_eClass eg)) + Just h -> liftIO (cpsKeys h) + -- Bypass the resident cache entirely: read the page straight from the store + -- and never insert into the (bounded) resident map, so a bulk traversal over + -- every class stays O(n) instead of O(n^2). + readDirect cid = do + eg <- gets id + case _classStore eg of + Nothing -> pureLookupClass cid + Just h -> liftIO (cpsLookup h cid) + writeDirect ec = do + eg <- gets id + case _classStore eg of + Nothing -> pureInsertClass ec + Just h -> liftIO (cpsInsert h ec) + streamRoots op budget exclude = do + eg <- gets id + case _classStore eg of + Nothing -> streamRootsFromDB op budget exclude + Just h -> liftIO (cpsStreamRoots h op budget exclude) + recordNode en eid = do + eg <- gets id + case _classStore eg of + Nothing -> pure () + Just h -> liftIO (cpsRecordNode h en eid) + lookupNode en = do + eg <- gets id + case _classStore eg of + Nothing -> gets (HashMap.lookup en . _eNodeToEClass) + Just h -> do + m <- gets (HashMap.lookup en . _eNodeToEClass) + case m of + Just eid -> pure (Just eid) + Nothing -> do + r <- liftIO (cpsNodeToClass h en) + case r of + Just eid -> do insertNode en eid + pure (Just eid) + Nothing -> pure Nothing + insertNode en eid = do + eg <- gets id + case _classStore eg of + Nothing -> modify' $ over eNodeToEClass (HashMap.insert en eid) + Just _ -> do modify' $ over eNodeToEClass (HashMap.insert en eid) + trimNodeCache + insertCanonical eid canon = do + eg <- gets id + case _classStore eg of + Nothing -> modify' $ over canonicalMap (IntMap.insert eid canon) + Just h -> do modify' $ over canonicalMap (IntMap.insert eid canon) + trimCanonicalCache + liftIO (cpsRecordCanonical h eid canon) + canonicalOf eid = do + eg <- gets id + case _classStore eg of + Nothing -> gets (IntMap.lookup eid . _canonicalMap) + Just h -> do + m <- gets (IntMap.lookup eid . _canonicalMap) + case m of + Just c -> pure (Just c) + Nothing -> do + r <- liftIO (cpsCanonicalOf h eid) + case r of + Just c -> do modify' $ over canonicalMap (IntMap.insert eid c) + trimCanonicalCache + pure (Just c) + Nothing -> pure Nothing + +-- * E-Graph basic supporting functions + +-- | returns an empty e-graph +emptyGraph :: EGraph +emptyGraph = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDB Nothing +{-# INLINE emptyGraph #-} + +-- | returns an empty e-graph DB +emptyDB :: EGraphDB +emptyDB = EDB + Set.empty + Set.empty + IntSet.empty + Map.empty + RangeSet.empty + RangeSet.empty + IntMap.empty + IntMap.empty + IntMap.empty + IntSet.empty + 0 + False + True + Map.empty +{-# INLINE emptyDB #-} + +-- | like 'emptyDB' but skips range-DB maintenance (pure simplify mode) +emptyDBNoTrack :: EGraphDB +emptyDBNoTrack = emptyDB{ _trackDBs = False } +{-# INLINE emptyDBNoTrack #-} + +-- | an empty e-graph that skips range-DB maintenance (pure simplify mode) +emptyGraphNoTrack :: EGraph +emptyGraphNoTrack = EGraph IntMap.empty HashMap.empty IntMap.empty emptyDBNoTrack Nothing +{-# INLINE emptyGraphNoTrack #-} + +-- | Creates a new e-class from an e-class id, a new e-node, +-- and the info of this e-class +createEClass :: EClassId -> ENode -> EClassData -> Int -> EClass +createEClass cId enode' info h = EClass cId (Set.singleton enode') Set.empty h info +{-# INLINE createEClass #-} + +-- | gets the canonical id of an e-class with full path compression +canonical :: ClassStore m => EClassId -> EGraphST m EClassId +canonical eclassId = do + mStep <- canonicalOf eclassId + case mStep of + Nothing -> canonError eclassId + Just oneStep + | oneStep == eclassId -> pure eclassId + | otherwise -> do + (root, chain) <- walk [eclassId] oneStep + -- compress the chain in the resident cache (cache-only: the store + -- keeps the authoritative semantic mappings recorded at insert + -- time, so eviction just loses the shortcut, never correctness). + modify' $ \eg -> eg{ _canonicalMap = + foldl' (\m' k -> IntMap.insert k root m') (_canonicalMap eg) chain } + pure root + where + walk :: ClassStore m => [EClassId] -> EClassId -> EGraphST m (EClassId, [EClassId]) + walk chain ecId = do + mNext <- canonicalOf ecId + case mNext of + Nothing -> canonError ecId + Just n + | n == ecId -> pure (ecId, chain) + | otherwise -> walk (ecId : chain) n + + canonError :: ClassStore m => EClassId -> EGraphST m a + canonError eid = do + m <- gets _canonicalMap + error $ "CANON_MISSING eid=" <> show eid <> " mapSize=" <> show (IntMap.size m) +{-# INLINE canonical #-} + +-- | canonize the e-node children +canonize :: (ClassStore m, HasCallStack) => ENode -> EGraphST m ENode +canonize (EVar ix) = pure (EVar ix) +canonize (EParam ix) = pure (EParam ix) +canonize (EConst x) = pure (EConst x) +canonize (EUni f t) = EUni f <$> canonical t +canonize (EBin op l r) = EBin op <$> canonical l <*> canonical r +-- re-map children to their canonical ids; IntMap keeps keys sorted, so +-- commutativity is structural, no rewrite rule required. +canonize (ENAry op m) = do + m' <- IntMap.fromListWith (+) <$> forM (IntMap.toList m) (\(c, n) -> do + c' <- canonical c + pure (c', n)) + pure (ENAry op m') +{-# INLINE canonize #-} + +-- | The children e-class ids of an e-node. +eChildren :: ENode -> [EClassId] +eChildren (EVar _) = [] +eChildren (EParam _) = [] +eChildren (EConst _) = [] +eChildren (EUni _ t) = [t] +eChildren (EBin _ l r) = [l, r] +eChildren (ENAry _ m) = expandedList m +{-# INLINE eChildren #-} + +toOp :: NOp -> Op +toOp EAdd = Add +toOp EMul = Mul +{-# INLINE toOp #-} + +-- | Operator shape key used to index the pattern database. ENAry maps back to +-- the corresponding binary operator shape so existing (binary) Add/Mul +-- patterns address the same trie. +eOpKey :: ENode -> SRTree () +eOpKey (EVar ix) = Var ix +eOpKey (EParam ix) = Param ix +eOpKey (EConst x) = Const x +eOpKey (EUni f _) = Uni f () +eOpKey (EBin op _ _) = Bin op () () +eOpKey (ENAry EAdd _) = Bin Add () () +eOpKey (ENAry EMul _) = Bin Mul () () +{-# INLINE eOpKey #-} + +-- | Convert an e-node (children still as e-class ids) into the equivalent +-- binary SRTree shape. NOTE: only called on non-ENary nodes; flattened +-- ENAry nodes have no binary skeleton (see 'naryTree' / the explicit ENAry +-- cases in the analyses). +fromENode :: ENode -> SRTree EClassId +fromENode (EVar ix) = Var ix +fromENode (EParam ix) = Param ix +fromENode (EConst x) = Const x +fromENode (EUni f t) = Uni f t +fromENode (EBin op l r) = Bin op l r +fromENode (ENAry _ _) = error "fromENode: ENAry has no binary skeleton" +{-# INLINE fromENode #-} + +-- | Right-fold a list of e-class child expressions into a binary Fix SRTree +-- for a flattened ENAry multiset (extraction). +naryTree :: NOp -> [Fix SRTree] -> Fix SRTree +naryTree op ts = normalizeSubDiv (foldr1 (\a b -> Fix (Bin (toOp op) a b)) ts) +{-# INLINE naryTree #-} + +-- | Re-render the internal negate/recip canonical forms back as Sub/Div so +-- extraction output keeps the familiar shape: `x + (-1)*y` -> `x - y`, +-- `x + (-3)` -> `x - 3` and `x * recip y` -> `x / y`. Sub and Div never +-- appear as e-nodes; they only reappear here during reconstruction. +normalizeSubDiv :: Fix SRTree -> Fix SRTree +normalizeSubDiv = cata alg + where + alg :: SRTree (Fix SRTree) -> Fix SRTree + alg (Bin Add l r) = case pick l r of + Just (pos, neg) -> Fix (Bin Sub pos neg) + Nothing -> Fix (Bin Add l r) + where + pick a b = case negated a of + Just t -> Just (b, t) + Nothing -> case negated b of + Just t -> Just (a, t) + Nothing -> Nothing + negated (Fix (Bin Mul (Fix (Const c)) t)) | c == -1 = Just t + negated (Fix (Bin Mul t (Fix (Const c)))) | c == -1 = Just t + negated (Fix (Const c)) | c < 0 = Just (Fix (Const (-c))) + negated _ = Nothing + alg (Bin Mul l r) = case pick l r of + Just (num, den) -> Fix (Bin Div num den) + Nothing -> Fix (Bin Mul l r) + where + pick a b = case a of + Fix (Uni Recip t) -> Just (b, t) + _ -> case b of + Fix (Uni Recip t) -> Just (a, t) + _ -> Nothing + alg t = Fix t + +-- | Convert a binary SRTree (children as e-class ids) into an e-node, +-- flattening Add/Mul into canonical ENAry multisets. +toENode :: (ClassStore m, HasCallStack) => SRTree EClassId -> EGraphST m ENode +toENode (Var ix) = pure (EVar ix) +toENode (Param ix) = pure (EParam ix) +toENode (Const x) = pure (EConst x) +toENode (Uni f t) = EUni f <$> canonical t +toENode (Bin Add l r) = mkENary EAdd [l, r] +toENode (Bin Mul l r) = mkENary EMul [l, r] +toENode (Bin op l r) = EBin op <$> canonical l <*> canonical r +toENode n = error $ "toENode: unsupported node " <> show n +{-# INLINE toENode #-} + +-- | Build a canonical ENAry from child ids: canonicalize children, absorb +-- nested same-op ENAry children (associativity), sort by key (commutativity). +mkENary :: (ClassStore m, HasCallStack) => NOp -> [EClassId] -> EGraphST m ENode +mkENary op cids = mkENaryM op (imFromList cids) + +-- | Build a canonical ENAry from a canonical multiset of child ids. +mkENaryM :: (ClassStore m, HasCallStack) => NOp -> IntMap Int -> EGraphST m ENode +mkENaryM op m = do + flat <- IntMap.unionsWith (+) <$> mapM (expandM op) (IntMap.toList m) + pure (ENAry op flat) + +-- | If the e-class of `cid` holds exactly one e-node and that node is an ENAry +-- of the same op, return its children scaled by `n` (flattening `n` +-- occurrences); otherwise return `n` copies of `cid`. Flattening is only sound +-- through a class with a single node: if the class were merged with other +-- nodes (e.g. `{Add[a,b], Mul[x,c]}`) flattening would silently pick one +-- representative and change the meaning of the term. +expandM :: (ClassStore m, HasCallStack) => NOp -> (EClassId, Int) -> EGraphST m (IntMap Int) +expandM op (cid, n) = do + ec <- getEClass cid + case Set.toList (_eNodes ec) of + [ENAry op' m'] | op' == op -> pure (IntMap.map (* n) m') + _ -> pure (IntMap.singleton cid n) + +-- | Reconstruct a binary Fix SRTree from an e-node, right-folding ENAry +-- into nested Bin Add/Mul. +enodeToTree :: (ClassStore m, HasCallStack) => ENode -> EGraphST m (Fix SRTree) +enodeToTree (EVar ix) = pure (Fix (Var ix)) +enodeToTree (EParam ix) = pure (Fix (Param ix)) +enodeToTree (EConst x) = pure (Fix (Const x)) +enodeToTree (EUni f t) = Fix . Uni f <$> getBestExpr t +enodeToTree (EBin op l r) = do + tl <- getBestExpr l + tr <- getBestExpr r + pure (Fix (Bin op tl tr)) +enodeToTree (ENAry op m) = do + ts <- mapM getBestExpr (expandedList m) + pure (naryTree op ts) +{-# INLINE enodeToTree #-} + +-- | gets an e-class with id `c` (auto-canonizes) +getEClass :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m EClass +getEClass c = do c' <- canonical c; getClass c' +{-# INLINE getEClass #-} + +-- | gets the best expression given the default cost function. Cycle-safe and +-- budgeted: see 'getBestExprBounded'. +getBestExpr :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree) +getBestExpr eid = getBestExprBounded eid + +-- | Like 'getBestExpr' but terminates on pathological graphs: a visited set +-- stops the expansion from re-entering an already-expanded class (a @_best@ +-- cycle arising from supersaturation/merges), and a node budget caps the total +-- expanded size (so an exponentially-shared DAG is truncated rather than +-- exploded). Both guards substitute a @Var 0@ placeholder for the part that +-- would otherwise blow up. On well-formed acyclic graphs with small bests +-- neither guard triggers, so the result is identical to the unbounded version. +-- This keeps out-of-core extraction (e.g. 'dbTop') bounded in memory. +getBestExprBounded :: (ClassStore m, HasCallStack) => EClassId -> EGraphST m (Fix SRTree) +getBestExprBounded eid = fst <$> expand Set.empty 0 eid + where + budget :: Int + budget = 200 + -- expand returns the tree and the running count of expanded nodes, so the + -- budget bounds the TOTAL size (not just the depth): an exponentially-shared + -- DAG is truncated instead of exploded. A revisited (cyclic) class or a + -- full budget yields a @Var 0@ placeholder. + expand :: ClassStore m => HashSet EClassId -> Int -> EClassId -> EGraphST m (Fix SRTree, Int) + expand _ n _ | n >= budget = pure (Fix (Var 0), n) + expand seen n eid + | Set.member eid seen = pure (Fix (Var 0), n) + | otherwise = do + best <- (_best . _info) <$> getEClass eid + let seen' = Set.insert eid seen + n0 = n + 1 + case best of + EVar ix -> pure (Fix (Var ix), n0) + EParam ix -> pure (Fix (Param ix), n0) + EConst x -> pure (Fix (Const x), n0) + EUni f t -> do (tt, n1) <- expand seen' n0 t + pure (Fix (Uni f tt), n1) + EBin op l r -> do + (tl, n1) <- expand seen' n0 l + (tr, n2) <- expand seen' n1 r + pure (Fix (Bin op tl tr), n2) + ENAry op m -> do + (xs, nEnd) <- goNary seen' n0 (IntMap.toAscList m) [] + pure (if null xs then (Fix (Var 0), nEnd) else (naryTree op xs, nEnd)) + -- build the ENAry children from the multiset WITHOUT materialising the + -- expanded multiplicity list: an enormous count (a pathological supersaturated + -- class) is capped per-child and by the total budget, so each copy counts + -- toward the budget and no giant list is ever allocated. + goNary seen n es acc + | n >= budget = pure (reverse acc, n) + | otherwise = case es of + [] -> pure (reverse acc, n) + ((c, cnt) : rest) -> do + (t, n1) <- expand seen n c + let take = min cnt (budget - n1 + 1) + n2 = n1 + (take - 1) + acc' = Prelude.replicate take t ++ acc + goNary seen n2 rest acc' + +-- | Creates a singleton trie from an e-class id +trie :: EClassId -> IntMap IntTrie -> IntTrie +trie eid = IntTrie +{-# INLINE trie #-} + +-- | Check whether an e-class is a constant value +isConst :: ClassStore m => EClassId -> EGraphST m Bool +isConst eid = do ec <- getEClass eid + case (_consts . _info) ec of + ConstVal _ -> pure True + _ -> pure False +{-# INLINE isConst #-} + +getFitness :: ClassStore m => EClassId -> EGraphST m (Maybe Double) +getFitness c = (_fitness . _info) <$> getEClass c +{-# INLINE getFitness #-} +getTheta :: ClassStore m => EClassId -> EGraphST m ([Target]) +getTheta c = (_theta . _info) <$> getEClass c +{-# INLINE getTheta #-} +getSize :: ClassStore m => EClassId -> EGraphST m Int +getSize c = (_size . _info) <$> getEClass c +{-# INLINE getSize #-} +isSizeOf :: (Int -> Bool) -> EClass -> Bool +isSizeOf p = p . _size . _info +{-# INLINE isSizeOf #-} +getBestFitness :: ClassStore m => EGraphST m (Maybe Double) +getBestFitness = do + mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB) + case mbec of + Just bec -> (_fitness . _info) <$> getEClass bec + Nothing -> pure Nothing +getDL :: ClassStore m => EClassId -> EGraphST m (Maybe Double) +getDL c = (_dl . _info) <$> getEClass c +{-# INLINE getDL #-} diff --git a/src/Algorithm/EqSat/Info.hs b/src/Algorithm/EqSat/Info.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Info.hs @@ -0,0 +1,222 @@ +----------------------------------------------------------------------------- +-- | +-- 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 +import Control.Monad.State +import Data.AEq (AEq ((~==))) +import Data.IntMap (IntMap) +import qualified Data.IntMap as IntMap +import Data.Map (Map) +import qualified Data.Map as Map +import Data.SRTree +import Data.SRTree.Eval (evalFun, evalOp, Target) +import Data.HashSet (HashSet) +import qualified Data.HashSet as Set +import qualified Data.Set as RangeSet +import qualified Data.IntSet as IntSet +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Queries + +import qualified Data.Set as TrueSet + +-- * 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) + +-- | Fetch consts, cost, and size for all children in a single state traversal +getChildrenData :: ClassStore m => [EClassId] -> EGraphST m [(Consts, Cost, Int)] +getChildrenData ids = do + ids' <- mapM canonical ids + mapM (\cid -> do + ec <- getEClass cid + let d = _info ec + pure (_consts d, _cost d, _size d)) ids' +{-# INLINE getChildrenData #-} + +-- | Calculate e-node data (constant values and cost) +makeAnalysis :: ClassStore m => CostFun -> ENode -> EGraphST m EClassData +makeAnalysis costFun enode = + do let cs = eChildren enode + childData <- getChildrenData cs + let (consts', costs', sizes) = unzip3 childData + consts = combineNode enode consts' + cost = costNode enode costs' + sz = sum sizes + enode' <- canonize enode + pure $ EData cost enode' consts Nothing Nothing [] (sz + 1) + where + -- ENAry folds children pairwise (constant folding over a multiset); the + -- binary skeleton cannot represent n children. + combineNode (ENAry op _) cs = foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) cs + combineNode _ cs = combineConsts (replaceChildren cs (fromENode enode)) + -- ENAry is a single flattened op node: op cost + sum of child costs. + costNode (ENAry op _) cs = costFun (Bin (toOp op) 0 0) + sum cs + costNode _ cs = costFun (replaceChildren cs (fromENode enode)) + +getChildrenMinHeight :: ClassStore m => ENode -> EGraphST m Int +getChildrenMinHeight enode = do + let children = eChildren enode + if null children then pure 0 else do + children' <- mapM canonical children + hs <- mapM (fmap _height . getEClass) children' + pure (minimum hs) + +-- | update the heights of each e-class +-- won't work if there's no root +calculateHeights :: ClassStore m => EGraphST m () +calculateHeights = + do queue <- findRootClasses + classes <- allKeys + let nClasses = length classes + forM_ classes (setHeight nClasses) -- set all heights to max possible height (number of e-classes) + forM_ queue (setHeight 0) -- set root e-classes height to zero + go queue (TrueSet.fromList queue) 1 -- next height is 1 + where + setHeight :: ClassStore m => Int -> EClassId -> EGraphST m () + setHeight x eId' = + do eId <- canonical eId' + ec <- getEClass eId + let ec' = over height (const x) ec + insertClass ec' + + setMinHeight :: ClassStore m => Int -> EClassId -> EGraphST m () + setMinHeight x eId' = -- set height to the minimum between current and x + do eId <- canonical eId' + h <- _height <$> getEClass eId + setHeight (min h x) eId + + getChildrenEC :: ClassStore m => EClassId -> EGraphST m [EClassId] + getChildrenEC ec' = do ec <- getEClass ec' + pure $ concatMap eChildren (_eNodes ec) + + 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 :: ClassStore m => CostFun -> ENode -> EGraphST m Cost +calculateCost f enode = + do let cs = eChildren enode + costs <- traverse (fmap (_cost . _info) . getEClass) cs + pure $ case enode of + ENAry op _ -> f (Bin (toOp op) 0 0) + sum costs + _ -> f (replaceChildren costs (fromENode enode)) + +-- | check whether an e-node evaluates to a const +calculateConsts :: ClassStore m => ENode -> EGraphST m Consts +calculateConsts enode = + do let cs = eChildren enode + consts <- traverse (fmap (_consts . _info) . getEClass) cs + let c = case enode of + ENAry op _ -> foldr1 (\a b -> combineConsts (Bin (toOp op) a b)) consts + _ -> combineConsts (replaceChildren consts (fromENode enode)) + case c of + ConstVal x | isNaN x -> pure (ConstVal x) + a -> pure a + +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 :: ClassStore m => EClassId -> Double -> [Target] -> EGraphST m () +insertFitness eId' fit params = + do eId <- canonical eId' + tree <- getBestExpr eId + let p = fromIntegral (length params) + let f_compl = countNodes tree * log (countUniqueTokens tree) + p * (log (2 * pi * exp(1 - log 3)) - log p) / 2.0 + ec <- getEClass eId + let oldFit = _fitness . _info $ ec + let newInfo = (_info ec){_fitness = Just fit, _theta = params} + newEc = ec{_info = newInfo} + sz = _size newInfo + insertClass newEc + case oldFit of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.delete eId) + . over (eDB . fitRangeDB) (insertRange eId fit) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty) + . over (eDB . dlRangeDB) (insertRange eId f_compl) + Just oldVal -> modify' $ over (eDB . fitRangeDB) (insertRange eId fit . removeRange eId oldVal) + . over (eDB . sizeFitDB) (IntMap.adjust (insertRange eId fit . removeRange eId oldVal) sz) + +insertDL :: ClassStore m => EClassId -> Double -> EGraphST m () +insertDL eId fit' = + do let fit = negate fit' + ec <- getEClass eId + let sz = _size . _info $ ec + newInfo = (_info ec){_dl = Just fit'} + newEc = ec{_info=newInfo} + insertClass newEc + modify' $ over (eDB . dlRangeDB) (insertRange eId fit) + . over (eDB . sizeDLDB) (IntMap.adjust (insertRange eId fit) sz . IntMap.insertWith RangeSet.union sz RangeSet.empty) + + 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,213 @@ +{-# 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 qualified Data.Set as RangeSet +import Control.Monad.State ( gets, modify' ) +import Control.Lens ( over ) +import Data.Maybe +import Data.SRTree (childrenOf) + +getEClassesThat :: ClassStore m => (EClass -> Bool) -> EGraphST m [EClassId] +getEClassesThat p = do + classes <- allClasses + pure [ _eClassId ec | ec <- classes, p ec ] + +updateFitness :: ClassStore m => Double -> EClassId -> EGraphST m () +updateFitness f ecId = do + ec <- getEClass ecId + let info = _info ec + insertClass ec{_info=info{_fitness = Just f}} + +-- | returns all the root e-classes (e-class without parents) +findRootClasses :: ClassStore m => EGraphST m [EClassId] +findRootClasses = do + classes <- allClasses + pure [ _eClassId ec | ec <- classes, isParent (_eClassId ec, ec) ] + where + isParent (k, v) = Prelude.null (_parents v) || (k `Set.member` (Set.map fst (_parents v))) + +-- | returns the e-class id with the best fitness that +-- is true to a predicate +getTopECLassThat :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopECLassThat b n p = do + let f = if b then _fitRangeDB else _dlRangeDB + gets (f . _eDB) + >>= go n [] + where + go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go 0 bests rt = pure bests + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests t + else if p ec + then go (m-1) (ecId:bests) t + else go m bests t + +getTopEClassInRange :: ClassStore m => Bool -> Int -> (EClass -> Double) -> [(Double, Double)] -> EGraphST m [EClassId] +getTopEClassInRange b n p range = do + let f = if b then _fitRangeDB else _dlRangeDB + gets (f . _eDB) + >>= go n [] range + where + inRange v (x, y) + | v >= x && v <= y = 0 + | v < x = -1 + | v > y = 1 + | otherwise = 1 + + go :: ClassStore m => Int -> [EClassId] -> [(Double, Double)] -> RangeTree Double -> EGraphST m [EClassId] + go _ bests [] _ = pure bests + go 0 bests (r:rs) rt = go n bests rs rt + go m bests (r:rs) rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests (r:rs) t + else do let v = p ec + case (v `inRange` r) of + 0 -> go (m-1) (ecId:bests) (r:rs) t + -1 -> go n bests rs (RangeSet.insert y t) + 1 -> go m bests (r:rs) t + +getTopECLassIn :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopECLassIn b n p ecs' = do + let f = if b then _fitRangeDB else _dlRangeDB + gets (f . _eDB) + >>= go n [] + where + ecs = Set.fromList ecs' + go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go 0 bests rt = pure bests + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests t + else if ecId `Set.member` ecs && p ec + then go (m-1) (ecId:bests) t + else go m bests t + +getTopECLassNotIn :: ClassStore m => Bool -> Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopECLassNotIn b n p ecs' = do + let f = if b then _fitRangeDB else _dlRangeDB + gets (f . _eDB) + >>= go n [] + where + ecs = Set.fromList ecs' + + go :: ClassStore m => Int -> [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go 0 bests rt = pure bests + go m bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go m bests t + else if not (ecId `Set.member` ecs) && p ec + then go (m-1) (ecId:bests) t + else go m bests t + +getAllEvaluatedEClasses :: ClassStore m => EGraphST m [EClassId] +getAllEvaluatedEClasses = do + gets (_fitRangeDB . _eDB) + >>= go [] + where + go :: ClassStore m => [EClassId] -> RangeTree Double -> EGraphST m [EClassId] + go bests rt = case RangeSet.maxView rt of + Nothing -> pure bests + Just (y, t) -> + let x = snd y + in do ecId <- canonical x + ec <- getEClass ecId + if (maybe True (isInfinite) . _fitness . _info $ ec) + then go bests t + else go (ecId:bests) t + +getTopEClassWithSize :: Monad m => Bool -> Int -> Int -> EGraphST m [EClassId] +getTopEClassWithSize b sz n = do + let fun = if b then _sizeFitDB else _sizeDLDB + gets (go n [] . (IntMap.!? sz) . fun . _eDB) + where + go _ bests Nothing = [] + go 0 bests (Just rt) = bests + go m bests (Just rt) = case RangeSet.maxView rt of + Nothing -> bests + Just ((f, x), t) -> if isInfinite f || isNaN f then go m bests (Just t) else go (m-1) (x:bests) (Just t) + +getTopFitEClassThat :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopFitEClassThat = getTopECLassThat True +getTopDLEClassThat :: ClassStore m => Int -> (EClass -> Bool) -> EGraphST m [EClassId] +getTopDLEClassThat = getTopECLassThat False +getTopFitEClassIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopFitEClassIn = getTopECLassIn True +getTopDLEClassIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopDLEClassIn = getTopECLassIn False +getTopFitEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopFitEClassNotIn = getTopECLassNotIn True +getTopDLEClassNotIn :: ClassStore m => Int -> (EClass -> Bool) -> [EClassId] -> EGraphST m [EClassId] +getTopDLEClassNotIn = getTopECLassNotIn False +getTopFitEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId] +getTopFitEClassWithSize = getTopEClassWithSize True +getTopDLEClassWithSize :: Monad m => Int -> Int -> EGraphST m [EClassId] +getTopDLEClassWithSize = getTopEClassWithSize False + +rebuildAllRanges :: ClassStore 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 :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double) +canonizeRange = fmap RangeSet.fromList . mapM (\(x, eid) -> (x,) <$> canonical eid) . RangeSet.toList + +rebuildRange :: ClassStore m => RangeTree Double -> EGraphST m (RangeTree Double) +rebuildRange rt = do + canonRt <- canonizeRange rt + pure $ snd $ go canonRt + where + go rt' = case RangeSet.maxView rt' of + Nothing -> (Set.empty, RangeSet.empty) + Just ((x, eid), rest) -> + let (seen, result) = go rest + in if Set.member eid seen + then (seen, result) + else (Set.insert eid seen, RangeSet.insert (x, eid) result) + diff --git a/src/Algorithm/EqSat/SearchSR.hs b/src/Algorithm/EqSat/SearchSR.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/SearchSR.hs @@ -0,0 +1,277 @@ +----------------------------------------------------------------------------- +-- | +-- 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 Data.SRTree.Eval (compileLoss) +import System.Random +import Control.Monad.State.Strict +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (mapConcurrently) +import Data.Maybe (catMaybes) +import Control.Exception (evaluate) +import qualified Control.DeepSeq as DeepSeq +import Algorithm.EqSat.Egraph +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.AD (ADBackEnd(..)) +import Algorithm.SRTree.AD.Unboxed (setMTPopParallel) +import qualified Data.IntMap as IM +import qualified Data.IntSet as IntSet +import qualified Data.SRTree.Random as Random +import Data.Function ( on ) +import Algorithm.SRTree.NonlinearOpt +import Control.Monad ( when, replicateM, forM, forM_ ) +import Numeric.Optimization.NLOPT +import Algorithm.EqSat.Info +import Algorithm.EqSat.Build +import Data.SRTree.Random +import Algorithm.EqSat.Queries +import Data.List ( maximumBy ) +import qualified Data.List as Data.List +import qualified Data.HashMap.Strict as HashMap +import qualified Data.Vector.Unboxed as V + +-- Environment of an e-graph with support to random generator and IO +type RndEGraph a = EGraphST (StateT StdGen IO) a + +io :: IO a -> RndEGraph a +io = lift . lift +{-# INLINE io #-} +rnd :: StateT StdGen IO a -> RndEGraph a +rnd = lift +{-# INLINE rnd #-} + +-- | Run an 'RndEGraph' action against a read-only egraph snapshot with the given +-- generator (for concurrent workers that do not mutate the shared egraph). +runRndEGraph :: EGraph -> StdGen -> RndEGraph a -> IO a +runRndEGraph eg g m = do + ((a, _), _) <- runStateT (runStateT m eg) g + pure a +{-# INLINE runRndEGraph #-} + +-- | Fit a batch of e-classes in parallel, then insert the results serially. +-- Semantics mirror 'updateIfNothing' (skip already-fitted) unless 'force' is +-- True. The shared 'StdGen' is split once; each worker gets its own generator, +-- so the global draw sequence differs from the serial search (acceptable). +-- While the batch runs, the MultiThread backend is switched to single-chunk so +-- cores go to the batch rather than oversubscribing the inner per-tree split. +fitBatch :: Bool + -> (Fix SRTree -> RndEGraph (Double, [Target])) + -> [EClassId] + -> RndEGraph () +fitBatch force fitFun ecs0 = do + ecs <- Prelude.mapM canonical ecs0 + jobs <- fmap catMaybes $ forM ecs $ \ec -> do + mf <- getFitness ec + if force || mf == Nothing + then do tree <- getBestExpr ec + pure (Just (ec, tree)) + else pure Nothing + case jobs of + [] -> pure () + _ -> do + nCaps <- io getNumCapabilities + g0 <- rnd get + let (seed, g1) = random g0 :: (Int, StdGen) + gs = [ mkStdGen (seed + fromIntegral i) | i <- [0 .. length jobs - 1] ] + jobsG = [ (ec, tree, g) | ((ec, tree), g) <- zip jobs gs ] + chunk k xs = [ [ xs !! j | j <- [i, i + k .. length xs - 1] ] | i <- [0 .. k - 1] ] + rnd (put g1) + eg <- get + io (setMTPopParallel True) + results <- io $ fmap concat (mapConcurrently (mapM (runJob eg fitFun)) (chunk nCaps jobsG)) + io (setMTPopParallel False) + forM_ results $ \(ec0, f, p) -> insertFitness ec0 f p + where + runJob :: EGraph -> (Fix SRTree -> RndEGraph (Double, [Target])) -> (EClassId, Fix SRTree, StdGen) -> IO (EClassId, Double, [Target]) + runJob eg fit' (ec, tree, g) = do + (f, p) <- runRndEGraph eg g (fit' tree) + f' <- evaluate (DeepSeq.force f) + p' <- evaluate (DeepSeq.force p) + pure (ec, f', p') + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +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 :: ADBackEnd -> Bool -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> Target -> (Double, Target) +fitnessFun backend skipVal nIter loss (x, y, mYErr) (x_val, y_val, mYErr_val) tree thetaOrig = + if isNaN val + then (-(1/0), theta) + else (val, theta) + where + nParams = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL Gaussian then 1 else 0 + (theta, lossVal, _) = minimizeNLL' VAR1 backend loss mYErr nIter x y tree thetaOrig + evalF a b c = negate $ compileLoss a (buildLoss loss (fromIntegral (V.length b)) tree) b c $ if nParams == 0 then thetaOrig else theta + -- at folds=1 the validation split is the training data itself, so the + -- train loss returned by minimizeNLL' already is the val loss; skipping + -- the separate compileLoss below avoids re-evaluating every expression. + val = if skipVal then negate lossVal else evalF x_val y_val mYErr_val + +--{-# INLINE fitnessFun #-} + +fitnessFunRep :: ADBackEnd -> Bool -> Int -> Int -> Loss -> DataSet -> DataSet -> Fix SRTree -> RndEGraph (Double, Target) +fitnessFunRep backend skipVal nRep nIter loss dataTrain dataVal tree = do + let nParams = countParamsUniq tree + if loss == NLL ROXY then 3 else if loss == NLL Gaussian then 1 else 0 + thetaOrigs <- replicateM nRep (rnd $ randomVec nParams) + pure $ maximumBy (\(x, _) (y, _) -> compare x y) $ Prelude.map (fitnessFun backend skipVal nIter loss dataTrain dataVal tree) thetaOrigs +--{-# INLINE fitnessFunRep #-} + + +fitnessMV :: ADBackEnd -> Bool -> Bool -> Int -> Int -> Loss -> [(DataSet, DataSet)] -> Fix SRTree -> RndEGraph (Double, [Target]) +fitnessMV backend skipVal shouldReparam nRep nIter loss dataTrainsVals _tree = do + let tree = if shouldReparam then relabelParams _tree else relabelParamsOrder _tree + response <- forM dataTrainsVals $ \(dt, dv) -> fitnessFunRep backend skipVal nRep nIter loss 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, [Target])) -> 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 <- (_consts . _info) <$> getEClass rndId + case constType of + NotConst -> pure $ Just rndId + _ -> pure Nothing + else pure Nothing + +getParetoEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopFitEClassWithSize i n) +getParetoDLEcsUpTo n maxSize = concat <$> forM [1..maxSize] (\i -> getTopDLEClassWithSize i n) + +getBestExprWithSize n = + do ec <- getTopFitEClassWithSize n 1 >>= traverse canonical + case ec of + (x:_) -> do bestFit <- getFitness x + bestP <- (_theta . _info) <$> getEClass x + pure [(x, bestFit)] + [] -> pure [] + +insertRndExpr maxSize rndTerm rndNonTerm = + do grow <- rnd toss + n <- rnd (randomFrom [if maxSize > 4 then 4 else 1 .. max 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 + mbec <- gets (fmap snd . getGreatest . _fitRangeDB . _eDB) + case mbec of + Just bec -> do bestFit <- (_fitness . _info) <$> getEClass bec + printExprFun 0 bec + Nothing -> pure () + +--paretoFront :: Int -> (Int -> EClassId -> RndEGraph ()) -> RndEGraph () +paretoFront fitFun maxSize printExprFun = go 1 0 (-(1.0/0.0)) + where + go :: Int -> Int -> Double -> RndEGraph [[String]] + go n ix f + | n > maxSize = pure [] + | otherwise = do + ecList <- getBestExprWithSize n + case ecList of + ((ec, Just f'):_) -> do + let improved = f' >= f && (not . isNaN) f' && (not . isInfinite) f' + ec' <- canonical ec + if improved + then do refit fitFun ec' + t <- printExprFun ix ec' + ts <- go (n+1) (ix + if improved then 1 else 0) (max f f') + pure (t:ts) + else go (n+1) (ix + if improved then 1 else 0) (max f f') + _ -> go (n+1) ix f + +evaluateUnevaluated fitFun = do + ec <- gets (IntSet.toList . _unevaluated . _eDB) + 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 ((HashMap.member en) . _eNodeToEClass) +doesNotExist en = gets ((not . HashMap.member en) . _eNodeToEClass) + +-- | check whether the partial tree defined by a list of ancestors will create +-- a non-existent expression when combined with a certain e-node. +doesNotExistGens :: [Maybe (EClassId -> ENode)] -> ENode -> RndEGraph Bool +doesNotExistGens [] en = gets ((not . HashMap.member en) . _eNodeToEClass) +doesNotExistGens (mGrand:grands) en = do b <- gets ((not . HashMap.member en) . _eNodeToEClass) + if b + then pure True + else case mGrand of + Nothing -> pure False + Just gf -> do ec <- gets ((HashMap.! en) . _eNodeToEClass) + en' <- canonize (gf ec) + doesNotExistGens grands en' + +-- | check whether combining a partial tree `parent` with the e-node `en'` +-- will create a new expression +checkToken parent en' = do en <- canonize en' + mEc <- gets (HashMap.lookup 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,284 @@ +{-# 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, + Condition (Condition), + NChild (Ch, MapP, Rest), + Pattern (Fixed, Hole, NAry, VarPat), + Rule (..), + Subst, + SubVal (SVMap, SVOne), + getInt, + ) +import Control.Monad.State.Strict (evalState) +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IM +import Data.Map (Map) +import qualified Data.Map as Map +import Data.SRTree + +-- | A constraint over a match's substitution: when applied to a substitution it +-- runs in the e-graph monad and fetches e-class data through 'ClassStore', so it +-- works on a paged (out-of-core) graph whose resident cache is bounded/empty. +type ConstrFun = Pattern -> Condition + +constrainOnVal :: (Consts -> Bool) -> Pattern -> Condition +constrainOnVal f (VarPat c) = Condition $ \subst -> do + let cid = getInt $ case Map.lookup (Right (fromEnum c)) subst of + Nothing -> error $ "CONSTRAINVAL_MISSING var=" <> show (fromEnum c) <> " substSize=" <> show (Map.size subst) + Just (SVOne v) -> v + Just (SVMap _) -> error $ "CONSTRAINVAL_REST_AS_SINGLE var=" <> show (fromEnum c) + ec <- getEClass cid + pure (f (_consts . _info $ ec)) +constrainOnVal _ _ = Condition $ \_ -> pure False + +-- TODO: aux functions to avoid repeated pattern in constraint creation +-- +-- 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 + +-- | e-class ids bound to a rest variable +restEidsOf :: Char -> Subst -> [EClassId] +restEidsOf c subst = case Map.lookup (Right (fromEnum c)) subst of + Just (SVMap m) -> expandedList m + _ -> [] + +-- | every e-class bound to a rest variable holds a valid value +allValidRest :: Char -> Condition +allValidRest c = Condition $ \subst -> do + let eids = restEidsOf c subst + validEid eid = getEClass eid >>= \ec -> + pure $ case _consts . _info $ ec of + ConstVal x -> not (isNaN x || isInfinite x) + _ -> True + and <$> mapM validEid eids + +-- basic algebraic rules +rewriteBasic :: [Rule] +rewriteBasic = + [ + -- B7/B8/C5: factor a common term out of a sum of products, and the + -- reverse (distribute), which make x*(y+z) and x*y+x*z equivalent. + NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1']) + , Ch (NAry EMul [Ch "x", Rest '2']) + , Rest '3' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1', Rest '2']) + ]) + , Rest '3' ] + , NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1']) + ]) + , Rest '2' ] + :=> + NAry EAdd [ MapP (NAry EMul [Ch "x", Ch Hole]) '1' + , Rest '2' ] + -- C5: x*y - z*x = x*(y - z) + , NAry EAdd [ Ch (NAry EMul [Ch "x", Rest '1']) + , Ch (NAry EMul [Ch (Fixed (Const (-1))), Ch "x", Ch "z"]) + , Rest '3' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Rest '1', Ch (negate (VarPat 'z'))]) + ]) + , Rest '3' ] + -- B1: group duplicate factors into a power (x*x = x^2) + , NAry EMul [Ch "x", Ch "x"] :=> "x" ** 2 + -- C9: binomial expansion of a closed 2-ary square + , ("x" + "y") ** 2 :=> "x" ** 2 + 2 * "x" * "y" + "y" ** 2 + -- C10: x^2 + x*y + ... = x*(x + y) + ... + , NAry EAdd [ Ch (Fixed (Bin Power (VarPat 'x') (Fixed (Const 2)))) + , Ch (NAry EMul [Ch "x", Rest '1']) + , Rest '2' ] + :=> + NAry EAdd [ Ch (NAry EMul [ Ch "x" + , Ch (NAry EAdd [Ch "x", Rest '1']) + ]) + , Rest '2' ] + ] + +-- rules for nonlinear functions +rewritesFun :: [Rule] +rewritesFun = + [ + log (exp "x") :=> "x" + -- C11: log(x*y*z*...) = log x + log y + ... + , log (NAry EMul [Rest '1']) :=> NAry EAdd [MapP (Fixed (Uni Log Hole)) '1'] + , log ("x" ** "y") :=> "y" * log "x" + , log (powabs "x" "y") :=> "y" * log (abs "x") + -- C12: abs(x*y*z*...) = abs x * abs y * ... + , abs (NAry EMul [Rest '1']) :=> NAry EMul [MapP (Fixed (Uni Abs Hole)) '1'] + , abs ("x" ** "y") :=> abs "x" ** "y" + , recip (recip "x") :=> "x" :| isNotZero "x" + -- C13: (x*y*z*...)^w = x^w * y^w * ... [was disabled: combinatorial blowup on (x*x)^t; the multiset matcher + matchCap bound that] + , (NAry EMul [Rest '1']) ** "z" :=> NAry EMul [MapP (Hole ** VarPat 'z') '1'] + , abs "x" ** "y" :=> "x" ** "y" :| isEven "y" + -- C14: sqrt(x*x) = abs x + , sqrt (NAry EMul [Ch "x", Ch "x"]) :=> abs "x" + ] + +-- Rules that reduces redundant parameters +constReduction :: [Rule] +constReduction = + [ + -- B3: 0 + rest = rest + NAry EAdd [Ch (Fixed (Const 0)), Rest '1'] :=> NAry EAdd [Rest '1'] + , "x" ** 1 :=> "x" + , powabs "x" 1 :=> abs "x" + + -- B9: x^y * x^z = x^(y+z) + , NAry EMul [Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin Power (VarPat 'x') (VarPat 'z')))] + :==: + Fixed (Bin Power (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')])) + :| isPositive "x" + -- B10: |x|^y * |x|^z = |x|^(y+z) (fixed: target used "y+x" instead of "y+z") + , NAry EMul [Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'y'))), Ch (Fixed (Bin PowerAbs (VarPat 'x') (VarPat 'z')))] + :=> + Fixed (Bin PowerAbs (VarPat 'x') (NAry EAdd [Ch (VarPat 'y'), Ch (VarPat 'z')])) + -- B11: (x^y)^z = x^(y*z) + , Fixed (Bin Power (Fixed (Bin Power (VarPat 'x') (VarPat 'y'))) (VarPat 'z')) + :==: + Fixed (Bin Power (VarPat 'x') (NAry EMul [Ch (VarPat 'y'), Ch (VarPat 'z')])) + :| isPositive "x" + , powabs (powabs "x" "y") "z" :=> powabs "x" ("y" * "z") + ] + +rewritesWithConstant :: [Rule] +rewritesWithConstant = + [ + "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") + -- B4: 0 * rest = 0 (provided every factor is valid) + , NAry EMul [Ch (Fixed (Const 0)), Rest '1'] :=> 0 :| allValidRest '1' + , 0 ** "x" :=> 0 :| isPositive "x" + , powabs 0 "x" :=> 0 + -- n-ary cancellation: x + y - x = y + , NAry EAdd [ Ch "a" + , Ch (NAry EMul [ Ch (Fixed (Const (-1.0))), Ch "a" ]) + , Rest 'r' ] + :=> NAry EAdd [Rest 'r'] + -- combining like terms: x + x = 2*x + , NAry EAdd [ Ch "a", Ch "a", Rest 'r' ] + :=> NAry EAdd [ Ch (2 * "a"), Rest 'r' ] + ] +rewritesWithParam :: [Rule] +rewritesWithParam = + [ + "x" - "x" :=> Fixed (Param 0) + , "x" / "x" :=> Fixed (Param 0) :| isNotZero "x" + , 1 ** "x" :=> Fixed (Param 0) + , powabs 1 "x" :=> Fixed (Param 0) + ] + +rewritesSimple :: [Rule] +rewritesSimple = rewriteBasic <> constReduction <> rewritesFun +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` emptyGraphNoTrack + +-- | simplifies with custom parameters +simplifyEqSat :: [Rule] -> CostFun -> Int -> Fix SRTree -> Fix SRTree +simplifyEqSat rwrts costFun it t = eqSat t rwrts costFun it `evalState` emptyGraph + +-- | apply a single step of merge-only using default rules +applyMergeOnlyDftl :: ClassStore m => CostFun -> EGraphST m () +applyMergeOnlyDftl costFun = applySingleMergeOnlyEqSat costFun rewrites diff --git a/src/Algorithm/EqSat/Store.hs b/src/Algorithm/EqSat/Store.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/EqSat/Store.hs @@ -0,0 +1,243 @@ +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveDataTypeable #-} + +module Algorithm.EqSat.Store + ( GraphRows(..) + , EClassRow(..) + , exportEGraph + , importEGraph + , mergeEGraph + , rebuildDBs + ) where + +import Control.Lens ( over ) +import Control.Monad ( forM, forM_, foldM ) +import Control.Monad.Identity ( Identity, runIdentity ) +import Control.Monad.State.Strict ( StateT, execStateT, modify', gets ) +import GHC.Generics ( Generic ) +import GHC.Stack ( HasCallStack ) + +import qualified Data.HashMap.Strict as HashMap +import Data.HashMap.Strict ( HashMap ) +import qualified Data.HashSet as Set +import qualified Data.IntMap.Strict as IntMap +import Data.IntMap.Strict ( IntMap ) +import qualified Data.IntSet as IntSet +import qualified Data.Set as RangeSet +import Data.List ( sortOn ) + +import Data.SRTree +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build + +-- | Row representation of the core (structural) state of an e-graph, +-- normalized for external storage (e.g. a relational DB). +data GraphRows = GraphRows + { _grCanonical :: IntMap EClassId -- ^ eid -> canonical representative (self-loop for roots) + , _grENodeToEClass :: HashMap ENode EClassId -- ^ canonical e-node -> its e-class + , _grEClasses :: IntMap EClassRow -- ^ canonical e-class id -> data row + , _grNextId :: Int -- ^ next free e-class id + , _grTrackDBs :: Bool -- ^ whether range DBs are maintained + } deriving (Show, Eq, Generic) + +-- | Per-e-class data row. +data EClassRow = EClassRow + { _rcNodes :: Set.HashSet ENode + , _rcParents :: Set.HashSet (EClassId, ENode) + , _rcHeight :: Int + , _rcInfo :: EClassData + } deriving (Show, Eq, Generic) + +-- | Export the core structural state of an e-graph into a normalised row format. +exportEGraph :: EGraph -> GraphRows +exportEGraph eg = GraphRows + { _grCanonical = _canonicalMap eg + , _grENodeToEClass = _eNodeToEClass eg + , _grEClasses = IntMap.map toRow (_eClass eg) + , _grNextId = _nextId (_eDB eg) + , _grTrackDBs = _trackDBs (_eDB eg) + } + where + toRow ec = EClassRow (_eNodes ec) (_parents ec) (_height ec) (_info ec) + +-- | Reconstruct an e-graph from normalised rows, rebuilding all derived indexes. +-- +-- Real e-graphs may carry stale @_eNodeToEClass@ entries left behind by +-- merges (a node pointing at a class whose canonical representative is +-- another class). Such entries are canonicalized at import: node -> class +-- values are routed through the canonical map and any non-root class rows +-- are dropped. Parent pointers are recomputed from the canonicalized node +-- map so they never reference dead classes. +importEGraph :: GraphRows -> Either String EGraph +importEGraph rows + | not (validate rows) = Left (validationMsg rows) + | otherwise = Right (runIdentity $ execStateT rebuildDBs (buildCore (canonicalize rows))) + +-- | Normalize stale rows: route node->class values through the canonical map +-- and drop non-root class rows. +-- +-- Parent pointers come from the stored @_rcParents@ when a class has any +-- (e.g. after a storage-layer round-trip through the @parent@ table); parent +-- class ids are routed through the canonical map so they never reference dead +-- classes. Classes without stored parents (legacy rows, hand-built rows) fall +-- back to recomputing parents from the canonicalized node map. +canonicalize :: GraphRows -> GraphRows +canonicalize rows = + let canon = _grCanonical rows + rep eid = IntMap.findWithDefault eid eid canon + nodeMap' = HashMap.map rep (_grENodeToEClass rows) + classes' = IntMap.filterWithKey + (\eid _ -> IntMap.lookup eid canon == Just eid) + (_grEClasses rows) + parents' = IntMap.fromListWith Set.union + [ (c, Set.singleton (eid, en)) + | (en, eid) <- HashMap.toList nodeMap' + , c <- eChildren en ] + stored' = IntMap.mapWithKey + (\_ r -> Set.map (\(pEid, pEn) -> (rep pEid, pEn)) (_rcParents r)) + classes' + fixRow eid r = + let stored = IntMap.findWithDefault Set.empty eid stored' + in r { _rcParents = if Set.null stored + then IntMap.findWithDefault Set.empty eid parents' + else stored } + in rows { _grENodeToEClass = nodeMap' + , _grEClasses = IntMap.mapWithKey fixRow classes' } + +buildCore :: GraphRows -> EGraph +buildCore rows = EGraph + { _canonicalMap = _grCanonical rows + , _eNodeToEClass = _grENodeToEClass rows + , _eClass = IntMap.mapWithKey mkEClass (_grEClasses rows) + , _eDB = (emptyDB){ _nextId = _grNextId rows, _trackDBs = _grTrackDBs rows } + , _classStore = Nothing + } + where + mkEClass eid r = EClass eid (_rcNodes r) (_rcParents r) (_rcHeight r) (_rcInfo r) + +rebuildDBs :: EGraphST Identity () +rebuildDBs = do + -- Rebuild the pattern database from the canonical e-node -> class mapping + nodes <- gets _eNodeToEClass + forM_ (HashMap.toList nodes) $ \(en, eid) -> addToDB en eid + + -- Rebuild range/size indexes from class info + classes <- gets _eClass + forM_ (IntMap.toList classes) $ \(eid, ec) -> do + let info = _info ec + sz = _size info + fit = _fitness info + dl = _dl info + modify' $ over (eDB . sizeDB) (IntMap.insertWith IntSet.union sz (IntSet.singleton eid)) + case fit of + Nothing -> modify' $ over (eDB . unevaluated) (IntSet.insert eid) + Just fn -> modify' $ over (eDB . fitRangeDB) (insertRange eid fn) + . over (eDB . sizeFitDB) (IntMap.insertWith RangeSet.union sz (RangeSet.singleton (fn, eid))) + case dl of + Nothing -> pure () + Just dn -> modify' $ over (eDB . dlRangeDB) (insertRange eid dn) + . over (eDB . sizeDLDB) (IntMap.insertWith RangeSet.union sz (RangeSet.singleton (dn, eid))) + +-- | Validate that the exported rows form a consistent graph. +-- +-- All referenced ids must be present in the canonical map. Node -> class +-- values and class rows may reference classes that are not their own +-- canonical representative (stale entries left behind by merges); those are +-- repaired by 'canonicalize' during import. +validate :: GraphRows -> Bool +validate rows = + let canon = _grCanonical rows + classes = _grEClasses rows + nodeIds = HashMap.keys (_grENodeToEClass rows) + extraIds = IntMap.keys classes + ++ HashMap.elems (_grENodeToEClass rows) + ++ concatMap eChildren nodeIds + inCanon = all (`IntMap.member` canon) extraIds + nextOk = _grNextId rows >= 0 + in inCanon && nextOk + +validationMsg :: GraphRows -> String +validationMsg rows + | not inCanon = "some e-node/e-class id is not present in the canonical map" + | not nextOk = "next id is negative" + | otherwise = "invalid GraphRows" + where + canon = _grCanonical rows + classes = _grEClasses rows + nodeIds = HashMap.keys (_grENodeToEClass rows) + extraIds = IntMap.keys classes + ++ HashMap.elems (_grENodeToEClass rows) + ++ concatMap eChildren nodeIds + inCanon = all (`IntMap.member` canon) extraIds + nextOk = _grNextId rows >= 0 + +-- | Return canonical e-class ids ordered children-before-parents (ascending height). +classOrder :: GraphRows -> Either String [EClassId] +classOrder rows = + Right $ map fst $ sortOn (_rcHeight . snd) $ IntMap.toAscList (_grEClasses rows) + +-- | Remap a B-e-graph's e-node into A's id-space using the correspondence map. +remapNode + :: GraphRows -- ^ rows of graph B (source) + -> IntMap EClassId -- ^ corr: B canonical id -> A id + -> ENode + -> Either String ENode +remapNode rowsB corr = go + where + canonB :: EClassId -> EClassId + canonB cid = IntMap.findWithDefault cid cid (_grCanonical rowsB) + + toA :: EClassId -> Either String EClassId + toA cid = + case IntMap.lookup (canonB cid) corr of + Just eidA -> Right eidA + Nothing -> Left ("child " <> show cid <> " of graph B not yet merged") + + go (EVar ix) = Right (EVar ix) + go (EParam ix) = Right (EParam ix) + go (EConst x) = Right (EConst x) + go (EUni f t) = EUni f <$> toA t + go (EBin op l r) = EBin op <$> toA l <*> toA r + go (ENAry op m) = do + m' <- foldM step IntMap.empty (IntMap.toList m) + Right (ENAry op m') + where + step acc (cid, n) = do + cidA <- toA cid + pure (IntMap.insertWith (+) cidA n acc) + +-- | Merge class ids by unioning their e-classes under the given cost function. +mergeClass :: HasCallStack => CostFun -> EClassId -> EClassId -> EGraphST Identity EClassId +mergeClass costFun x y = + if x == y then pure x else merge costFun x y + +-- | Structurally merge graph @b@ into a copy of graph @a@. +-- +-- The e-nodes of @b@ are canonicalized under @a@'s id space, deduplicated +-- against @a@'s existing content, and equivalent classes are unioned. Cost and +-- best of newly introduced content are computed with @costFun@ (i.e. merging +-- adopts @a@'s cost function). Dataset-specific values (fitness/DL/theta) are +-- NOT transferred: they are per-dataset data managed by the storage layer. +mergeEGraph :: HasCallStack => CostFun -> EGraph -> EGraph -> Either String EGraph +mergeEGraph costFun a b = + let rowsB = exportEGraph b + in case classOrder rowsB of + Left err -> Left err + Right order -> Right (runIdentity $ execStateT (step IntMap.empty order) a) + where + step :: IntMap EClassId -> [EClassId] -> EGraphST Identity () + step _ [] = rebuild costFun + step corr (bCanon : rest) = do + let ec = _grEClasses rowsB IntMap.! bCanon + resolved <- forM (Set.toList (_rcNodes ec)) $ \en -> + case remapNode rowsB corr en of + Left err -> pure (Left err) + Right enA -> Right <$> add costFun enA + case sequence resolved of + Left err -> error ("mergeEGraph: " <> err) -- pre-validated + Right [] -> step corr rest + Right (x : xs) -> do + rep <- foldM (mergeClass costFun) x xs + step (IntMap.insert bCanon rep corr) rest + rowsB = exportEGraph b diff --git a/src/Algorithm/SRTree/AD.hs b/src/Algorithm/SRTree/AD.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD.hs @@ -0,0 +1,32 @@ +----------------------------------------------------------------------------- +-- | +-- 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 + ( compileFunAndGrad + , ADBackEnd(..) + ) where + +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as V +import Data.SRTree +import Algorithm.SRTree.AD.Unboxed + +data ADBackEnd = SingleThread | MultiThread deriving (Read, Show) + +compileFunAndGrad :: ADBackEnd -> [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double) -> Fix SRTree -> V.Vector Double -> (Double, V.Vector Double) +compileFunAndGrad SingleThread xss ys mYerr tree = + let ct = compileTree xss ys mYerr tree + in \theta -> evalGradVec ct theta +compileFunAndGrad MultiThread xss ys mYerr tree = + let cts = compileTreeMulti xss ys mYerr tree + in \theta -> evalGradMulti cts theta diff --git a/src/Algorithm/SRTree/AD/CompiledAD.hs b/src/Algorithm/SRTree/AD/CompiledAD.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD/CompiledAD.hs @@ -0,0 +1,39 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.AD.CompiledAD +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Automatic Differentiation for Expression trees +-- +----------------------------------------------------------------------------- + +module Algorithm.SRTree.AD.CompiledAD + ( CompiledTree(..) + ) where + +import Data.SRTree.Internal +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector as VB + +-- --------------------------------------------------------------------- +-- Public entry point -- same signature/behaviour as before. +-- --------------------------------------------------------------------- +data CompiledTree = CompiledTree + { ctNodes :: !(VB.Vector (SRTree Int)) -- id -> node, children already resolved to ids + , ctRoot :: !Int + , ctDyn :: !(VU.Vector Bool) -- id -> depends on theta? + , ctStatic :: VU.Vector Double -- flat [staticSlot * m + row]; only static nodes + , ctStaticBase :: !(VU.Vector Int) -- id -> staticSlot * m (0 for dynamic ids and Var leaves) + , ctM :: !Int + , ctNPred :: !Int -- root + 1 (stride for flat static) + , ctKind :: !(VU.Vector Int) -- id -> node kind: 0 Var, 1 Param, 2 Const, 3 Uni, 4 Bin + , ctArg :: !(VU.Vector Int) -- id -> Param: param ix; Var: var ix (-1 = y, -2 = yErr); Uni: child id; Bin: left id + , ctArg2 :: !(VU.Vector Int) -- id -> Bin: right id; else 0 + , ctFcode :: !(VU.Vector Int) -- id -> Uni: fromEnum Function + , ctOcode :: !(VU.Vector Int) -- id -> Bin: fromEnum Op + , ctVars :: !(VB.Vector (VU.Vector Double)) -- leaf source columns xss ++ [y, yErr] (referenced, not copied) + } diff --git a/src/Algorithm/SRTree/AD/Unboxed.hs b/src/Algorithm/SRTree/AD/Unboxed.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/AD/Unboxed.hs @@ -0,0 +1,965 @@ +{-# language FlexibleInstances, DeriveFunctor #-} +{-# language ScopedTypeVariables #-} +{-# language RankNTypes #-} +{-# language ViewPatterns #-} +{-# language FlexibleContexts #-} +{-# language BangPatterns #-} +{-# language TypeApplications #-} +{-# language MultiWayIf #-} +{-# LANGUAGE LambdaCase #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.SRTree.AD +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables +-- +-- Automatic Differentiation for Expression trees +-- +----------------------------------------------------------------------------- + +module Algorithm.SRTree.AD.Unboxed + ( compileTree + , compileTreeMulti + , evalGradMulti + , evalGrad + , evalGradVec + , evalLossVec + , CompiledTree(..) + , setMTPopParallel + ) where + +import Control.Monad (forM_, foldM, when, unless) +import Control.Monad.ST +import Data.STRef (newSTRef, readSTRef, modifySTRef') +import Data.Bifunctor (bimap, first, second) +import Data.SRTree.Derivative ( derivative ) +import Data.SRTree.Eval + ( Target, Theta, Columns, evalFun, evalOp, replicateAs ) +import Data.SRTree.Internal +import Data.SRTree.Print (showExpr) +import Data.SRTree.Recursion ( cataM, cata, accu ) +import qualified Data.Vector.Storable as V +import qualified Data.Vector.Storable.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VMB +import Debug.Trace (trace, traceShow) +import qualified Data.IntMap.Strict as IntMap +import Data.List ( foldl', foldl1' ) +import Data.Maybe (isJust, fromMaybe) + +import Control.Monad.State.Strict +import Control.Monad.Identity + + +import Data.List (transpose) +import System.IO.Unsafe (unsafePerformIO) +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (forConcurrently) +import Control.Exception (evaluate) +import Data.IORef (IORef, newIORef, writeIORef, readIORef) + +import qualified Data.Map.Strict as Map +import Algorithm.SRTree.AD.CompiledAD + +compileTree :: [VU.Vector Double] -> VU.Vector Double -> Maybe (VU.Vector Double) -> Fix SRTree -> CompiledTree +compileTree xss ys mYErr tree = + CompiledTree { ctNodes = nodes, ctRoot = root, ctDyn = dynArr, ctStatic = staticArr, ctStaticBase = staticBaseArr, ctM = m, ctNPred = root + 1 + , ctKind = kindArr, ctArg = argArr, ctArg2 = arg2Arr, ctFcode = fcodeArr, ctOcode = ocodeArr, ctVars = vars } + where + -- yErr is only defined when mYErr is present (a tree referencing Var -2 + -- always pairs with mYErr = Just e, see the likelihood loss wrappers). The + -- ctVars list must stay well-defined for every column even when mYErr is + -- Nothing -- the Accelerate leaf array concatenates the whole list -- so a + -- missing yErr is represented by a zero column rather than the bottom + -- `fromJust mYErr` (which the old static-array copy path could keep lazy). + yErr = case mYErr of + Just e -> e + Nothing -> VU.replicate m 0 + m = VU.length ys + -- Leaf source columns, referenced (never copied per tree): a static Var + -- leaf reads feature column ix (arg), y (arg = -1), or yErr (arg = -2) + -- straight from these run-fixed vectors instead of a materialized copy in + -- staticArr. ctVars ix = xss !! ix, ctVars nFeats = y, ctVars (nFeats+1) + -- = yErr. + vars = VB.fromList (xss <> [ys, yErr]) + nFeats = VB.length vars - 2 + + -- Rewrite x ** 2.0 into the unary Square kernel (x*x, fcode 17): + -- the loss wrap ((tree - y) ** 2) / m is the single hottest subgraph in + -- every NLopt call, and replacing the per-element pow with a multiply + -- avoids the slow ** (x**2.0 == x*x exactly, and the derivative 2x + -- matches), so no numerical semantics change. + tree' = rewritePowSq tree + + -- state: (structural CSE map, id -> node, id -> isDynamic, counter) + (_, int2key, dynMap, (subtract 1) -> root) = + cataM leftToRight alg tree' + `execState` (Map.empty, IntMap.empty, IntMap.empty, 0) + + nodes = VB.fromList (IntMap.elems int2key) + dynArr = VU.fromList (IntMap.elems dynMap) + stride = root + 1 + -- static nodes in ascending (topological) id order, so a single bottom-up + -- sweep fills every column before its parent. Dynamic nodes are omitted + -- entirely, and so are Var leaves (their values are read directly from the + -- run-fixed `vars` columns, see the eval kernels): their static slots were + -- zeros that evalGrad*/forwardPassRange never read, so the flat array + -- shrinks from stride * m to #static * m (a handful of feature/const + -- columns per tree instead of all nodes). + staticKeys = [k | k <- [0 .. root], not (VU.unsafeIndex dynArr k), not (isVarLeaf k)] + nStatic = length staticKeys + isVarLeaf k = case VB.unsafeIndex nodes k of { Var _ -> True; _ -> False } + -- id -> static slot base (slot * m); 0 for dynamic ids (never read) + staticBaseArr = VU.create $ do + arr <- VUM.replicate (root + 1) 0 + forM_ (zip staticKeys [0 ..]) $ \(k, slot) -> + VUM.write arr k (slot * m) + pure arr + -- flat [slot * m + row]; computed in a single bottom-up sweep over the + -- static ids (a child always gets a smaller id than its parent, since + -- cataM assigns the id only after both children are built), writing each + -- static node's column directly into the flat array. This fuses the old + -- per-node VU.map/VU.zipWith intermediates into the array. + staticArr = VU.create $ do + arr <- VUM.replicate (nStatic * m) 0 + let slice slot = VUM.slice (slot * m) m arr + slotOf k = VU.unsafeIndex staticBaseArr k `div` m + -- Resolve a static child @c@ to its source column ONCE per + -- column (hoisted out of the row loop): Var leaves are not + -- materialized in staticArr, so their column is the run-fixed + -- `vars` vector; every other static node is a column already + -- written into the (mutable) arr (children always have smaller + -- ids). `Left` = pure vector (Var leaf), `Right` = mutable slice. + staticSrc c + | isVarLeaf c = Left (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex argArr c))) + | otherwise = Right (slice (slotOf c)) + -- Read row @i from a hoisted source (see staticSrc). Called per + -- element, but the Left/Right tag is fixed per column, so GHC + -- keeps the dispatch cheap and no slice/leaf lookup is repeated. + readSrc s i = case s of + Left v -> pure (VU.unsafeIndex v i) + Right m -> VUM.unsafeRead m i + mapStatic f t k = go 0 + where + dst = slice (slotOf k) + src = staticSrc t + go !i | i >= m = pure () + | otherwise = do + x <- readSrc src i + VUM.unsafeWrite dst i (evalFun f x) + go (i + 1) + zipStatic op l r k = go 0 + where + dst = slice (slotOf k) + srcL = staticSrc l + srcR = staticSrc r + go !i | i >= m = pure () + | otherwise = do + xl <- readSrc srcL i + xr <- readSrc srcR i + VUM.unsafeWrite dst i (evalOp op xl xr) + go (i + 1) + forM_ (zip staticKeys [0 ..]) $ \(k, slot) -> + case VB.unsafeIndex nodes k of + -- Var leaves are excluded from staticKeys (their columns live + -- in `vars`), so they never reach this sweep. + Const v -> VUM.set (slice slot) v + Uni f t -> mapStatic f t k + Bin op l r -> zipStatic op l r k + Param _ -> pure () + Var _ -> pure () + pure arr + + -- compact unboxed per-id code arrays (length root+1) so the hot row loop + -- never touches the boxed `nodes` vector nor dispatches through the + -- function-returning evalOp/evalFun + kindArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Var _ -> 0 + Param _ -> 1 + Const _ -> 2 + Uni _ _ -> 3 + Bin _ _ _ -> 4 + argArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Var ix -> ix + Param ix -> ix + Uni _ t -> t + Bin _ l _ -> l + Const _ -> 0 + arg2Arr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Bin _ _ r -> r + _ -> 0 + fcodeArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Uni f _ -> fromEnum f + _ -> 0 + ocodeArr = VU.generate (root + 1) $ \k -> case int2key IntMap.! k of + Bin op _ _ -> fromEnum op + _ -> 0 + + leftToRight (Uni f mt) = Uni f <$> mt + leftToRight (Bin f ml mr) = Bin f <$> ml <*> mr + leftToRight (Var ix) = pure (Var ix) + leftToRight (Param ix) = pure (Param ix) + leftToRight (Const c) = pure (Const c) + + alg = insertKey + + graph (a, _, _, _) = a + isDynSt k (_, _, d, _) = d IntMap.! k + + insEntry key isD (a, b, d, c) = + ( Map.insert key c a + , IntMap.insert c key b + , IntMap.insert c isD d + , c + 1 ) + + -- a node depends on theta iff it IS a Param, or any child does + nodeIsDynamic (Param _) = pure True + nodeIsDynamic (Var _) = pure False + nodeIsDynamic (Const _) = pure False + nodeIsDynamic (Uni _ t) = gets (isDynSt t) + nodeIsDynamic (Bin _ l r) = (||) <$> gets (isDynSt l) <*> gets (isDynSt r) + + -- Data.Map is unreliable with NaN-valued keys (Ord Double is not a valid + -- total order for NaN: insert(Const NaN) then member/lookup can disagree), + -- and eqsat constant folding can yield Const NaN nodes. So do a single + -- direct lookup; on a miss, return the fresh id that insEntry assigns + -- instead of looking the key back up. Repeated NaN nodes simply get + -- separate ids (no CSE), which is harmless since their static value is + -- recomputed identically. + insertKey key = do + cached <- gets (Map.lookup key . graph) + case cached of + Just v -> pure v + Nothing -> do + d <- nodeIsDynamic key + fresh <- state $ \st@(_, _, _, c) -> let st' = insEntry key d st in (c, st') + pure fresh + +-- Rewrite (a) Bin Power t (Const 2.0) into the unary Square kernel and +-- (b) Bin Div t (Const c) into Bin Mul t (Const (1/c)). Both are exact at +-- the Double level (x ** 2.0 == x * x; x / c == x * (1/c) up to one ulp) +-- and replace the slow per-element pow()/div with a multiply. The loss +-- wrap ((tree - y) ** 2) / m appears in every NLopt objective/gradient +-- call, so these two rewrites are worth a measurable fraction of the AD +-- time. +rewritePowSq :: Fix SRTree -> Fix SRTree +rewritePowSq = cata alg + where + alg :: SRTree (Fix SRTree) -> Fix SRTree + alg (Bin Power t (Fix (Const 2.0))) = Fix (Uni Square t) + alg (Bin Div t (Fix (Const c))) | c /= 0 = Fix (Bin Mul t (Fix (Const (recip c)))) + alg n = Fix n + +-- --------------------------------------------------------------------- +-- Static-child source resolution. A static Var leaf is NOT materialized +-- into ctStatic anymore: its value column lives in ctVars (= xss ++ [y, +-- yErr]) and is read directly at the absolute row (base 0, so the chunk +-- start s0 positions the read). Every other static node is a computed +-- column inside ctStatic at ctStaticBase k. +-- --------------------------------------------------------------------- + +-- | Map a Var leaf's arg (feature ix, or -1 = y, -2 = yErr) to an index +-- into ctVars = xss ++ [y, yErr]. +leafSrcIdx :: Int -> Int -> Int +leafSrcIdx nFeats a | a >= 0 = a + | a == -1 = nFeats + | otherwise = nFeats + 1 +{-# INLINE leafSrcIdx #-} + +-- | Resolve the (source vector, base) of a static child node @k@, where the +-- row value is read at @src (base + i)@. +resolveStatic :: VU.Vector Double + -> VB.Vector (VU.Vector Double) + -> VU.Vector Int + -> VU.Vector Int + -> VU.Vector Int + -> Int -> Int -> Int + -> (VU.Vector Double, Int) +resolveStatic static vars kind arg staticBase nFeats s0 k = + if VU.unsafeIndex kind k == 0 + then (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex arg k)), s0) + else (static, VU.unsafeIndex staticBase k + s0) +{-# INLINE resolveStatic #-} + +-- --------------------------------------------------------------------- +-- Per-theta evaluation: the hot path, called once per NLopt objective/ +-- gradient call. Forward pass only recomputes dynamic nodes (ids are +-- already topologically ordered, so a single left-to-right fold works). +-- Backward pass is the same recursive shape as the original calcGrad, +-- except it stops immediately on any non-dynamic node -- that subtree +-- has no Param in it, so it can never contribute to the gradient. +-- --------------------------------------------------------------------- + +-- Row-fused evaluation: instead of storing one full length-m array per +-- node (which meant ~2 * #nodes large allocations per objective/gradient +-- call), we walk the m data rows one at a time and, for each row, run the +-- forward pass and the reverse-mode backward pass over small per-node +-- scratch arrays of Double (length root+1). This mirrors what the fused +-- Accelerate/LLVM kernel does (one pass per row, no big intermediate +-- arrays) while staying in plain ST: allocation drops from O(nodes * m) +-- to O(nodes + params), and the tight inner loops are all unboxed. +evalGrad :: CompiledTree -> V.Vector Double -> (Double, V.Vector Double) +evalGrad ct theta = runST $ do + fwd <- VUM.new (root + 1) -- node id -> forward value, current row + adj <- VUM.new (root + 1) -- node id -> adjoint (dL/dnode), current row + gradM <- VUM.replicate p 0 -- accumulated per-parameter gradient + objRef <- newSTRef 0 + + let -- forward pass for a single row: fills `fwd` for ids 0..root + forwardLoop !row !key + | key > root = pure () + | otherwise = do + v <- if not (VU.unsafeIndex dyn key) + then if VU.unsafeIndex kind key == 0 + then pure (VU.unsafeIndex (VB.unsafeIndex vars (leafSrcIdx nFeats (VU.unsafeIndex arg key))) row) + else pure (VU.unsafeIndex static (VU.unsafeIndex staticBase key + row)) + else case VU.unsafeIndex kind key of + 1 -> pure (V.unsafeIndex theta (VU.unsafeIndex arg key)) + 3 -> do x <- VUM.unsafeRead fwd (VU.unsafeIndex arg key) + pure (evalFunCode (VU.unsafeIndex fcode key) x) + 4 -> do xl <- VUM.unsafeRead fwd (VU.unsafeIndex arg key) + xr <- VUM.unsafeRead fwd (VU.unsafeIndex arg2 key) + pure (evalOpCode (VU.unsafeIndex ocode key) xl xr) + _ -> error "evalGrad: unreachable" + VUM.unsafeWrite fwd key v + forwardLoop row (key + 1) + + -- backward pass for a single row: ids are visited from root down + -- to 0, which is a valid reverse-topological order since every + -- child id is smaller than its parent's id by construction. + backwardLoop !key + | key < 0 = pure () + | otherwise = do + when (VU.unsafeIndex dyn key) $ do + v <- VUM.unsafeRead adj key + case VU.unsafeIndex kind key of + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + xl <- VUM.unsafeRead fwd l + xr <- VUM.unsafeRead fwd r + fg <- VUM.unsafeRead fwd key + let (dl, dr) = diffScalarCode (VU.unsafeIndex ocode key) v xl xr fg + VUM.unsafeModify adj (+ dl) l + VUM.unsafeModify adj (+ dr) r + 3 -> do + let t = VU.unsafeIndex arg key + x <- VUM.unsafeRead fwd t + VUM.unsafeModify adj (+ v * derivFunCode (VU.unsafeIndex fcode key) x) t + 1 -> VUM.unsafeModify gradM (+ v) (VU.unsafeIndex arg key) + _ -> pure () + backwardLoop (key - 1) + + rowLoop !row + | row >= m = pure () + | otherwise = do + forwardLoop row 0 + rootVal <- VUM.unsafeRead fwd root + modifySTRef' objRef (+ rootVal) + when (VU.unsafeIndex dyn root) $ do + VUM.set adj 0 + VUM.unsafeWrite adj root 1 + backwardLoop root + rowLoop (row + 1) + + rowLoop 0 + + obj <- readSTRef objRef + gradFrozen <- VU.unsafeFreeze gradM + pure (obj, V.convert gradFrozen) + where + root = ctRoot ct + m = ctM ct + p = V.length theta + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + +-- --------------------------------------------------------------------- +-- Node-outer (vectorized-over-rows) evaluation: mirrors reverseModeGraph's +-- shape (one full length-m column per node, node-major loops) so the inner +-- loops are fused per node over all rows, with the static/dynamic pattern +-- decided once per node instead of once per row. Uses the same flat +-- [staticSlot * m + row] layout and compact op-code dispatch as `evalGrad`, but +-- trades the O(nodes + params) scratch of the row-fused version for the +-- O(nodes * m) fwd/adj columns of the massiv-style whole-column kernel. +evalGradVec :: CompiledTree -> V.Vector Double -> (Double, V.Vector Double) +evalGradVec ct theta = runST $ do + -- Per-chunk buffers of O(stride * chunk) instead of one O(stride * m) + -- allocation per call: the fwd/adj matrices are streamed one chunk of + -- `chunk` rows at a time, so the per-call allocation drops ~m/chunk x + -- (and the working set stays L3-resident). The chunk partition does not + -- change any value: each row is independent, the objective row sums + -- accumulate in order and the gradient accumulates row-sums per chunk. + fwd <- VUM.new (stride * chunk) -- [node * nb + i]; dynamic columns written before read + adj <- VUM.replicate (stride * chunk) 0 -- [node * nb + i] + gradM <- VUM.replicate p 0 + + let go !start !acc + | start >= m = do + gradFrozen <- VU.unsafeFreeze gradM + pure (acc, V.convert gradFrozen) + | otherwise = do + let nb = min chunk (m - start) + s0 = start + forwardPassRange ct theta fwd s0 nb + -- objective contribution = sum over this chunk's rows of root + s <- if VU.unsafeIndex dyn root + then {-# SCC "objSumFwd" #-} sumCol fwd (root * nb) nb + else {-# SCC "objSumStatic" #-} sumStatic (VU.unsafeIndex staticBase root + s0) nb + -- seed the root adjoint: d(obj)/d(root value) = 1 per row + unless (s0 == 0) $ VUM.set adj 0 -- reuse the buffer; keep it clean + when (VU.unsafeIndex dyn root) $ {-# SCC "seedAdj" #-} VUM.set (VUM.slice (root * nb) nb adj) 1 + -- backward: nodes from root down to 0 (valid reverse-topological order) + let goBwd !key + | key < 0 = pure () + | otherwise = do + bwdNode key + goBwd (key - 1) + + bwdNode key + | not (VU.unsafeIndex dyn key) = pure () -- no Param in subtree + | otherwise = case VU.unsafeIndex kind key of + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + oc = VU.unsafeIndex ocode key + dl = VU.unsafeIndex dyn l + dr = VU.unsafeIndex dyn r + kb = key * nb + lb = l * nb + rb = r * nb + case (dl, dr) of + (True, True) -> {-# SCC "bwdBinTT" #-} bwdBin nb fwd adj 0 oc kb lb rb (static, 0) + (True, False) -> {-# SCC "bwdBinTS" #-} bwdBin nb fwd adj 1 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 r) + (False, True) -> {-# SCC "bwdBinST" #-} bwdBin nb fwd adj 2 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 l) + (False, False) -> pure () -- no dynamic children to propagate to + 3 -> do + let t = VU.unsafeIndex arg key + fc = VU.unsafeIndex fcode key + kb = key * nb + tb = t * nb + if VU.unsafeIndex dyn t + then {-# SCC "bwdUni" #-} bwdUni nb fwd adj fc kb tb + else pure () -- static child: no Param below, nothing to accumulate + 1 -> do + let a = VU.unsafeIndex arg key + kb = key * nb + {-# SCC "bwdParam" #-} do + s' <- sumCol adj kb nb + VUM.unsafeModify gradM (+ s') a + _ -> pure () + goBwd root + go (start + nb) (acc + s) + + go 0 0 + where + root = ctRoot ct + m = ctM ct + p = V.length theta + stride = root + 1 + chunk = 1024 + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + + sumCol v vbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = VUM.unsafeRead v (vbase + i) >>= \vv -> go (i + 1) (acc + vv) + sumStatic sbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = go (i + 1) (acc + VU.unsafeIndex static (sbase + i)) + +-- --------------------------------------------------------------------- +-- Forward-only objective evaluation: runs the forward pass and the row +-- sum but skips the adjoint/backward pass. Used where only the objective +-- value is needed (reporting loss / R2 metrics, the validation fitness in +-- the search), avoiding the ~2/3 of evalGradVec's work that computes the +-- gradient. +-- --------------------------------------------------------------------- +-- Chunked loss evaluation: runs the same node-outer forward pass as +-- `evalGradVec` (static columns precomputed in `ctStatic`, op codes +-- dispatched once per node into INLINE kernels) but only over a chunk of +-- `chunk` rows at a time with a per-call buffer of O(stride * chunk) +-- instead of O(stride * m). The chunk partition does not change any value +-- (each row is computed independently, the row sums accumulate in order), +-- but it cuts the per-call allocation ~30x so this is cheap enough for the +-- val-eval hot path that runs once per explored expression. +evalLossVec :: CompiledTree -> V.Vector Double -> Double +evalLossVec ct theta = runST $ do + buf <- VUM.new (stride * chunk) + go buf 0 0 + where + root = ctRoot ct + m = ctM ct + stride = root + 1 + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + chunk = 4096 + + go :: VUM.MVector s Double -> Int -> Double -> ST s Double + go buf !start !acc + | start >= m = pure acc + | otherwise = do + let nb = min chunk (m - start) + forwardPassRange ct theta buf start nb + s <- if VU.unsafeIndex dyn root + then sumCol buf (root * nb) nb + else sumStatic (VU.unsafeIndex staticBase root + start) nb + go buf (start + nb) (acc + s) + + sumCol buf vbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = VUM.unsafeRead buf (vbase + i) >>= \vv -> go (i + 1) (acc + vv) + sumStatic sbase !n = go 0 0 + where go !i !acc | i >= n = pure acc + | otherwise = go (i + 1) (acc + VU.unsafeIndex static (sbase + i)) + +-- Forward pass shared by evalGradVec and evalLossVec: fills the `fwd` +-- columns of every dynamic node (ids are topologically ordered, so one +-- left-to-right sweep computes all of them; static columns are already in +-- `ctStatic`). The op/function codes are dispatched once per node and the +-- INLINE loop helpers run a tight fused kernel over the rows. +-- +-- `s0`/`nb` select a range of rows [start, start+nb): with nb = m, start = 0 +-- this is the full-matrix pass used by evalGradVec; evalLossVec calls it on +-- chunks of rows with a stride*nb buffer. The fwd buffer is indexed +-- [key * nb + i], static columns are read at [slot(key) * m + s0 + i]. +forwardPassRange :: CompiledTree -> V.Vector Double -> VUM.MVector s Double -> Int -> Int -> ST s () +forwardPassRange ct theta fwd s0 nb = goFwd 0 + where + root = ctRoot ct + m = ctM ct + kind = ctKind ct + arg = ctArg ct + arg2 = ctArg2 ct + fcode = ctFcode ct + ocode = ctOcode ct + dyn = ctDyn ct + static = ctStatic ct + staticBase = ctStaticBase ct + vars = ctVars ct + nFeats = VB.length vars - 2 + + goFwd !key + | key > root = pure () + | otherwise = do + if VU.unsafeIndex dyn key + then case VU.unsafeIndex kind key of + 1 -> {-# SCC "fwdParam" #-} VUM.set (VUM.slice (key * nb) nb fwd) (V.unsafeIndex theta (VU.unsafeIndex arg key)) + 3 -> do + let t = VU.unsafeIndex arg key + fc = VU.unsafeIndex fcode key + kb = key * nb + tb = t * nb + if VU.unsafeIndex dyn t + then {-# SCC "fwdUniD" #-} fwdUniD nb fwd fc kb tb + else pure () -- a dynamic Uni always has a dynamic child + 4 -> do + let l = VU.unsafeIndex arg key + r = VU.unsafeIndex arg2 key + oc = VU.unsafeIndex ocode key + dl = VU.unsafeIndex dyn l + dr = VU.unsafeIndex dyn r + kb = key * nb + lb = l * nb + rb = r * nb + case (dl, dr) of + (True, True) -> {-# SCC "fwdBinTT" #-} fwdBin nb fwd 0 oc kb lb rb (static, 0) + (True, False) -> {-# SCC "fwdBinTS" #-} fwdBin nb fwd 1 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 r) + (False, True) -> {-# SCC "fwdBinST" #-} fwdBin nb fwd 2 oc kb lb rb (resolveStatic static vars kind arg staticBase nFeats s0 l) + (False, False) -> pure () -- unreachable: a dynamic Bin always has a dynamic child + _ -> pure () + else pure () -- static node: column already in `static` + goFwd (key + 1) + + -- Forward binary kernels: `combo` 0=TT, 1=TS, 2=ST (SS is unreachable + -- for dynamic nodes). The opcode is dispatched ONCE per node; the loop + -- helpers are INLINE with the literal operator so each row iteration + -- is a tight fused kernel with no per-element `case oc of` dispatch. + -- `stSrc` is the (source vector, base) of the static child (either a + -- run-fixed leaf column from `vars` or a computed column of `static`), + -- used by the TS/ST variants; `nb` is the number of rows in this chunk. +fwdBin :: Int -> VUM.MVector s Double -> Int -> Int -> Int -> Int -> Int -> (VU.Vector Double, Int) -> ST s () +fwdBin nb fwd combo oc kb lb rb stSrc = case (combo, oc) of + (0, 0) -> fwdTT nb fwd (+) kb lb rb + (0, 1) -> fwdTT nb fwd (-) kb lb rb + (0, 2) -> fwdTT nb fwd (*) kb lb rb + (0, 3) -> fwdTT nb fwd (/) kb lb rb + (0, 4) -> fwdTT nb fwd (**) kb lb rb + (0, 5) -> fwdTT nb fwd (\l r -> abs l ** r) kb lb rb + (0, 6) -> fwdTT nb fwd (\l r -> l / sqrt (1 + r * r)) kb lb rb + (1, 0) -> fwdTS nb stSrc fwd (+) kb lb + (1, 1) -> fwdTS nb stSrc fwd (-) kb lb + (1, 2) -> fwdTS nb stSrc fwd (*) kb lb + (1, 3) -> fwdTS nb stSrc fwd (/) kb lb + (1, 4) -> fwdTS nb stSrc fwd (**) kb lb + (1, 5) -> fwdTS nb stSrc fwd (\l r -> abs l ** r) kb lb + (1, 6) -> fwdTS nb stSrc fwd (\l r -> l / sqrt (1 + r * r)) kb lb + (2, 0) -> fwdST nb stSrc fwd (+) kb rb + (2, 1) -> fwdST nb stSrc fwd (-) kb rb + (2, 2) -> fwdST nb stSrc fwd (*) kb rb + (2, 3) -> fwdST nb stSrc fwd (/) kb rb + (2, 4) -> fwdST nb stSrc fwd (**) kb rb + (2, 5) -> fwdST nb stSrc fwd (\l r -> abs l ** r) kb rb + (2, 6) -> fwdST nb stSrc fwd (\l r -> l / sqrt (1 + r * r)) kb rb + _ -> pure () +{-# INLINE fwdBin #-} + +-- Backward binary kernels: same dispatch structure, `diff` is the local +-- (dl/dchild, dr/dchild) rule keyed on the opcode. `nb` is the number of +-- rows in this chunk, `stSrc` is the (source vector, base) of the static +-- child (either a run-fixed leaf column from `vars` or a computed column of +-- `static`), used by the TS/ST variants. +bwdBin :: Int -> VUM.MVector s Double -> VUM.MVector s Double -> Int -> Int -> Int -> Int -> Int -> (VU.Vector Double, Int) -> ST s () +bwdBin nb fwd adj combo oc kb lb rb stSrc = case (combo, oc) of + (0, 0) -> bwdTT nb fwd adj (\dx _ _ _ -> (dx, dx)) kb lb rb + (0, 1) -> bwdTT nb fwd adj (\dx _ _ _ -> (dx, negate dx)) kb lb rb + (0, 2) -> bwdTT nb fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb lb rb + (0, 3) -> bwdTT nb fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb lb rb + (0, 4) -> bwdTT nb fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb lb rb + (0, 5) -> bwdTT nb fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb rb + (0, 6) -> bwdTT nb fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb lb rb + (1, 0) -> bwdTS nb stSrc fwd adj (\dx _ _ _ -> (dx, dx)) kb lb + (1, 1) -> bwdTS nb stSrc fwd adj (\dx _ _ _ -> (dx, negate dx)) kb lb + (1, 2) -> bwdTS nb stSrc fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb lb + (1, 3) -> bwdTS nb stSrc fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb lb + (1, 4) -> bwdTS nb stSrc fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb lb + (1, 5) -> bwdTS nb stSrc fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb + (1, 6) -> bwdTS nb stSrc fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb lb + (2, 0) -> bwdST nb stSrc fwd adj (\dx _ _ _ -> (dx, dx)) kb rb + (2, 1) -> bwdST nb stSrc fwd adj (\dx _ _ _ -> (dx, negate dx)) kb rb + (2, 2) -> bwdST nb stSrc fwd adj (\dx fx gy _ -> (dx * gy, dx * fx)) kb rb + (2, 3) -> bwdST nb stSrc fwd adj (\dx _ gy fg -> (dx / gy, dx * (negate fg / gy))) kb rb + (2, 4) -> bwdST nb stSrc fwd adj (\dx fx gy fg -> (fixNaN (dx * gy * fg / fx), fixNaN (dx * fg * log fx))) kb rb + (2, 5) -> bwdST nb stSrc fwd adj (\dx fx gy fg -> + let v2 = abs fx in (fixNaN (dx * (fx * gy) * fg / (v2 * v2)), fixNaN (dx * fg * log (abs fx)))) kb lb + (2, 6) -> bwdST nb stSrc fwd adj (\dx fx gy _ -> + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy)) kb rb + _ -> pure () +{-# INLINE bwdBin #-} + +fwdTT nb fwd op kb lb rb = forRows nb $ \i -> do + xl <- VUM.unsafeRead fwd (lb + i) + xr <- VUM.unsafeRead fwd (rb + i) + VUM.unsafeWrite fwd (kb + i) (op xl xr) +{-# INLINE fwdTT #-} + +fwdTS nb (src, base) fwd op kb lb = forRows nb $ \i -> do + xl <- VUM.unsafeRead fwd (lb + i) + VUM.unsafeWrite fwd (kb + i) (op xl (VU.unsafeIndex src (base + i))) +{-# INLINE fwdTS #-} + +fwdST nb (src, base) fwd op kb rb = forRows nb $ \i -> do + xr <- VUM.unsafeRead fwd (rb + i) + VUM.unsafeWrite fwd (kb + i) (op (VU.unsafeIndex src (base + i)) xr) +{-# INLINE fwdST #-} + +bwdTT nb fwd adj diff kb lb rb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xl <- VUM.unsafeRead fwd (lb + i) + xr <- VUM.unsafeRead fwd (rb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (gl, gr) = diff v xl xr fg + a <- VUM.unsafeRead adj (lb + i) + VUM.unsafeWrite adj (lb + i) (a + gl) + b <- VUM.unsafeRead adj (rb + i) + VUM.unsafeWrite adj (rb + i) (b + gr) +{-# INLINE bwdTT #-} + +bwdTS nb (src, base) fwd adj diff kb lb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xl <- VUM.unsafeRead fwd (lb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (gl, _) = diff v xl (VU.unsafeIndex src (base + i)) fg + a <- VUM.unsafeRead adj (lb + i) + VUM.unsafeWrite adj (lb + i) (a + gl) +{-# INLINE bwdTS #-} + +bwdST nb (src, base) fwd adj diff kb rb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + xr <- VUM.unsafeRead fwd (rb + i) + fg <- VUM.unsafeRead fwd (kb + i) + let (_, gr) = diff v (VU.unsafeIndex src (base + i)) xr fg + b <- VUM.unsafeRead adj (rb + i) + VUM.unsafeWrite adj (rb + i) (b + gr) +{-# INLINE bwdST #-} + +-- Forward unary kernels (dynamic child): the function code is dispatched +-- ONCE per node and the loop helper is INLINE with the literal function, +-- so each row iteration is a tight fused kernel with no per-element +-- `case fc of` / closure build (a dynamic Uni node always has a dynamic +-- child, so there is no static-child variant here). +fwdUniD :: Int -> VUM.MVector s Double -> Int -> Int -> Int -> ST s () +fwdUniD nb fwd fc kb tb = case fc of + 0 -> fwdUniD' nb fwd (\x -> x) kb tb + 1 -> fwdUniD' nb fwd abs kb tb + 2 -> fwdUniD' nb fwd sin kb tb + 3 -> fwdUniD' nb fwd cos kb tb + 4 -> fwdUniD' nb fwd tan kb tb + 5 -> fwdUniD' nb fwd sinh kb tb + 6 -> fwdUniD' nb fwd cosh kb tb + 7 -> fwdUniD' nb fwd tanh kb tb + 8 -> fwdUniD' nb fwd asin kb tb + 9 -> fwdUniD' nb fwd acos kb tb + 10 -> fwdUniD' nb fwd atan kb tb + 11 -> fwdUniD' nb fwd asinh kb tb + 12 -> fwdUniD' nb fwd acosh kb tb + 13 -> fwdUniD' nb fwd atanh kb tb + 14 -> fwdUniD' nb fwd sqrt kb tb + 15 -> fwdUniD' nb fwd (\x -> sqrt (abs x)) kb tb + 16 -> fwdUniD' nb fwd (\x -> signum x * abs x ** (1 / 3)) kb tb + 17 -> fwdUniD' nb fwd (\x -> x * x) kb tb + 18 -> fwdUniD' nb fwd log kb tb + 19 -> fwdUniD' nb fwd (\x -> log (abs x)) kb tb + 20 -> fwdUniD' nb fwd exp kb tb + 21 -> fwdUniD' nb fwd recip kb tb + 22 -> fwdUniD' nb fwd (\x -> x * x * x) kb tb + _ -> pure () +{-# INLINE fwdUniD #-} + +fwdUniD' nb fwd f kb tb = forRows nb $ \i -> do + x <- VUM.unsafeRead fwd (tb + i) + VUM.unsafeWrite fwd (kb + i) (f x) +{-# INLINE fwdUniD' #-} + +-- Backward unary kernel: derivative of the function, dispatched once per +-- node and inlined into the accumulation loop. `nb` is the number of rows +-- in the current chunk. +bwdUni :: Int -> VUM.MVector s Double -> VUM.MVector s Double -> Int -> Int -> Int -> ST s () +bwdUni nb fwd adj fc kb tb = case fc of + 0 -> bwdUni' nb fwd adj (\_ -> 1) kb tb + 1 -> bwdUni' nb fwd adj (\x -> x / abs x) kb tb + 2 -> bwdUni' nb fwd adj cos kb tb + 3 -> bwdUni' nb fwd adj (negate . sin) kb tb + 4 -> bwdUni' nb fwd adj (\x -> 1 / (cos x * cos x)) kb tb + 5 -> bwdUni' nb fwd adj cosh kb tb + 6 -> bwdUni' nb fwd adj sinh kb tb + 7 -> bwdUni' nb fwd adj (\x -> 1 - tanh x * tanh x) kb tb + 8 -> bwdUni' nb fwd adj (\x -> 1 / sqrt (1 - x * x)) kb tb + 9 -> bwdUni' nb fwd adj (\x -> -1 / sqrt (1 - x * x)) kb tb + 10 -> bwdUni' nb fwd adj (\x -> 1 / (1 + x * x)) kb tb + 11 -> bwdUni' nb fwd adj (\x -> 1 / sqrt (1 + x * x)) kb tb + 12 -> bwdUni' nb fwd adj (\x -> 1 / (sqrt (x - 1) * sqrt (x + 1))) kb tb + 13 -> bwdUni' nb fwd adj (\x -> 1 / (1 - x * x)) kb tb + 14 -> bwdUni' nb fwd adj (\x -> 1 / (2 * sqrt x)) kb tb + 15 -> bwdUni' nb fwd adj (\x -> x / (2 * abs x ** (3 / 2))) kb tb + 16 -> bwdUni' nb fwd adj (\x -> 1 / (3 * (x * x) ** (1 / 3))) kb tb + 17 -> bwdUni' nb fwd adj (\x -> 2 * x) kb tb + 18 -> bwdUni' nb fwd adj recip kb tb + 19 -> bwdUni' nb fwd adj recip kb tb + 20 -> bwdUni' nb fwd adj exp kb tb + 21 -> bwdUni' nb fwd adj (\x -> -1 / (x * x)) kb tb + 22 -> bwdUni' nb fwd adj (\x -> 3 * x * x) kb tb + _ -> pure () +{-# INLINE bwdUni #-} + +bwdUni' nb fwd adj f kb tb = forRows nb $ \i -> do + v <- VUM.unsafeRead adj (kb + i) + x <- VUM.unsafeRead fwd (tb + i) + c <- VUM.unsafeRead adj (tb + i) + VUM.unsafeWrite adj (tb + i) (c + v * f x) +{-# INLINE bwdUni' #-} +-- Unboxed ST loop over the m data rows; always inlined so the per-node +-- bodies above are fused into a single tail-recursive kernel per node. +forRows :: Int -> (Int -> ST s ()) -> ST s () +forRows !n f = go 0 + where + go !i | i >= n = pure () + | otherwise = f i >> go (i + 1) +{-# INLINE forRows #-} +evalOpCode :: Int -> Double -> Double -> Double +evalOpCode 0 = (+) +evalOpCode 1 = (-) +evalOpCode 2 = (*) +evalOpCode 3 = (/) +evalOpCode 4 = (**) +evalOpCode 5 = \l r -> abs l ** r +evalOpCode 6 = \l r -> l / sqrt (1 + r * r) +evalOpCode _ = error "evalOpCode: bad op code" +{-# INLINE evalOpCode #-} + +evalFunCode :: Int -> Double -> Double +evalFunCode 0 = id +evalFunCode 1 = abs +evalFunCode 2 = sin +evalFunCode 3 = cos +evalFunCode 4 = tan +evalFunCode 5 = sinh +evalFunCode 6 = cosh +evalFunCode 7 = tanh +evalFunCode 8 = asin +evalFunCode 9 = acos +evalFunCode 10 = atan +evalFunCode 11 = asinh +evalFunCode 12 = acosh +evalFunCode 13 = atanh +evalFunCode 14 = sqrt +evalFunCode 15 = \x -> sqrt (abs x) +evalFunCode 16 = \x -> signum x * abs x ** (1 / 3) +evalFunCode 17 = \x -> x * x +evalFunCode 18 = log +evalFunCode 19 = \x -> log (abs x) +evalFunCode 20 = exp +evalFunCode 21 = recip +evalFunCode 22 = \x -> x * x * x +evalFunCode _ = error "evalFunCode: bad function code" +{-# INLINE evalFunCode #-} + +derivFunCode :: Int -> Double -> Double +derivFunCode 0 = const 1 +derivFunCode 1 = \x -> x / abs x +derivFunCode 2 = cos +derivFunCode 3 = negate . sin +derivFunCode 4 = \x -> 1 / (cos x * cos x) +derivFunCode 5 = cosh +derivFunCode 6 = sinh +derivFunCode 7 = \x -> 1 - tanh x * tanh x +derivFunCode 8 = \x -> 1 / sqrt (1 - x * x) +derivFunCode 9 = \x -> -1 / sqrt (1 - x * x) +derivFunCode 10 = \x -> 1 / (1 + x * x) +derivFunCode 11 = \x -> 1 / sqrt (1 + x * x) +derivFunCode 12 = \x -> 1 / (sqrt (x - 1) * sqrt (x + 1)) +derivFunCode 13 = \x -> 1 / (1 - x * x) +derivFunCode 14 = \x -> 1 / (2 * sqrt x) +derivFunCode 15 = \x -> x / (2 * abs x ** (3 / 2)) +derivFunCode 16 = \x -> 1 / (3 * (x * x) ** (1 / 3)) +derivFunCode 17 = (* 2) +derivFunCode 18 = recip +derivFunCode 19 = recip +derivFunCode 20 = exp +derivFunCode 21 = \x -> -1 / (x * x) +derivFunCode 22 = \x -> 3 * x * x +derivFunCode _ = error "derivFunCode: bad function code" +{-# INLINE derivFunCode #-} + +-- Pure local-derivative rules keyed on fromEnum Op, scalar version (same +-- math as the original vectorized `diffPure`, applied per-row above). +diffScalarCode :: Int -> Double -> Double -> Double -> Double -> (Double, Double) +diffScalarCode 0 dx _ _ _ = (dx, dx) +diffScalarCode 1 dx _ _ _ = (dx, negate dx) +diffScalarCode 2 dx fx gy _ = (dx * gy, dx * fx) +diffScalarCode 3 dx _ gy fg = (dx / gy, dx * (negate fg / gy)) +diffScalarCode 4 dx fx gy fg = + ( fixNaN (dx * gy * fg / fx) + , fixNaN (dx * fg * log fx) ) +diffScalarCode 5 dx fx gy fg = + let v2 = abs fx + in ( fixNaN (dx * (fx * gy) * fg / (v2 * v2)) + , fixNaN (dx * fg * log (abs fx)) ) +diffScalarCode 6 dx fx gy _ = + let dxl = dx * (recip . sqrt . (+1) . (^(2::Int))) gy + dxy = fx * gy * dxl ^ (3::Int) + in (dxl, dxy) +diffScalarCode _ _ _ _ _ = error "diffScalarCode: bad op code" +{-# INLINE diffScalarCode #-} + +fixNaN :: Double -> Double +fixNaN x = if isNaN x then 0 else x +{-# INLINE fixNaN #-} + +-- --------------------------------------------------------------------- +-- Drop-in-compatible wrapper -- same signature as your original function. +-- Use this ONLY to verify correctness against your existing implementation +-- (e.g. QuickCheck / golden tests comparing outputs). It gets you ZERO +-- speedup on its own, since it calls compileTree fresh every time, same +-- as before. The actual win requires changing the NLopt-facing call site. +-- --------------------------------------------------------------------- + +--reverseModeGraphO :: [V.Vector Double] -> V.Vector Double -> Maybe (V.Vector Double) -> V.Vector Double -> Fix SRTree -> (V.Vector Double, V.Vector Double) +--reverseModeGraphO xss ys mYErr theta tree = evalGrad (compileTree xss ys mYErr tree) theta + +-- | Safely chunk an unboxed vector into 'n' roughly equal parts. +chunkVector :: Int -> VU.Vector Double -> [VU.Vector Double] +chunkVector numChunks v + | VU.null v = [] + | otherwise = + let n = VU.length v + chunkSize = max 1 (n `div` numChunks) + go vec | VU.null vec = [] + | VU.length vec <= chunkSize = [vec] + | otherwise = let (h, t) = VU.splitAt chunkSize vec + in h : go t + in go v + +-- | Compiles the tree for multiple data chunks independently. +compileTreeMulti :: [VU.Vector Double] + -> VU.Vector Double + -> Maybe (VU.Vector Double) + -> Fix SRTree + -> [CompiledTree] +compileTreeMulti xss ys mYErr tree = + let nRows = VU.length ys + minChunkSize = 2000 + numChunks = max 1 (min cap (nRows `div` minChunkSize)) + cap = if mtSingleChunk then 1 else unsafePerformIO getNumCapabilities + ysChunks = chunkVector numChunks ys + -- transpose groups the chunks by slice rather than by feature + xssChunks = Data.List.transpose (map (chunkVector numChunks) xss) + errChunks = case mYErr of + Just e -> map Just (chunkVector numChunks e) + Nothing -> replicate (length ysChunks) Nothing + in [ compileTree xs y err tree | (xs, y, err) <- zip3 xssChunks ysChunks errChunks ] + +-- | When True, the MultiThread backend compiles/evaluates each tree on a +-- single chunk so a higher-level population-parallel driver (eggp's fitness +-- batch) owns the cores instead of oversubscribing the per-tree chunk split. +{-# NOINLINE mtSingleChunk #-} +mtSingleChunk :: Bool +mtSingleChunk = unsafePerformIO (readIORef mtParGate) + +mtParGate :: IORef Bool +mtParGate = unsafePerformIO (newIORef False) +{-# NOINLINE mtParGate #-} + +-- | Enable/disable single-chunk (non-oversubscribing) mode for the MultiThread +-- backend; called around a population-parallel fitness batch. +setMTPopParallel :: Bool -> IO () +setMTPopParallel b = writeIORef mtParGate b + +-- | Evaluates the gradient across all compiled chunks in parallel. +-- Each chunk is evaluated by the fast node-outer `evalGradVec` kernel on +-- its own slice of the data. The kernel is now chunked internally (O(stride +-- * 1024) per-call buffers, L3-resident) and is compute-bound rather than +-- memory-bandwidth-bound, so splitting the data into one chunk per core and +-- running the kernels concurrently scales almost linearly. The objective +-- and gradient accumulate across chunks (same math per row; only the FP +-- summation order across chunk boundaries differs). +evalGradMulti :: [CompiledTree] -> V.Vector Double -> (Double, V.Vector Double) +evalGradMulti [ct] theta = evalGradVec ct theta +evalGradMulti cts theta = unsafePerformIO $ do + results <- forConcurrently cts $ \ct -> evaluate (evalGradVec ct theta) + let totalObj = sum $ map fst results + totalGrad = foldl1' (V.zipWith (+)) (map snd results) + pure (totalObj, totalGrad) diff --git a/src/Algorithm/SRTree/Compile.hs b/src/Algorithm/SRTree/Compile.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Compile.hs @@ -0,0 +1,106 @@ +{-# LANGUAGE GADTs #-} + +module Algorithm.SRTree.Compile where + +import Data.SRTree +import Data.SRTree.Eval (compileLoss, Target, Columns, Theta) +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Generic as G +import Algorithm.SRTree.AD +import Algorithm.SRTree.Utils +import Algorithm.SRTree.Likelihoods (Distribution(..), Loss(..), buildLoss, hessianNLL) +import Algorithm.SRTree.NonlinearOpt (minimizeNLL, minimizeNLLWithFixedParam) +import Data.SRTree.Recursion (cata) + +data EvalTree = EvalTree { + ctDist :: Distribution, + ctLoss :: Theta -> Double, + ctAD :: VS.Vector Double -> (Double, VS.Vector Double), + ctOptimizer :: Target -> Target, + ctOptimizerFixed :: Int -> Target -> Target, + ctNLL :: Target -> Double, + ctGradNLL :: Target -> (Double, Target), + ctHessianNLL :: Target -> Columns, + ctTree :: Fix SRTree, + ctRows :: Int, + ctVar :: Double +} + +-- | Compile a tree and store it in a CompiledTree data structure +compileTree :: Distribution -> Columns -> Target -> Maybe Target -> Fix SRTree -> EvalTree +compileTree dist xss ys mYerr tree = EvalTree { + ctDist = dist, + ctLoss = compileLoss xss tree ys mYerr, + ctAD = compileFunAndGrad MultiThread xss ys mYerr tree, + ctOptimizer = fst3 . minimizeNLL MultiThread (NLL dist) mYerr 100 xss ys tree, + ctOptimizerFixed = minimizeNLLWithFixedParam MultiThread (NLL dist) mYerr 100 xss ys tree, + ctNLL = compileLoss xss lossTree ys mYerr, + ctGradNLL = \theta -> let fg = compileFunAndGrad MultiThread xss ys mYerr lossTree + (obj, gradStorable) = fg (G.convert theta) + in (obj, G.convert gradStorable), + ctHessianNLL = hessianNLL dist mYerr xss ys tree, + ctTree = tree, + ctRows = n, + ctVar = let ym = U.sum ys / fromIntegral n + in U.foldr (\yi acc -> acc + (yi - ym)^2) 0 ys +} + where + n = U.length ys + lossTree = buildLoss (NLL dist) (fromIntegral n) tree + fst3 (a, _, _) = a + +data EvaluatedTree = EvaluatedTree { + valLoss :: Double, + valTheta :: Theta, + valRows :: Double, + valParams :: Double, + valTree :: Fix SRTree, + valLogParams :: Double, + valLogParamsLattice :: Double, + valVar :: Double +} + +evaluateTree :: EvalTree -> Target -> [[Double]] -> Theta -> EvaluatedTree +evaluateTree et fisher hessian theta = EvaluatedTree { + valLoss = ctLoss et theta, + valTheta = theta, + valRows = fromIntegral (ctRows et), + valParams = fromIntegral (U.length theta), + valTree = ctTree et, + valLogParams = logParameters fisher theta, + valLogParamsLattice = logParametersLatt hessian fisher theta, + valVar = ctVar et +} + + +-- log of the parameters complexity +logParameters :: U.Vector Double -> Target -> Double +logParameters fisher theta = -(p / 2) * log 3 + 0.5 * logFisher + logTheta + where + (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (U.toList theta) (U.toList fisher) + +-- same as above but for the Lattice +logParametersLatt :: [[Double]] -> U.Vector Double -> Target -> Double +logParametersLatt hessian fisher theta = 0.5 * p * (1 - log 3) + 0.5 * log detFisher + where + detFisher = det $ map U.fromList hessian + + (logTheta, logFisher, p) = foldr addIfSignificant (0, 0, 0) $ zip (U.toList theta) (U.toList fisher) + +addIfSignificant (v, f) (acc_v, acc_f, acc_p) + | isSignificant v f = (acc_v + log (abs v), acc_f + log f, acc_p + 1) + | otherwise = (acc_v, acc_f, acc_p) +{-# INLINE addIfSignificant #-} + +isSignificant v f = abs (v / sqrt(12 / f) ) >= 1 +{-# INLINE isSignificant #-} + +fixParam :: Int -> Double -> Fix SRTree -> Fix SRTree +fixParam ix val = cata alg + where + alg (Param i) | i == ix = Fix $ Const val + | i > ix = Fix $ Param (i-1) + | otherwise = Fix $ Param i + alg other = Fix other +{-# INLINE fixParam #-} diff --git a/src/Algorithm/SRTree/ConfidenceIntervals.hs b/src/Algorithm/SRTree/ConfidenceIntervals.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/ConfidenceIntervals.hs @@ -0,0 +1,422 @@ +{-# language ViewPatterns, ScopedTypeVariables, MultiWayIf, FlexibleContexts #-} +------------------------------------------------------------------------------- +-- | +-- Module : Algorithm.SRTree.ConfidenceIntervals +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to optimize the parameters of an expression. +------------------------------------------------------------------------------- +module Algorithm.SRTree.ConfidenceIntervals where + +import Statistics.Distribution ( ContDistr(quantile) ) +import Statistics.Distribution.StudentT ( studentT ) +import Statistics.Distribution.FDistribution ( fDistribution ) +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Generic as G +import Data.SRTree +import Data.SRTree.Eval +import Data.SRTree.Recursion ( cata ) +import Algorithm.SRTree.Likelihoods +import Algorithm.SRTree.Compile +import Data.List ( sortOn, nubBy ) +import Data.Maybe ( listToMaybe ) +import Algorithm.SRTree.Utils +import Numeric.Optimization.NLOPT +import System.IO.Unsafe ( unsafePerformIO ) +import Control.Monad.Catch ( catch, SomeException ) + +import Debug.Trace ( trace ) + +-- | profile likelihood algorithms: Bates (classical), ODE (faster), Constrained (fastest) +-- The Constrained approach returns only the endpoints. +data PType = Bates | ODE | Constrained deriving (Show, Read, Eq) + +-- | Confidence Interval using Laplace approximation or profile likelihood. +data CIType = Laplace BasicStats | Profile BasicStats [ProfileT] + +-- | Basic stats of the data: covariance of parameters, correlation, standard errors +data BasicStats = MkStats + { _cov :: Columns + , _corr :: Columns + , _stdErr :: Target + } deriving (Eq, Show) + +-- | a confience interval is composed of the point estimate (`est_`), lower bound (`_lower_`) +-- and upper bound (`upper_`) +data CI = CI + { est_ :: Double + , lower_ :: Double + , upper_ :: Double + } deriving (Eq, Show, Read) + +-- | A profile likelihood is composed of a vector of tau values that traces the likelihood, +-- the matrix of thetas for each profile, the local optima, and two splines that converts +-- taus to theta and vice-versa. +data ProfileT = ProfileT + { _taus :: Target + , _thetas :: Columns + , _opt :: Double + , _tau2theta :: Double -> Double + , _theta2tau :: Double -> Double + } + +-- shows the CI with n places +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 -> Target -> Double -> [CI] +paramCI (Laplace stats) nSamples theta alpha = zipWith3 CI (U.toList theta) lows highs + where + -- the Laplace approximation is theta +/- t(1-alpha/2) * standard error + k = U.length theta + t = quantile (studentT . fromIntegral $ nSamples - k) (1 - alpha / 2.0) + stdErr = _stdErr stats + lows = U.toList $ U.zipWith (-) theta $ U.map (*t) stdErr + highs = U.toList $ U.zipWith (+) theta $ U.map (*t) stdErr + +paramCI (Profile stats profiles) nSamples _ alpha = zipWith3 CI theta lows highs + where + -- for the profile likelihood we use the square root of the F-distribution with (1-alpha) + k = length theta + t = sqrt $ quantile (fDistribution k (fromIntegral $ nSamples - k)) (1 - alpha) + stdErr = _stdErr stats + lows = map (`_tau2theta` (-t)) profiles + highs = map (`_tau2theta` t) profiles + theta = map _opt profiles + +-- | calculates the prediction confidence interval using Laplace approximation or profile likelihood. +-- predictionCI +predictionCI :: CIType -> Distribution -> (Columns -> Target) -> (Columns -> [Target]) -> (CI -> Target -> Fix SRTree -> (Double -> Double, Double)) -> Columns -> Fix SRTree -> Target -> Double -> [CI] -> [CI] +predictionCI (Laplace stats) _ predFun jacFun _ xss tree theta alpha _ = zipWith3 CI yhat lows highs + where + yhat = U.toList $ predFun xss + jac' = jacFun xss + k = U.length theta + n = length yhat + t = quantile (studentT . fromIntegral $ n - k) (1 - alpha / 2.0) + + covMat = toRowMajor (_cov stats) + nCov = k - 1 + + lows = zipWith (-) yhat $ map (*t) resStdErr + highs = zipWith (+) yhat $ map (*t) resStdErr + + getResStdError row = + sqrt $ U.sum $ U.generate nCov $ \i -> + (row U.! i) * U.sum (U.zipWith (*) row (U.slice (i * k) nCov covMat)) + resStdErr = map (getResStdError . U.slice 0 nCov) (getRows jac') + +predictionCI (Profile _ _) dist predFun _ profFun xss tree theta alpha estPIs = zipWith3 f estPIs yhat xss' + where + yhat = U.toList $ predFun xss + k = U.length theta + n = length yhat + t = sqrt $ quantile (fDistribution k (fromIntegral $ n - k)) (1 - alpha) + + theta0 = calcTheta0 dist tree + xss' = getRows xss + + f estPI yh xs = let + t' = replaceParam0 tree $ evalVar xs theta0 + (spline, yh') = profFun estPI (theta U.// [(0, yh)]) t' + in CI yh' (spline (-t)) (spline t) + +-- inverse function of the distributions +inverseDist :: Floating p => Distribution -> p -> p +inverseDist Gaussian y = y +inverseDist Bernoulli y = log (y/(1-y)) +inverseDist Poisson y = log y +inverseDist _ y = 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 (Y ix) = Fix $ Y ix + alg (Uni g t) = Fix $ Uni g t + alg (Bin op l r) = Fix $ Bin op l r + +evalVar :: Target -> Fix SRTree -> Fix SRTree +evalVar xs = cata alg + where + alg (Var ix) = Fix $ Const (xs U.! ix) + alg (Param ix) = Fix $ Param ix + alg (Const c) = Fix $ Const c + alg (Y ix) = Fix $ Y ix + alg (Uni g t) = Fix $ Uni g t + alg (Bin op l r) = Fix $ Bin op l r + +calcTheta0 :: Distribution -> Fix SRTree -> Fix SRTree +calcTheta0 dist tree = case cata alg tree of + Left g -> g $ inverseDist dist (Fix $ Param 0) + Right _ -> error "No theta0?" + 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 (Y ix) = Right $ Fix $ Y ix + alg (Uni g t) = case t of + Left f -> Left $ f . evalInverse g + Right v -> Right $ evalFun g v + alg (Bin op l r) = case l of + Left f -> case r of + Left _ -> error "This shouldn't happen!" + Right v -> Left $ f . invright op v + Right vl -> case r of + Left g -> Left $ g . invleft op vl + Right vr -> Right $ evalOp op vl vr + +-- calculate the profile likelihood of every parameter +getAllProfiles :: PType -> EvalTree -> Target -> Target -> [CI] -> Double -> [ProfileT] +getAllProfiles ptype et theta stdErr estCIs alpha = getAll 0 [] + where + k = U.length theta + n = ctRows et + tau_max = sqrt $ quantile (fDistribution k (n - k)) (1 - 0.01) + tau_max' = sqrt $ quantile (fDistribution k (n - k)) (1 - alpha) + + profFun ix = case ptype of + Bates -> getProfile et theta (stdErr U.! ix) tau_max ix + ODE -> getProfileODE et theta (stdErr U.! ix) (estCIs !! ix) tau_max ix + Constrained -> getProfileCnstr et theta (stdErr U.! ix) tau_max' ix + + getAll ix acc | ix == k = acc + | ix == k-1 && ptype == Constrained && ctDist et == Gaussian = case getProfileODE et theta (stdErr U.! ix) (estCIs !! ix) tau_max ix of + Left t -> getAllProfiles ptype et t stdErr estCIs alpha + Right p -> getAll (ix + 1) (acc <> [p]) + | otherwise = case profFun ix of + Left t -> getAllProfiles ptype et t stdErr estCIs alpha + Right p -> getAll (ix + 1) (acc <> [p]) + +-- calculates the profile likelihood of a single parameter +getProfile :: EvalTree -> Target -> Double -> Double -> Int -> Either Target ProfileT +getProfile et theta stdErr_i tau_max ix + | stdErr_i == 0.0 = pure $ ProfileT (U.fromList [-tau_max, tau_max]) [theta, theta] (theta U.! ix) (const (theta U.! ix)) (const tau_max) + | otherwise = + do negDelta <- go kmax (-stdErr_i / 8) 0 1 mempty + posDelta <- go kmax (stdErr_i / 8) 0 1 p0 + let (taus', thetas') = negDelta <> posDelta + taus = U.fromList taus' + thetas = thetas' + (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix + pure $ ProfileT taus thetas optTh tau2theta theta2tau + where + p0 = ([0], [theta_opt]) + kmax = 300 + nll_opt = ctNLL et theta_opt + theta_opt = ctOptimizer et theta + optTh = theta_opt U.! ix + minimizer = ctOptimizerFixed et ix + + go 0 delta _ _ acc = Right acc + go k delta t inv_slope acc@(taus, thetas) + | isNaN inv_slope = Right acc + | nll_cond < nll_opt = Left theta_t + | abs tau > tau_max = Right acc' + + | otherwise = go (k-1) delta (t + inv_slope) inv_slope' acc' + where + t_delta = (theta_opt U.! ix) + delta * (t + inv_slope) + theta_delta = updateS theta_opt [(ix, t_delta)] + theta_t = minimizer theta_delta + (nll_cond, grad) = ctGradNLL et theta_t + zv = grad U.! ix + inv_slope' = min 4.0 . max 0.0625 . abs $ (tau / (stdErr_i * zv)) + tau = signum delta * sqrt (2*nll_cond - 2*nll_opt) + acc' = if nll_cond == nll_opt || maybe False (tau ==) (listToMaybe taus) || isNaN tau + then acc + else (tau:taus, theta_t:thetas) + +-- 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 :: EvalTree -> Target -> Double -> Double -> Int -> Either Target ProfileT +getProfileCnstr et theta stdErr_i tau_max ix + | stdErr_i == 0.0 = pure $ ProfileT taus thetas theta_i (const theta_i) (const tau_max) + | otherwise = pure $ ProfileT taus thetas theta_i tau2theta (const tau_max) + where + taus = U.fromList [-tau_max, tau_max] + thetas = [theta, theta] + theta_i = theta U.! ix + getPoint = getEndPoint et theta tau_max ix + leftPt = getPoint True + rightPt = getPoint False + tau2theta tau = if tau < 0 then leftPt else rightPt + +getEndPoint :: EvalTree -> Target -> Double -> Int -> Bool -> Double +getEndPoint et theta tau_max ix isLeft = + case minimizeAugLag problem (G.convert theta_opt) of + Right sol -> solutionParams sol VS.! ix + Left _ -> theta_opt U.! ix + where + n = U.length theta + + theta_opt = ctOptimizer et theta + nll_opt = ctNLL et theta_opt + loss_crit = nll_opt + tau_max + + loss = subtract loss_crit . ctNLL et . G.convert + obj = (if isLeft then id else negate) . (VS.! ix) + + stop = ObjectiveRelativeTolerance 1e-4 :| [MaximumEvaluations 1000] + 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 :: EvalTree -> Target -> Double -> CI -> Double -> Int -> Either Target ProfileT +getProfileODE et theta stdErr_i estCI tau_max ix + | stdErr_i == 0.0 = pure dflt + | otherwise = let (taus', thetas') = solLeft <> ([0], [theta_opt]) <> solRight + taus = U.fromList taus' + thetas = thetas' + (tau2theta, theta2tau) = createSplines taus thetas stdErr_i tau_max ix + in pure $ ProfileT taus thetas optTh tau2theta theta2tau + where + dflt = ProfileT (U.fromList [-tau_max, tau_max]) [theta, theta] (theta U.! ix) (const (theta U.! ix)) (const tau_max) + theta_opt = ctOptimizer et theta + grader = snd . ctGradNLL et + nll_opt = ctNLL et theta_opt + optTh = theta_opt U.! ix + p = U.length theta + p' = p + 1 + + odeFun gamma _ u = + let grad = grader u + w = ctHessianNLL et u + m = [ U.generate p' (\i -> + if i < p && j < p then (w !! j) U.! i + else if i == ix || j == ix then 1 + else 0 + ) + | j <- [0 .. p'-1] ] + v = U.snoc (U.map (*(-gamma)) grad) 1 + dotTheta = unsafePerformIO $ luSolve m v + in U.init dotTheta + + minRange = max (abs (upper_ estCI - optTh)) (abs (lower_ estCI - optTh)) + scanRange = max minRange (tau_max * abs stdErr_i) + nPts = max 50 (min 100 (ceiling (scanRange / minRange * 49) + 1)) + tsHi = linSpace nPts (optTh, optTh + scanRange) + tsLo = linSpace nPts (optTh, optTh - scanRange) + scanOn sig = foldMap (calcTau sig) . f . scanl (rk (odeFun sig)) (optTh, theta_opt) + where f = if sig==1 then id else reverse + solRight = scanOn 1 tsHi + solLeft = scanOn (-1) tsLo + calcTau s t = let nll_i = ctNLL et (snd t) + z = signum ((snd t U.! ix) - optTh) * sqrt (2 * nll_i - 2 * nll_opt) + in if z == 0 || isNaN z then ([], []) else ([z], [snd t]) + +rk :: (Double -> Target -> Target) -> (Double, Target) -> Double -> (Double, Target) +rk f (t, y) t' = (t', U.zipWith5 (\y0 k1 k2 k3 k4 -> y0 + h/6 * (k1 + 2*k2 + 2*k3 + k4)) y k1 k2 k3 k4) + where + h = t' - t + k1 = f t y + k2 = f (t + 0.5*h) (U.zipWith (\y0 k -> y0 + 0.5*h*k) y k1) + k3 = f (t + 0.5*h) (U.zipWith (\y0 k -> y0 + 0.5*h*k) y k2) + k4 = f (t + 1.0*h) (U.zipWith (\y0 k -> y0 + 1.0*h*k) y k3) +{-# INLINE rk #-} + +-- tau0, tau1 theta0, thetaX = tau1 theta0 / tau0 +getStatsFromModel :: Distribution -> Maybe Target -> Columns -> Target -> Fix SRTree -> Target -> BasicStats +getStatsFromModel dist mYerr xss ys tree theta = MkStats cov corr stdErr + where + k = U.length theta + n = U.length ys + nParams = fromIntegral k + ident = fromRowMajor k k (U.generate (k * k) (\ix -> let (i, j) = ix `divMod` k in if i == j then 1.0 else 0.0)) + + hess = hessianNLL dist mYerr xss ys tree theta + + fexcept :: SomeException -> IO Columns + fexcept e = trace ("cov NegDef" <> show (toRowMajor hess)) $ pure ident + + cov = unsafePerformIO $ catch (invChol hess) fexcept + + covMat = toRowMajor cov + stdErr = U.generate k (\ix -> sqrt $ covMat U.! (ix * k + ix)) + + stdErrSq = case outer stdErr stdErr of + Right v -> v + Left _ -> [] + + stdErrSqMat = toRowMajor stdErrSq + corr = fromRowMajor k k $ U.generate (k * k) (\ix -> covMat U.! ix / stdErrSqMat U.! ix) + +-- Create splines for profile-t +createSplines :: Target -> Columns -> Double -> Double -> Int -> (Double -> Double, Double -> Double) +createSplines taus thetas se tau_max ix + | n < 2 = (genSplineFun [(-tau_max, -se), (tau_max, se)], genSplineFun [(-se, 0), (se, 1)]) + | otherwise = (tau2theta, theta2tau) + where + n = U.length taus + cols = getCol ix thetas + nubOnFirst = nubBy (\x y -> fst x == fst y) + tau2theta = genSplineFun $ nubOnFirst $ sortOnFirst taus cols + theta2tau = genSplineFun $ nubOnFirst $ sortOnFirst cols taus + +getCol :: Int -> Columns -> Target +getCol ix mtx = U.generate (length mtx) (\j -> (mtx !! j) U.! ix) +{-# inline getCol #-} + +sortOnFirst :: Target -> Target -> [(Double, Double)] +sortOnFirst xs ys = sortOn fst $ zip (U.toList xs) (U.toList ys) +{-# inline sortOnFirst #-} + +splinesSketches :: Double -> Target -> Target -> (Double -> Double) -> (Double -> Double) +splinesSketches tauScale (U.toList -> tau) (U.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 + (prof1, prof2) = (profs !! ix1, profs !! ix2) + (tau2theta1, theta2tau1) = (_tau2theta prof1, _theta2tau prof1) + (tau2theta2, theta2tau2) = (_tau2theta prof2, _theta2tau prof2) + + tauScale = sqrt (fromIntegral nParams * quantile (fDistribution nParams (fromIntegral nPoints - fromIntegral nParams)) (1 - alpha)) + splineG1 = splinesSketches tauScale (_taus prof2) (getCol ix1 (_thetas prof2)) theta2tau1 + splineG2 = splinesSketches tauScale (_taus prof1) (getCol ix2 (_thetas prof1)) theta2tau2 + + angles = [ (0, splineG2 1), (splineG1 1, 0), (pi, splineG2 (-1)), (splineG1 (-1), pi) ] + applyIfNeg (x, y) = if y < 0 then (-x, -y) else (x ,y) + points' = [applyIfNeg ((x+y)/2, x - y) | (x, y) <- angles] + points = sortOn fst $ points' <> maybe [] (\(x,y) -> [(x + 2*pi, y)]) (listToMaybe points') + splineAD = genSplineFun points + + fmod a b = a - b * fromIntegral (truncate (a / b)) + + tot = 100 + go 100 = [] + go ix = (p, q) : go (ix+1) + where + ai = fromIntegral ix * 2 * pi / 99 - pi + di = splineAD ai + t1i = tauScale * cos (ai + di) + t2i = tauScale * cos (ai - di) + p = tau2theta1 t1i + q = tau2theta2 t2i + diff --git a/src/Algorithm/SRTree/Likelihoods.hs b/src/Algorithm/SRTree/Likelihoods.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Likelihoods.hs @@ -0,0 +1,329 @@ +{-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE UnboxedTuples #-} + +----------------------------------------------------------------------------- +-- | +-- Module : AlgorithV.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 (..) + , Loss (..) + , readLoss + , Target + , Columns + , buildDistLoss + , buildLoss + , buildPredictor + , fisherNLL + , getSErr + , hessianNLL + ) + where + +import Data.SRTree +import Data.SRTree.Recursion ( cata, accu ) +import Data.SRTree.Derivative (deriveByParam, deriveByVar, derivative, derivOp) +import Data.SRTree.Eval +import qualified Data.IntMap.Strict as IntMap +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM + +import GHC.IO (unsafePerformIO) +import Data.Maybe +import Text.Read (readMaybe) + +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Unboxed.Mutable as VM +import Control.Concurrent (getNumCapabilities) +import Control.Concurrent.Async (forConcurrently) + +import Debug.Trace +import Data.SRTree.Print +import Control.Monad.State.Strict +import Control.Monad.Identity + +import Data.SRTree.Print +import qualified Data.Vector.Generic as G + +-- | Supported distributions for negative log-likelihood. +-- | HGaussian is Gaussian with heteroscedasticity, where the error should be provided. +data Distribution = Gaussian | HGaussian | Bernoulli | Poisson | ROXY | LeastSquares + deriving (Show, Read, Enum, Bounded, Eq) + +-- | Loss functions used to build the per-row optimization objective (see +-- 'buildLoss'), to be used by e.g. "Algorithm.SRTree.Opt". 'NLL' wraps a +-- 'Distribution' to use its negative log-likelihood as the loss -- +-- including the plain \'MSE\' and \'LOG10\' losses, reached via @NLL MSE@ +-- and @NLL LOG10@ respectively (kept on 'Distribution', rather than +-- duplicated here, since Haskell does not allow two data constructors +-- with the same name -- 'MSE' and 'LOG10' -- to coexist in the same +-- module). +data Loss = MSE | LOG10 | MAE | MAPE | Pinball Double | NLL Distribution + deriving (Show, Read, Eq) + +instance Enum Loss where + fromEnum MSE = 0 + fromEnum LOG10 = 1 + fromEnum MAE = 2 + fromEnum MAPE = 3 + fromEnum (Pinball _) = 4 + fromEnum (NLL dist) = 5 + fromEnum dist + + toEnum 0 = MSE + toEnum 1 = LOG10 + toEnum 2 = MAE + toEnum 3 = MAPE + toEnum 4 = Pinball 0.95 + toEnum x | x >= 5 = NLL (toEnum (x-5)) + +instance Bounded Loss where + minBound = MSE + maxBound = NLL ROXY + +-- | Parse a loss from its CLI string. Accepts both the direct 'Loss' +-- names ('MSE', 'LOG10', 'MAE', 'MAPE', @Pinball tau@) and the bare +-- 'Distribution' names ('Gaussian', 'HGaussian', 'Bernoulli', 'Poisson', +-- 'ROXY', 'LeastSquares'), which are wrapped in 'NLL'. +readLoss :: String -> Maybe Loss +readLoss s = case readMaybe s of + Just l -> Just l + Nothing -> NLL <$> (readMaybe s :: Maybe Distribution) + +-- | logistic function +logistic :: Floating a => a -> a +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 :: Target -> Double +negSum = negate . V.sum +{-# inline negSum #-} + +checkAssumptions :: Distribution -> Maybe Target -> Target -> Bool +checkAssumptions Gaussian _ _ = True +checkAssumptions HGaussian (Just yErr) _ = True +checkAssumptions HGaussian Nothing _ = False +checkAssumptions Bernoulli _ ys = V.all (\x -> x /= 0 && x /= 1) ys +checkAssumptions Poisson _ ys = V.all (>0) ys +checkAssumptions LeastSquares _ _ = True +checkAssumptions ROXY mYerr ys = isJust mYerr + +-- WARNING: pass tree with parameters +-- TODO: handle error similar to ROXY + +-- | Builds the per-row negative log-likelihood expression for a given +-- 'Distribution', to be summed across rows (e.g. by +-- 'Algorithm.SRTree.AD.evalGradMulti') and differentiated by automatic +-- differentiation. The special variable index @-1@ refers to the target +-- ('ys') and @-2@ to the target's measurement error ('yErr'), following +-- the convention used by "Algorithm.SRTree.AD". +-- +-- 'buildLoss' delegates to this function for the @'NLL' dist@ loss. +buildDistLoss :: Distribution -> Double -> Fix SRTree -> Fix SRTree +buildDistLoss Gaussian m tree = (square(tree - var (-1)) * (e (negate (param p)))) + (((param p))) + where + square = Fix . Uni Square + e = Fix. Uni Exp + p = countParamsUniq tree +buildDistLoss HGaussian m tree = (tree - var (-1)) ** 2 / var (-2) + constv m * log (2*pi* var (-2)) +buildDistLoss Poisson m tree = var (-1) * log (var (-1)) + exp tree - var (-1) * tree +buildDistLoss Bernoulli m tree = log (1 + exp (negate tree)) + (1 - var (-1)) * tree +buildDistLoss LeastSquares m tree = ((tree - var (-1)) ** 2) / constv m +buildDistLoss ROXY m tree = neglogP + where + p = countParamsUniq tree + f = log (abs tree) / log 10 + fprime = deriveByVar 0 tree / (log 10 * tree) * var 0 * log 10 + logX = var 1 + logY = var 2 + logXErr = var 3 + logYErr = var 4 + sig = param p + 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 + +-- | Builds the per-row loss expression for a given 'Loss', to be summed +-- across rows (e.g. by 'Algorithm.SRTree.AD.evalGradMulti') and +-- differentiated by automatic differentiation. Same special variable +-- convention as 'buildDistLoss'. +buildLoss :: Loss -> Double -> Fix SRTree -> Fix SRTree +buildLoss MSE m tree = ((tree - var (-1)) ** 2) / constv m +buildLoss LOG10 m tree = (((log (y / tree')) / log 10) ** 2) / constv m + where + tree' = (tree + sqrt(tree^2 + 1e-10)) + y = (var (-1) + sqrt(var (-1) ^ 2 + 1e-10)) + +buildLoss MAE m tree = abs (tree - var (-1)) / constv m + +-- | Mean absolute percentage error. A small epsilon is added to the +-- denominator's magnitude to avoid division by zero when the target is +-- (close to) zero. +buildLoss MAPE m tree = (abs (tree - var (-1)) / (abs (var (-1)) + constv 1e-8)) / constv m + +-- | Pinball (quantile) loss for a residual @r = y - yhat@: +-- @tau * r@ if @r >= 0@, @(tau - 1) * r@ otherwise. Both cases are +-- captured in closed form by @0.5 * ((2*tau - 1) * r + abs r)@, which +-- avoids branching in the symbolic tree. +buildLoss (Pinball tau) m tree = ((constv (2*tau - 1) * r + abs r) / 2) / constv m + where r = var (-1) - tree + +buildLoss (NLL dist) m tree = buildDistLoss dist m tree + +-- | Builds the predictor expression from a fitted model tree by applying +-- the inverse link function implied by the 'Distribution': @exp@ for +-- 'Poisson', the logistic function for 'Bernoulli', and the identity +-- otherwise. +buildPredictor :: Distribution -> Fix SRTree -> Fix SRTree +buildPredictor Poisson tree = exp tree +buildPredictor Bernoulli tree = 1 / (1 + exp (negate tree)) +buildPredictor _ tree = tree + +-- | Fisher information of negative log-likelihood +fisherNLL :: Distribution -> Maybe Target -> Columns -> Target -> Fix SRTree -> Target -> Target +fisherNLL ROXY mYerr xss ys tree theta = V.generate p finiteDiff + where + m = V.length ys + p = V.length theta + loss = compileLoss xss (buildDistLoss ROXY (fromIntegral m) tree) ys mYerr + f = loss theta + eps = 1e-6 + finiteDiff ix = unsafePerformIO $ do + theta' <- V.thaw theta + v <- VM.read theta' ix + VM.write theta' ix (v + eps) + thetaPlus <- V.freeze theta' + VM.write theta' ix (v - eps) + thetaMinus <- V.freeze theta' + let fPlus = loss thetaPlus + fMinus = loss thetaMinus + pure $ (fPlus + fMinus - 2*f)/(eps*eps) +fisherNLL Gaussian mYerr xss ys tree theta = V.generate p finiteDiff + where + m = V.length ys + p = V.length theta + loss = compileLoss xss (buildDistLoss Gaussian (fromIntegral m) tree) ys mYerr + f = loss theta + eps = 1e-6 + finiteDiff ix = unsafePerformIO $ do + theta' <- V.thaw theta + v <- VM.read theta' ix + VM.write theta' ix (v + eps) + thetaPlus <- V.freeze theta' + VM.write theta' ix (v - eps) + thetaMinus <- V.freeze theta' + let fPlus = loss thetaPlus + fMinus = loss thetaMinus + pure $ (fPlus + fMinus - 2*f)/(eps*eps) +fisherNLL dist mYerr xss ys tree theta = V.generate p build + where + build ix = let dtdix = deriveByParam ix t' + d2tdix2 = deriveByParam ix dtdix + f' = eval dtdix + f'' = eval d2tdix2 + in V.sum $ phi' * f'^2 - res * f'' + --case dist of + -- Gaussian -> V.sum . (/(theta V.! (p-1))) $ phi' * f'^2 - res * f'' + -- _ -> V.sum $ phi' * f'^2 - res * f'' + m = V.length ys + p = V.length theta + t' = fst $ floatConstsToParam tree + eval = \t -> compile xss t theta + yhat = eval t' + res = ys - phi + yErr = case mYerr of + Nothing -> V.replicate m est + Just e -> e + est = fromIntegral (m - p) + + (phi, phi') = case dist of + Gaussian -> (yhat, V.replicate m 1) + LeastSquares -> (yhat, V.replicate m 1) + Bernoulli -> (logistic yhat, phi*(V.replicate m 1 - phi)) + Poisson -> (exp yhat, phi) + +-- | 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 Target -> Columns -> Target -> Fix SRTree -> Target -> Columns +hessianNLL ROXY mYerr xss ys tree theta = undefined +hessianNLL Gaussian mYerr xss ys tree theta = [V.generate p (build iy) | iy <- [0..p-1]] + where + build iy ix = let dtdix = deriveByParam ix tree + dtdiy = deriveByParam iy tree + d2tdixy = deriveByParam iy dtdix + fx = eval dtdix + fy = eval dtdiy + fxy = eval d2tdixy + in if ix < p-1 && iy < p-1 + then V.sum . (/yErr) $ fx * fy - res * fxy + else if ix == p-1 && iy == p-1 + then (*0.5) . V.sum . (/ yErr ) $ res*res + else if ix == p-1 + then V.sum . (/yErr) $ res * fy + else V.sum . (/yErr) $ res * fx + m = V.length ys + p = V.length theta + yErr :: Target + yErr = V.replicate m $ exp (theta V.! (p-1)) / est + yhat = eval tree + res = ys - yhat + eval = \t -> compile xss t theta + est = fromIntegral (m - p + 1) + +hessianNLL dist mYerr xss ys tree theta = [V.generate p (build iy) | iy <- [0..p-1]] + where + build iy ix = let dtdix = deriveByParam ix t' + dtdiy = deriveByParam iy t' + d2tdixy = deriveByParam iy dtdix + fx = eval dtdix + fy = eval dtdiy + fxy = eval d2tdixy + in case dist of + Gaussian -> V.sum . (/yErr) $ phi' * fx * fy - res * fxy + _ -> V.sum $ phi' * fx * fy - res * fxy + + m = V.length ys + p = V.length theta + t' = tree -- relabelParams tree -- $ floatConstsToParam tree + eval = \t -> compile xss t theta + yErr = case mYerr of + Nothing -> V.replicate m est + Just e -> e + est = fromIntegral (m - p) + yhat = eval t' + res = ys - phi + + (phi, phi') = case dist of + Gaussian -> (yhat, V.replicate m 1) + LeastSquares -> (yhat, V.replicate m 1) + Bernoulli -> (logistic yhat, phi*(V.replicate m 1 - phi)) + Poisson -> (exp yhat, phi) + diff --git a/src/Algorithm/SRTree/ModelSelection.hs b/src/Algorithm/SRTree/ModelSelection.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/ModelSelection.hs @@ -0,0 +1,198 @@ +{-# 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 + ( bic + , aic + , evidence + , fractionalBayesFactor + , mdl + , mdlLatt + , mdlFreq + , logFunctional + , logFunctionalFreq + , ModelEval (..) + , module Algorithm.SRTree.Compile + ) where + +import Algorithm.SRTree.Utils ( det ) +import Algorithm.SRTree.Likelihoods + ( fisherNLL, hessianNLL + , Distribution(..), Loss(..), buildDistLoss + ) +import Data.SRTree +import Data.SRTree.Eval (Target, Columns, compileLoss) +import Data.SRTree.Recursion (cata) +import qualified Data.Vector.Unboxed as U +import Algorithm.SRTree.Compile + +import Debug.Trace + +-- | Bayesian information criterion +bic :: EvaluatedTree -> Double +bic et = valParams et * log (valRows et) + 2 * valLoss et +{-# INLINE bic #-} + +-- | Akaike information criterion +aic :: EvaluatedTree -> Double +aic et = 2 * valParams et + 2 * valLoss et +{-# INLINE aic #-} + +-- | Evidence +evidence :: EvaluatedTree -> Double +evidence et = (1 - b) * valLoss et - valParams et / 2 * log b + where + b = 1 / sqrt (valRows et) +{-# INLINE evidence #-} + +fractionalBayesFactor :: EvaluatedTree -> Double +fractionalBayesFactor et = (1 - b) * valLoss et - valParams et / 2 * log b + f_compl + valParams et / 2 * log(2*pi*nup) + where + b = 1 / sqrt (valRows et) + nup = exp(1 - log 3) + f_compl = countNodes (valTree et) * log (countUniqueTokens (valTree et)) +{-# 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 :: EvaluatedTree -> Double +mdl et = valLoss et + logFunctional (valTree et) + valLogParams et +{-# INLINE mdl #-} + +-- | MDL Lattice as described in +-- Bartlett, Deaglan, Harry Desmond, and Pedro Ferreira. "Priors for symbolic regression." Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 2023. +mdlLatt :: EvaluatedTree -> Double +mdlLatt et = valLoss et + logFunctional (valTree et) + valLogParamsLattice et +{-# INLINE mdlLatt #-} + +-- | same as `mdl` but weighting the functional structure by frequency calculated using a wiki information of +-- physics and engineering functions +mdlFreq :: EvaluatedTree -> Double +mdlFreq et = valLoss et + logFunctionalFreq (valTree et) + valLogParams et +{-# INLINE mdlFreq #-} + +-- | The possible metrics used to evaluate\/select a fitted model, +-- ranging from plain loss functions ('EvalLoss', wrapping any 'Loss' -- +-- including a distribution's negative log-likelihood via @EvalLoss (NLL +-- dist)@) to the error metrics and model-selection criteria already +-- provided by this module ('RMSE', 'R2', 'AIC', 'BIC', 'Evidence', 'FBF', +-- 'MDL', 'MDLLatt', 'MDLFreq'). +data ModelEval + = RMSE + | R2 + | AIC + | BIC + | Evidence + | FBF + | MDL + | MDLLatt + | MDLFreq + | EvalLoss Loss + deriving (Show, Read, Eq) + +instance Enum ModelEval where + fromEnum RMSE = 0 + fromEnum R2 = 1 + fromEnum AIC = 2 + fromEnum BIC = 3 + fromEnum Evidence = 4 + fromEnum FBF = 5 + fromEnum MDL = 6 + fromEnum MDLLatt = 7 + fromEnum MDLFreq = 8 + fromEnum (EvalLoss l) = 9 + fromEnum l + + toEnum 0 = RMSE + toEnum 1 = R2 + toEnum 2 = AIC + toEnum 3 = BIC + toEnum 4 = Evidence + toEnum 5 = FBF + toEnum 6 = MDL + toEnum 7 = MDLLatt + toEnum 8 = MDLFreq + toEnum x | x >= 9 = EvalLoss (toEnum (x-9)) + +instance Bounded ModelEval where + minBound = RMSE + maxBound = EvalLoss maxBound + +-- | Evaluates the requested 'ModelEval' metric. +-- +-- for 'RMSE', and 'R2' the tree must have been compiled +-- with MSE loss. + +evalModelSelection :: ModelEval -> EvaluatedTree -> Double +evalModelSelection (EvalLoss MAE) et = valLoss et +evalModelSelection (EvalLoss MAPE) et = valLoss et +evalModelSelection (EvalLoss (Pinball tau)) et = valLoss et +evalModelSelection (EvalLoss (NLL dist)) et = valLoss et +evalModelSelection RMSE et = sqrt (valLoss et) -- assumes MSE +evalModelSelection R2 et = 1 - (valRows et * valLoss et) / valVar et -- assumes MSE +evalModelSelection AIC et = aic et +evalModelSelection BIC et = bic et +evalModelSelection Evidence et = evidence et +evalModelSelection FBF et = fractionalBayesFactor et +evalModelSelection MDL et = mdl et +evalModelSelection MDLLatt et = mdlLatt et +evalModelSelection MDLFreq et = mdlFreq et +{-# INLINE evalModelSelection #-} + +-- log of the functional complexity +logFunctional :: Fix SRTree -> Double +logFunctional tree = countNodes tree * log (countUniqueTokens tree') + foldr (\c acc -> log (abs c) + acc) 0 consts + log(2) * numberOfConsts + where + tree' = fst $ floatConstsToParam tree + consts = getIntConsts tree + numberOfConsts = fromIntegral $ length consts +{-# 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 #-} + + +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 +{-# 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,98 @@ +{-# LANGUAGE BangPatterns #-} +----------------------------------------------------------------------------- +-- | +-- Module : Algorithm.SRTree.Opt +-- Copyright : (c) Fabricio Olivetti 2021 - 2024 +-- License : BSD3 +-- Maintainer : fabricio.olivetti@gmail.com +-- Stability : experimental +-- Portability : ConstraintKinds +-- +-- Functions to optimize the parameters of an expression. +-- +----------------------------------------------------------------------------- +module Algorithm.SRTree.NonlinearOpt + where + +import Algorithm.SRTree.Likelihoods +import Numeric.Optimization.NLOPT +import Data.Bifunctor (bimap, second) +import Data.SRTree (Fix (..), SRTree (..), floatConstsToParam, relabelParams, countNodes, convertProtectedOps) +import Data.SRTree.Eval +import Algorithm.SRTree.AD + +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed.Mutable as VM +import qualified Data.Vector.Generic as G + +import qualified Data.IntMap.Strict as IntMap +import Data.SRTree.Recursion +import Control.Monad.State.Strict +import Control.Monad.Identity + +import Debug.Trace + +minimizeNLLWith :: (VS.Vector Double -> (Double, VS.Vector Double)) -> (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> Int -> Target -> (Target, Double, Int) +minimizeNLLWith funAndGrad alg niter t0 + | niter == 0 = (t0, f, 0) + | n == 0 = (t0, f, 0) + | otherwise = (t_opt', fst (funAndGrad t_opt), nEvs) + where + t0' = G.convert t0 + n = V.length t0 + + (f, _) = funAndGrad t0' -- if there's no parameter or no iterations + + algorithm = alg funAndGrad (Just $ VectorStorage $ fromIntegral n) + stop = ObjectiveRelativeTolerance 1e-6 :| [ObjectiveAbsoluteTolerance 1e-6, MaximumEvaluations (fromIntegral niter)] + problem = LocalProblem (fromIntegral n) stop algorithm + (t_opt, nEvs) = case minimizeLocal problem t0' of + Right sol -> (solutionParams sol, nEvals sol) + Left e -> (t0', 0) + t_opt' = G.convert t_opt +{-# INLINE minimizeNLLWith #-} + +-- | minimizes the negative log-likelihood of the expression +minimizeNLL' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Target -> (Target, Double, Int) +minimizeNLL' alg backend dist mYerr niter xss ys tree t0 = minimizeNLLWith funAndGrad alg niter t0 + where + m = V.length ys + tree' = buildLoss dist (fromIntegral m) tree + funAndGrad = compileFunAndGrad backend xss ys mYerr tree' + + +minimizeNLL :: ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Target -> (Target, Double, Int) +minimizeNLL = minimizeNLL' TNEWTON + +minimizeNLLWithFixedParam' :: (ObjectiveD -> (Maybe VectorStorage) -> LocalAlgorithm) -> ADBackEnd -> Loss -> Maybe Target -> Int -> Columns -> Target -> Fix SRTree -> Int -> Target -> Target +minimizeNLLWithFixedParam' alg backend dist mYerr' niter xss' ys' tree ix t0 = result + where + m = V.length ys' + tree' = buildLoss dist (fromIntegral m) tree + fixedVal = t0 V.! ix + p = V.length t0 + + evalFull = compileFunAndGrad backend xss' ys' mYerr' tree' + + wrapRed thRed = let (lo, hi) = VS.splitAt ix thRed + in (lo `VS.snoc` fixedVal) VS.++ hi + unwrapRed th = let (lo, hi) = VS.splitAt ix th + in lo VS.++ VS.tail hi + + wrap thRed = let (lo, hi) = V.splitAt ix thRed in (lo `V.snoc` fixedVal) V.++ hi + unwrap th = let (lo, hi) = V.splitAt ix th in lo V.++ V.tail hi + + fgRed :: VS.Vector Double -> (Double, VS.Vector Double) + fgRed thRed = + let thFull = wrapRed thRed + (nll, gradFull) = evalFull thFull + gradRed = unwrapRed gradFull + in (nll, gradRed) + + t0Red = unwrap t0 + (tRawRed,_,_) = minimizeNLLWith fgRed alg niter t0Red + result = wrap tRawRed + +minimizeNLLWithFixedParam = minimizeNLLWithFixedParam' TNEWTON + diff --git a/src/Algorithm/SRTree/Utils.hs b/src/Algorithm/SRTree/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Algorithm/SRTree/Utils.hs @@ -0,0 +1,320 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +module Algorithm.SRTree.Utils where + +import qualified Data.Vector.Unboxed as U +import qualified Data.Vector.Unboxed.Mutable as UM +import Control.Monad +import Control.Monad.Catch +import Control.Monad.Primitive +import Control.Monad.IO.Class +import System.IO.Unsafe + +-- taken from https://hackage.haskell.org/package/cubicspline-0.1.2 +import Control.Arrow +import Data.List (unfoldr) + +import Data.SRTree.Eval +import Debug.Trace (traceShow) + +-- | Internal helper to get dimensions (rows, columns) +matSize :: Columns -> (Int, Int) +matSize [] = (0, 0) +matSize cs@(c:_) = (U.length c, length cs) + +getRows :: Columns -> [Target] +getRows mtx + | n == 0 = [] + | otherwise = [ U.fromListN n [ c U.! i | c <- mtx ] | i <- [0 .. m - 1] ] + where (m, n) = matSize mtx +{-# INLINE getRows #-} + +getCols :: Columns -> [Target] +getCols = id +{-# INLINE getCols #-} + +appendRow :: MonadThrow m => Columns -> Target -> m Columns +appendRow xs v = pure $ zipWith U.snoc xs (U.toList v) +{-# INLINE appendRow #-} + +appendCol :: MonadThrow m => Columns -> Target -> m Columns +appendCol xs v = pure $ xs ++ [v] +{-# INLINE appendCol #-} + +updateS :: Target -> [(Int, Double)] -> Target +updateS vec new = vec U.// new + +linSpace :: Int -> (Double, Double) -> [Double] +linSpace num (lo, hi) = Prelude.take num $ iterate (\x -> x + step) lo + where step = (hi - lo) / (fromIntegral num - 1) +{-# INLINE linSpace #-} + +outer :: (MonadThrow m) => Target -> Target -> m Columns +outer arr1 arr2 + | U.null arr1 || U.null arr2 = pure [] + | otherwise = pure [ U.map (* (arr2 U.! j)) arr1 | j <- [0 .. U.length arr2 - 1] ] +{-# INLINE outer #-} + +-- | Flatten list of column vectors to a row-major U.Vector Double +toRowMajor :: Columns -> U.Vector Double +toRowMajor cols = U.generate (m * n) (\ix -> let (i, j) = ix `divMod` n in (cols !! j) U.! i) + where (m, n) = matSize cols + +-- | Restore a row-major continuous U.Vector Double back to Columns +fromRowMajor :: Int -> Int -> U.Vector Double -> Columns +fromRowMajor m n vec = [ U.generate m (\i -> vec U.! (i * n + j)) | j <- [0 .. n - 1] ] + +unsafeRead :: PrimMonad m => Int -> UM.MVector (PrimState m) Double -> (Int, Int) -> m Double +unsafeRead stride arr (i, j) = UM.unsafeRead arr (i * stride + j) +{-# INLINE unsafeRead #-} + +unsafeWrite :: PrimMonad m => Int -> UM.MVector (PrimState m) Double -> (Int, Int) -> Double -> m () +unsafeWrite stride arr (i, j) val = UM.unsafeWrite arr (i * stride + j) val +{-# INLINE unsafeWrite #-} + +det :: Columns -> Double +det mtx + | m == 0 || n == 0 = 1 + | otherwise = (^2) $ product [ (toRowMajor l) U.! (i * n + i) | i <- [0 .. m - 1] ] + where + (m, n) = matSize mtx + (l, _) = unsafePerformIO (lu mtx) + +detChol :: Columns -> Double +detChol mtx + | m == 0 || n == 0 = 1 + | otherwise = (^2) $ product [ (toRowMajor cho) U.! (i * m + i) | i <- [0 .. m - 1] ] + where + (m, n) = matSize mtx + cho = unsafePerformIO (cholesky mtx) +{-# INLINE det #-} + +rangedLinearDotProd :: PrimMonad m => Int -> Int -> Int -> UM.MVector (PrimState m) Double -> m Double +rangedLinearDotProd r1 r2 len arr = go 0 0 + where + go !acc k + | k < len = do + x <- UM.unsafeRead arr (r1 + k) + y <- UM.unsafeRead arr (r2 + k) + go (acc + x * y) (k + 1) + | otherwise = pure acc +{-# INLINE rangedLinearDotProd #-} + +data NegDef = NegDef deriving Show +instance Exception NegDef + +cholesky :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m Columns +cholesky arr + | m /= n = error $ "cholesky dimension mismatch " <> show m <> " X " <> show n + | m == 0 = pure [] + | otherwise = do + l <- UM.new (m * m) + let orig = toRowMajor arr + forM_ [0 .. m - 1] $ \i -> + forM_ [0 .. m - 1] $ \j -> + if i < j then unsafeWrite m l (i, j) 0 + else do + let cur = orig U.! (i * m + j) + rowI = i * m + rowJ = j * m + xjj <- UM.unsafeRead l (rowJ + j) + tot <- rangedLinearDotProd rowI rowJ j l + let delta = cur - tot + if i == j + then if delta <= 0 + then throwM NegDef + else UM.unsafeWrite l (rowI + j) (sqrt delta) + else UM.unsafeWrite l (rowI + j) (delta / xjj) + frozen <- U.unsafeFreeze l + pure $ fromRowMajor m m frozen + where (m, n) = matSize arr +{-# INLINE cholesky #-} + +invChol :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m Columns +invChol arr = do + lMtx <- cholesky arr + let (m, _) = matSize arr + mtx <- U.thaw (toRowMajor lMtx) + forM_ [0 .. m - 1] $ \i -> do + lII <- unsafeRead m mtx (i, i) + unsafeWrite m mtx (i, i) (1 / lII) + forM_ [0 .. i - 1] $ \j -> do + tot <- rangedLinearDotProd (i * m + j) (j * m + j) (i - j) mtx + unsafeWrite m mtx (j, i) ((-tot) / lII) + unsafeWrite m mtx (i, j) 0 + + mm <- UM.replicate (m * m) 0 + forM_ [0 .. m - 1] $ \i -> do + dii <- rangedLinearDotProd (i * m + i) (i * m + i) (m - i) mtx + unsafeWrite m mm (i, i) dii + forM_ [i + 1 .. m - 1] $ \j -> do + dij <- rangedLinearDotProd (i * m + j) (j * m + j) (m - j) mtx + unsafeWrite m mm (i, j) dij + unsafeWrite m mm (j, i) dij + frozen <- U.unsafeFreeze mm + pure $ fromRowMajor m m frozen +{-# INLINE invChol #-} + +lu :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> m (Columns, Columns) +lu mtx = do + let (m, n) = matSize mtx + orig = toRowMajor mtx + u <- UM.replicate (m * n) 0 + forM_ [0 .. min m n - 1] $ \i -> unsafeWrite n u (i, i) 1 + l <- UM.replicate (m * n) 0 + + let buildLVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do + lik <- unsafeRead n l (i, k) + ukj <- unsafeRead n u (k, j) + go (k+1) (s + lik * ukj) + s' <- go 0 0 + unsafeWrite n l (i, j) ((orig U.! (i * n + j)) - s') + + buildL !i !j = when (i /= m) $ do + buildLVal i j + buildL (i+1) j + + buildUVal !i !j = do + let go !k !s + | k == j = pure s + | otherwise = do + ljk <- unsafeRead n l (j, k) + uki <- unsafeRead n u (k, i) + go (k+1) (s + ljk * uki) + s' <- go 0 0 + ljj <- unsafeRead n l (j, j) + unsafeWrite n u (j, i) (((orig U.! (j * n + i)) - s') / ljj) + + buildU !i !j = when (i /= n) $ do + buildUVal i j + buildU (i+1) j + + buildLU !j = when (j /= n && j /= m) $ do + buildL j j + buildU j j + buildLU (j+1) + + buildLU 0 + finalL <- U.unsafeFreeze l + finalU <- U.unsafeFreeze u + pure (fromRowMajor m n finalL, fromRowMajor m n finalU) + +forwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +forwardSub a b = do + let m = U.length b + n = length a + aMat = toRowMajor a + x <- UM.replicate m 0 + let coeff !i !j !s + | j == i = pure s + | otherwise = do + let aij = aMat U.! (i * n + j) + xj <- UM.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i /= m) $ do + let bi = b U.! i + aii = aMat U.! (i * n + i) + c <- coeff i 0 0 + UM.unsafeWrite x i ((bi - c) / aii) + go (i+1) + go 0 + U.unsafeFreeze x + +backwardSub :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +backwardSub a b = do + let m = U.length b + n = length a + aMat = toRowMajor a + x <- UM.replicate m 0 + let coeff !i !j !s + | j == m = pure s + | otherwise = do + let aij = aMat U.! (i * n + j) + xj <- UM.unsafeRead x j + coeff i (j+1) (s + aij * xj) + go !i = when (i >= 0) $ do + let bi = b U.! i + aii = aMat U.! (i * n + i) + c <- coeff i (i+1) 0 + UM.unsafeWrite x i ((bi - c) / aii) + go (i-1) + go (m-1) + U.unsafeFreeze x + +luSolve :: (PrimMonad m, MonadThrow m, MonadIO m) => Columns -> Target -> m Target +luSolve a b = do + (l, u) <- lu a + forwardSub l b >>= backwardSub u + +type PolyCos = (Double, Double, Double) + +cubicSplineCoefficients :: [(Double, Double)] -> [PolyCos] +cubicSplineCoefficients xs = Prelude.zip3 x y z' + where + x = map fst xs + y = map snd xs + xdiff = zipWith (-) (tail x) x + xdiff' = U.fromList xdiff + + dydx :: U.Vector Double + dydx = U.fromList $ Prelude.zipWith3 (\y0 y1 xd -> (y0 - y1) / xd) (tail y) y xdiff + + n = length x + + w :: [Double] + w = 0 : nextW 1 w + where + nextW ix (wi : t) + | ix == n - 1 = [] + | otherwise = + let m = (xdiff' U.! (ix - 1)) * (2 - wi) + 2 * (xdiff' U.! ix) + wn = (xdiff' U.! ix) / m + in wn : nextW (ix + 1) t + + z :: [Double] + z = 0 : nextZ 1 z + where + nextZ ix (zi : t) + | ix == n - 1 = [0] + | otherwise = + let m = (xdiff' U.! (ix - 1)) * (2 - (w !! (ix - 1))) + 2 * (xdiff' U.! ix) + zn = (6 * ((dydx U.! ix) - (dydx U.! (ix - 1))) - (xdiff' U.! (ix - 1)) * zi) / m + in zn : nextZ (ix + 1) t + + z' :: [Double] + z' = Prelude.reverse $ 0 : [z !! i - w !! i * z !! (i + 1) | i <- [n - 2, n - 3 .. 0]] + +chunkBy :: Int -> [t] -> [[t]] +chunkBy n = unfoldr go + where + go [] = Nothing + go x = Just $ splitAt n x + +genSplineFun :: [(Double, Double)] -> Double -> Double +genSplineFun pts x + | length xs < 2 = x + | x < head xs = y1 + (x - x1) * (y2 - y1) / (x2 - x1) + | x > last xs = y_1 + (x - x_1) * (y_n - y_1) / (x_n - x_1) + | otherwise = go xs $ zip coefs (tail coefs) + where + xs = map fst pts + ys = map snd pts + coefs = cubicSplineCoefficients pts + x1 = head xs; y1 = head ys + x2 = xs !! 1; y2 = ys !! 1 + x_1 = xs !! (len - 2); y_1 = ys !! (len - 2) + x_n = last xs; y_n = last ys + len = length xs + + evalAt (a1, b1, c1) (a2, b2, c2) y = + let hi1 = a2 - a1 + in c1 / (6 * hi1) * (a2 - y)^3 + c2 / (6 * hi1) * (y - a1)^3 + + (b2 / hi1 - c2 * hi1 / 6) * (y - a1) + (b1 / hi1 - c1 * hi1 / 6) * (a2 - y) + + go [x1, x2] [(c1, c2)] = evalAt c1 c2 x + go (x1 : x2 : xs') ((c1, c2) : cs) + | x >= x1 && x <= x2 = evalAt c1 c2 x + | otherwise = go (x2 : xs') cs diff --git a/src/Data/SRTree.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,416 @@ +{-# language ImportQualifiedPost #-} +{-# language ViewPatterns #-} +{-# language OverloadedStrings #-} +{-# language BlockArguments #-} +{-# language ExplicitForAll #-} +{-# language BangPatterns #-} +{-# language LambdaCase #-} +{-# language RankNTypes, ScopedTypeVariables #-} +----------------------------------------------------------------------------- +-- | +-- 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(..), splitFileNameParams, getRows, getColumns ) + where + +import Codec.Compression.GZip (decompress) +import Data.ByteString.Char8 qualified as B +import Data.ByteString.Lazy qualified as BS +import Data.List (delete, find, intercalate) +import Data.Maybe (fromJust) +import Data.Ratio ((%)) +import Data.Vector.Unboxed (Vector) +import qualified Data.Vector as VB +import qualified Data.Vector.Unboxed as V +import System.FilePath (takeExtension) +import Text.Read (readMaybe) +import Control.Monad.State.Strict +import System.Random +import qualified Data.Vector.Primitive as VP +import Data.Foldable qualified as Foldable +import Data.Primitive.Array qualified as Array +import Control.Monad.ST (runST) +import Control.Monad.ST.Strict (ST) + +-- a dataset is a triple (X, y, y_error) +type DataSet = ([Vector Double], Vector Double, Maybe (Vector Double)) + +-- | Loads a list of list of bytestrings to a matrix of double +loadMtx :: [[B.ByteString]] -> [Vector Double] +loadMtx [] = [] +loadMtx rows = map V.fromList + $ foldr (zipWith (:) . map parseDouble) (replicate ncols []) rows + where ncols = length (head rows) +{-# INLINE loadMtx #-} + +-- | Powers of ten as exact 'Integer's, precomputed once and shared by every +-- 'parseDouble' call. The per-value @10 ^ k@ exponentiation previously ran a +-- growing-Integer multiply loop on every parsed number, which showed up as a +-- measurable chunk of the corpus-load allocation. The table is the exact same +-- integer, so conversions stay bit-identical. +maxPow10 :: Int +maxPow10 = 400 + +pow10 :: VB.Vector Integer +pow10 = VB.generate (maxPow10 + 1) (\k -> 10 ^ k) +{-# NOINLINE pow10 #-} + +-- | @10^k@ as an exact 'Integer'; falls back to direct exponentiation for +-- exponents beyond the precomputed range (only reachable with absurd inputs). +pow10E :: Int -> Integer +pow10E k | k >= 0 && k <= maxPow10 = VB.unsafeIndex pow10 k + | otherwise = 10 ^ k +{-# INLINE pow10E #-} + +-- | Fast decimal double parser over a 'B.ByteString'. Handles an optional +-- sign, a fractional part and an optional 'e'/'E' exponent. The mantissa is +-- accumulated exactly as an 'Integer' and converted to 'Double' through a +-- single 'fromRational', which matches the correctly-rounded result of 'read'. +-- Falls back to 'read' (the slow Show-derived parser) for anything it can't +-- parse (NaN, Infinity, hex floats, etc.), so behavior is unchanged for odd +-- input. +parseDouble :: B.ByteString -> Double +parseDouble bs = case go 0 1 0 False 0 of + Just (m, s, nd, e) + -- when e >= nd the rational m * 10^e / 10^nd is an exact integer, so a + -- single fromInteger is bit-identical to fromRational (which would only + -- gcd-reduce it) but skips the rational machinery entirely. + | e >= nd -> fromInteger (s * (m * pow10E (e - nd))) + -- otherwise the value is m / 10^(nd-e); keep fromRational so the single + -- rounding matches `read` exactly (a Double division by a rounded power + -- of ten would be off by up to an ulp). + | otherwise -> fromRational (s * m % (pow10E (nd - e))) + Nothing -> read (B.unpack bs) + where + n = B.length bs + -- i: index, sgn: +/-1, acc: accumulated mantissa digits (exact Integer), + -- dot: whether a '.' has been seen, nd: number of digits following the + -- decimal point, expo: signed integer exponent from the 'e' tail + go :: Int -> Integer -> Integer -> Bool -> Int -> Maybe (Integer, Integer, Int, Int) + go !i !sgn !acc !dot !nd + | i >= n = Just (acc, sgn, nd, 0) + | otherwise = + let c = fromEnum (B.index bs i) + in case c of + 45 -> if i == 0 then go (i+1) (-sgn) acc dot nd else Nothing -- '-' + 43 -> if i == 0 then go (i+1) sgn acc dot nd else Nothing -- '+' + 46 -> if dot then Nothing else go (i+1) sgn acc True nd -- '.' + _ | c >= 48 && c <= 57 -> + let d = fromIntegral (c - 48) :: Integer + nd' = if dot then nd + 1 else nd + in go (i+1) sgn (acc * 10 + d) dot nd' + | (c == 101 || c == 69) && i > 0 -> -- 'e' / 'E' + parseExp (i+1) sgn acc dot nd + | otherwise -> Nothing + -- parse the (optional) exponent tail: an optional sign then digits + parseExp :: Int -> Integer -> Integer -> Bool -> Int -> Maybe (Integer, Integer, Int, Int) + parseExp !i !sgn !acc !dot !nd + | i >= n = Just (acc, sgn, nd, 0) + | otherwise = + let c = fromEnum (B.index bs i) + in case c of + 45 -> expDig (i+1) sgn acc dot nd (-1) 0 -- '-' + 43 -> expDig (i+1) sgn acc dot nd 1 0 -- '+' + _ -> expDig i sgn acc dot nd 1 0 + where + -- es: exponent sign (+/-1); e: accumulated exponent magnitude + expDig :: Int -> Integer -> Integer -> Bool -> Int -> Int -> Int -> Maybe (Integer, Integer, Int, Int) + expDig !i !sgn !acc !dot !nd !es !e + | i >= n = Just (acc, sgn, nd, es * e) + | otherwise = + let c = fromEnum (B.index bs i) + in if c >= 48 && c <= 57 + then expDig (i+1) sgn acc dot nd es (e * 10 + fromIntegral (c - 48)) + else Nothing + +-- | Returns true if the extension is .gz +isGZip :: FilePath -> Bool +isGZip = (== ".gz") . takeExtension +{-# 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 . toStrict . 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 + -- lazy -> strict without going through a [Word8]/[Char] list (the old + -- B.pack . map toEnum . BS.unpack round trip allocated ~1GB on a 14MB + -- CSV); BS.toStrict is a single O(n) copy. + toStrict = BS.toStrict +{-# INLINE readFileToLines #-} + +-- | Splits the parameters from the filename +-- 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.null filename = ("", replicate 6 B.empty) + | otherwise = (B.unpack fname, take 6 params) + where + (fname : params') = B.split ':' filename + -- fill up the empty parameters with an empty string + 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) + | otherwise = (st_ix, end_ix + 1) + 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 (([Vector Double], Vector Double, [Vector Double], Vector Double), (Maybe (Vector Double), Maybe (Vector Double)), String, String) +loadDataset filename hasHeader = do + csv <- readFileToLines fname + pure $ processData csv params hasHeader + where + (fname, params) = splitFileNameParams filename + +-- support function that does everything for loadDataset +processData :: [[B.ByteString]] -> [B.ByteString] -> Bool -> (([Vector Double], Vector Double, [Vector Double], Vector Double), (Maybe (Vector Double), Maybe (Vector Double)), String, String) +processData csv params hasHeader = ((x_train, y_train, x_val, y_val) , (y_err_train, y_err_val), varnames, targetname) + where + ncols = length $ head csv + 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 = map (datum !!) ixs + y = datum !! iy + y_err = datum !! iy_err + + x_train = map (V.take end . V.drop st) x + y_train = V.take end . V.drop st $ y + x_val = map (V.drop (st + end)) x + y_val = V.drop (st + end) y + + y_err_train = if iy_err == -1 then Nothing else Just $ (V.take end . V.drop st) y_err + y_err_val = if iy_err == -1 then Nothing else Just $ (V.take end . V.drop st) y_err +{-# inline processData #-} + +chunksOf :: Int -> [e] -> [[e]] +chunksOf i ls = Prelude.map (Prelude.take i) (build (splitter ls)) + where + splitter :: [e] -> ([e] -> a -> a) -> a -> a + splitter [] _ n = n + splitter l c n = l `c` splitter (Prelude.drop i l) c n + build :: ((a -> [a] -> [a]) -> [a] -> [a]) -> [a] + build g = g (:) [] + +splitData :: DataSet -> Int -> State StdGen (DataSet, DataSet) +splitData (x, y, mYErr) k = do + if k == 1 + then pure ((x, y, mYErr), (x, y, mYErr)) + else do + ixs' <- (state . shuffle) [0 .. sz-1] + let ixs = chunksOf k ixs' + + let tr_ix = [ix | ixs_i <- ixs, ix <- Prelude.tail ixs_i] + val_ix = [ix | ixs_i <- ixs, let ix = Prelude.head ixs_i] + (x_tr, x_te) = getX tr_ix val_ix x + (y_tr, y_te) = getY tr_ix val_ix y + + mY = fmap (getY tr_ix val_ix) mYErr + (y_err_tr, y_err_te) = (fmap fst mY, fmap snd mY) + pure ((x_tr, y_tr, y_err_tr), (x_te, y_te, y_err_te)) + where + sz = V.length y + + getX :: [Int] -> [Int] -> [Vector Double] -> ([Vector Double], [Vector Double]) + getX tr_ix val_ix xs = ( [ V.fromList [x V.! ix | ix <- tr_ix] | x <- xs ] + , [ V.fromList [x V.! ix | ix <- val_ix] | x <- xs ] + ) + getY :: [Int] -> [Int] -> Vector Double -> (Vector Double, Vector Double) + getY tr_ix val_ix ys = ( V.fromList [ys V.! ix | ix <- tr_ix] + , V.fromList [ys V.! ix | ix <- val_ix] + ) + +getTrain :: ((a, b1, c1, d1), (c2, b2), c3, d2) -> (a, b1, c2) +getTrain ((a, b, _, _), (c, _), _, _) = (a,b,c) + +getX :: DataSet -> [Vector Double] +getX (a, _, _) = a + +getTarget :: DataSet -> Vector Double +getTarget (_, b, _) = b + +getError :: DataSet -> Maybe (Vector Double) +getError (_, _, c) = c + +loadTrainingOnly fname b = getTrain <$> loadDataset fname b + +-- | Shuffles a list, taken from list-shuffle +shuffle :: (RandomGen g) => [a] -> g -> ([a], g) +shuffle list gen0 = + runST do + array <- listToMutableArray list + gen1 <- shuffleN (Array.sizeofMutableArray array - 1) array gen0 + array1 <- Array.unsafeFreezeArray array + pure (Foldable.toList array1, gen1) + +listToMutableArray :: forall a s. [a] -> ST s (Array.MutableArray s a) +listToMutableArray list = do + array <- Array.newArray (length list) undefined + let writeElems :: Int -> [a] -> ST s () + writeElems !i = \case + [] -> pure () + x : xs -> do + Array.writeArray array i x + writeElems (i + 1) xs + writeElems 0 list + pure array +{-# INLINE listToMutableArray #-} + +shuffleN :: forall a g s. (RandomGen g) => Int -> Array.MutableArray s a -> g -> ST s g +shuffleN n0 array = + go 0 + where + go :: Int -> g -> ST s g + go !i gen0 + | i >= n = pure gen0 + | otherwise = do + let (j, gen1) = uniformR (i, m) gen0 + swapArrayElems i j array + go (i + 1) gen1 + + n = min n0 m + m = Array.sizeofMutableArray array - 1 +{-# SPECIALIZE shuffleN :: Int -> Array.MutableArray s a -> StdGen -> ST s StdGen #-} + +-- Swap two elements in a mutable array. +swapArrayElems :: Int -> Int -> Array.MutableArray s a -> ST s () +swapArrayElems i j array = do + x <- Array.readArray array i + y <- Array.readArray array j + Array.writeArray array i y + Array.writeArray array j x +{-# INLINE swapArrayElems #-} diff --git a/src/Data/SRTree/Derivative.hs b/src/Data/SRTree/Derivative.hs new file mode 100644 --- /dev/null +++ b/src/Data/SRTree/Derivative.hs @@ -0,0 +1,140 @@ +{-# 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 + , derivOp + ) + 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 #-} + +-- | Returns (d(Output)/d(Left), d(Output)/d(Right)) +-- used for AD +derivOp :: Op -> Double -> Double -> (Double, Double) +derivOp Add _ _ = (1.0, 1.0) +derivOp Sub _ _ = (1.0, -1.0) +derivOp Mul v1 v2 = (v2, v1) +derivOp Div v1 v2 = (1.0 / v2, -(v1) / (v2 * v2)) +-- e.g., Power: d(x^y)/dx = y*x^(y-1), d(x^y)/dy = x^y * ln(x) +derivOp Power v1 v2 = (v2 * (v1 ** (v2 - 1)), (v1 ** v2) * log v1) +derivOp _ _ _ = (0.0, 0.0) -- Add remaining ops +{-# INLINE derivOp #-} + +-- | Symbolic derivative by a variable +deriveByVar :: Int -> Fix SRTree -> Fix SRTree +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,393 @@ +{-# LANGUAGE LambdaCase, BangPatterns #-} + +----------------------------------------------------------------------------- +-- | +-- 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 + ( evalOp + , evalFun + , cbrt + , inverseFunc + , invertibles + , evalInverse + , invright + , invleft + , replicateAs + , Target, Theta, Columns + , compile + , compileLoss + ) + where + +import Data.SRTree.Internal +import Data.SRTree.Recursion (Fix (..), cata) +import Data.Vector.Unboxed (Vector) +import qualified Data.Vector.Unboxed as V +import Control.Monad.ST (runST) +import qualified Data.Vector as VB -- Boxed vector for instructions +import qualified Data.Vector.Unboxed.Mutable as VM +import Control.Concurrent.Async (forConcurrently_) +import System.IO.Unsafe (unsafePerformIO) +import Control.Concurrent (getNumCapabilities) +import Data.Maybe (fromJust) + +-- | Vector of target values +type Target = Vector Double +-- | Vector of parameter values. Needs to be strict to be readily accesible. +type Theta = Vector Double +-- | Matrix of features values +type Columns = [Vector Double] + +-- A multi-threaded replacement for V.sum +sumParallel :: Int -> (Int -> Double) -> Double +sumParallel n f = unsafePerformIO $ do + numThreads <- getNumCapabilities + let chunkSize = n `quot` numThreads + + -- 1. Allocate a single block of unboxed memory EXACTLY ONCE + out <- VM.unsafeNew numThreads + + -- 2. Spawn threads. Each thread gets a unique ID and a slice of memory. + forConcurrently_ [0 .. numThreads - 1] $ \tId -> do + let !start = tId * chunkSize + -- The last thread cleans up the remainder + !end = if tId == numThreads - 1 then n else start + chunkSize + + -- 3. The inner thread loop. Strict, unboxed, and bounds-check free. + let loop !i !acc + | i >= end = return acc + | otherwise = loop (i + 1) (acc + f i) + + total <- loop start 0.0 + VM.unsafeWrite out tId total + + -- 4. Instantly cast the mutable memory to an immutable Vector (O(1) cost) + totals <- V.unsafeFreeze out + return (V.sum totals) +{-# NOINLINE sumParallel #-} + +-- Improve quality of life with Num and Floating instances for our matrices +instance Num Target where + (+) = V.zipWith (+) + (-) = V.zipWith (-) + (*) = V.zipWith (*) + abs = V.map abs + signum = V.map signum + fromInteger = V.singleton . fromInteger + negate = V.map negate + +instance Floating Target where + pi = V.singleton pi + exp = V.map exp + log = V.map log + sqrt = V.map sqrt + sin = V.map sin + cos = V.map cos + tan = V.map tan + asin = V.map asin + acos = V.map acos + atan = V.map atan + sinh = V.map sinh + cosh = V.map cosh + tanh = V.map tanh + asinh = V.map asinh + acosh = V.map acosh + atanh = V.map atanh + (**) = V.zipWith (**) +instance Fractional Target where + fromRational = V.singleton . fromRational + (/) = V.zipWith (/) + recip = V.map recip + +-- We change the Dynamic type to evaluate a single scalar at a specific row index (Int) +data Staged = + Scl Double + | Static (Vector Double) + | Dynamic (Vector Double -> Int -> Double) -- (Theta -> RowIndex -> Result) + +-- A multi-threaded replacement for V.generate +generateParallel :: Int -> (Int -> Double) -> V.Vector Double +generateParallel n f = unsafePerformIO $ do + numThreads <- getNumCapabilities + let chunkSize = n `quot` numThreads + + -- 1. Allocate a single block of unboxed memory EXACTLY ONCE + out <- VM.unsafeNew n + + -- 2. Spawn threads. Each thread gets a unique ID and a slice of memory. + forConcurrently_ [0 .. numThreads - 1] $ \tId -> do + let !start = tId * chunkSize + -- The last thread cleans up the remainder + !end = if tId == numThreads - 1 then n else start + chunkSize + + -- 3. The inner thread loop. Strict, unboxed, and bounds-check free. + let loop !i + | i >= end = return () + | otherwise = do + -- Write directly to the shared memory pointer + VM.unsafeWrite out i (f i) + loop (i + 1) + + loop start + + -- 4. Instantly cast the mutable memory to an immutable Vector (O(1) cost) + V.unsafeFreeze out +{-# NOINLINE generateParallel #-} + +compileLoss :: [Vector Double] -> Fix SRTree -> Target -> Maybe Target -> (Vector Double -> Double) +compileLoss dataset tree y mYerr = + case cata alg tree of + Scl c -> \_ -> V.sum $ V.replicate n c + Static v -> \_ -> V.sum v + -- We only allocate memory EXACTLY ONCE here at the top level + --Dynamic f -> \th -> V.generate n (f th) + Dynamic f -> \th -> V.sum (V.generate n (f th)) + where + n = V.length (head dataset) + yErr = fromJust mYerr + + alg :: SRTree Staged -> Staged + + -- 1. Base Cases + alg (Const c) = Scl c + alg (Var (-1)) = Static y + alg (Var (-2)) = Static yErr + alg (Var i) = Static (dataset !! i) + alg (Param i) = Dynamic (\th !idx -> th `V.unsafeIndex` i) + + -- 2. Univariate Functions + alg (Uni f (Scl c)) = Scl (evalFun f c) + alg (Uni f (Static v)) = Static (V.map (evalFun f) v) + + -- We map the function over the scalar result of the inner closure + alg (Uni f (Dynamic g)) = let !rawFun = evalFun f in Dynamic (\th !i -> rawFun (g th i)) + + -- 3. Binary Functions + alg (Bin op (Scl c1) (Scl c2)) = Scl (evalOp op c1 c2) + alg (Bin op (Scl c) (Static v)) = Static (V.map (evalOp op c) v) + alg (Bin op (Static v) (Scl c)) = Static (V.map (\c2 -> evalOp op c2 c) v) + alg (Bin op (Static v1) (Static v2)) = Static (V.zipWith (evalOp op) v1 v2) + + -- 4. Dynamic Combinations (The Core Optimization) + + alg (Bin op (Scl c) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp c (g th i)) + + alg (Bin op (Dynamic g) (Scl c)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) c) + + -- When combining a Static array with a Dynamic closure, + -- we use unsafeIndex to fetch the static value at row 'i' directly. + alg (Bin op (Static v) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (v `V.unsafeIndex` i) (g th i)) + + alg (Bin op (Dynamic g) (Static v)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) (v `V.unsafeIndex` i)) + + alg (Bin op (Dynamic g1) (Dynamic g2)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g1 th i) (g2 th i)) + + +compile :: [Vector Double] -> Fix SRTree -> (Vector Double -> Vector Double) +compile dataset tree = + case cata alg tree of + Scl c -> \_ -> V.replicate n c + Static v -> \_ -> v + -- We only allocate memory EXACTLY ONCE here at the top level + --Dynamic f -> \th -> V.generate n (f th) + Dynamic f -> \th -> V.generate n (f th) + where + n = V.length (head dataset) + + alg :: SRTree Staged -> Staged + + -- 1. Base Cases + alg (Const c) = Scl c + alg (Var i) = Static (dataset !! i) + -- Look at this! No more V.replicate. It just fetches the scalar directly. + alg (Param i) = Dynamic (\th !idx -> th `V.unsafeIndex` i) + alg (Y i) = undefined -- this shouldn't be called + + -- 2. Univariate Functions + alg (Uni f (Scl c)) = Scl (evalFun f c) + alg (Uni f (Static v)) = Static (V.map (evalFun f) v) + + -- We map the function over the scalar result of the inner closure + alg (Uni f (Dynamic g)) = let !rawFun = evalFun f in Dynamic (\th !i -> rawFun (g th i)) + + -- 3. Binary Functions + alg (Bin op (Scl c1) (Scl c2)) = Scl (evalOp op c1 c2) + alg (Bin op (Scl c) (Static v)) = Static (V.map (evalOp op c) v) + alg (Bin op (Static v) (Scl c)) = Static (V.map (\c2 -> evalOp op c2 c) v) + alg (Bin op (Static v1) (Static v2)) = Static (V.zipWith (evalOp op) v1 v2) + + -- 4. Dynamic Combinations (The Core Optimization) + + alg (Bin op (Scl c) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp c (g th i)) + + alg (Bin op (Dynamic g) (Scl c)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) c) + + -- When combining a Static array with a Dynamic closure, + -- we use unsafeIndex to fetch the static value at row 'i' directly. + alg (Bin op (Static v) (Dynamic g)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (v `V.unsafeIndex` i) (g th i)) + + alg (Bin op (Dynamic g) (Static v)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g th i) (v `V.unsafeIndex` i)) + + alg (Bin op (Dynamic g1) (Dynamic g2)) = + let !rawOp = evalOp op in Dynamic (\th !i -> rawOp (g1 th i) (g2 th i)) + + +-- returns a vector with the same number of rows as xss and containing a single repeated value. +replicateAs :: Columns -> Double -> Target +replicateAs xss c = let m = V.length (head xss) in V.replicate m c +{-# INLINE replicateAs #-} + +-- | Evaluates the tree given a vector of variable values, a vector of parameter values and a function that takes a Double and change to whatever type the variables have. This is useful when working with datasets of many values per variables. +evalTree :: Columns -> Theta -> Fix SRTree -> Target +evalTree xss params = cata $ + \case + Var ix -> xss !! ix + Param ix -> replicateAs xss $ params V.! ix + Const c -> replicateAs xss c + Y _ -> undefined + Uni g t -> evalFun g t + Bin op l r -> evalOp op l r +{-# INLINE evalTree #-} + +-- evaluates an operator +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,13 @@ {-# language FlexibleInstances, DeriveFunctor #-} {-# language ScopedTypeVariables #-} {-# language RankNTypes #-} -{-# language ViewPatterns #-} +{-# language OverloadedStrings #-} +{-# language LambdaCase #-} +{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-} ----------------------------------------------------------------------------- -- | -- 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,55 +23,59 @@ , 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 ) +import GHC.Generics (Generic) +import Control.DeepSeq (NFData) -- | Tree structure to be used with Symbolic Regression algorithms. -- This structure is a fixed point of a n-ary tree. data SRTree val = - Var Int -- ^ index of the variables - | Param Int -- ^ index of the parameter - | Const Double -- ^ constant value, can be converted to a parameter + Var {-# UNPACK #-} !Int -- ^ index of the variables + | Param {-# UNPACK #-} !Int -- ^ index of the parameter + | Const {-# UNPACK #-} !Double -- ^ constant value, can be converted to a parameter + | Y {-# UNPACK #-} !Int -- ^ index of the target variable, always 0 for now + -- | IConst Int -- TODO: integer constant + -- | RConst Ratio -- TODO: rational constant | Uni Function val -- ^ univariate function | Bin Op val val -- ^ binary operator - deriving (Show, Eq, Ord, Functor) + deriving (Show, Eq, Ord, Functor, Generic, NFData) -- | Supported operators -data Op = Add | Sub | Mul | Div | Power - deriving (Show, Read, Eq, Ord, Enum) +data Op = Add | Sub | Mul | Div | Power | PowerAbs | AQ + deriving (Show, Read, Eq, Ord, Enum, Generic, NFData) -- | Supported functions data Function = @@ -88,12 +94,42 @@ | ACosh | ATanh | Sqrt + | SqrtAbs | Cbrt | Square | Log + | LogAbs | Exp - deriving (Show, Read, Eq, Ord, Enum) + | Recip + | Cube + deriving (Show, Read, Eq, Ord, Enum, Generic, NFData) +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 +138,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 +198,9 @@ l / r = Fix $ Bin Div l r {-# INLINE (/) #-} + recip = Fix . Uni Recip + {-# INLINE recip #-} + fromRational = Fix . Const . fromRational {-# INLINE fromRational #-} @@ -188,6 +247,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 +282,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 +294,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 +351,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 +364,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 +393,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 +406,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.SRTree.Eval +import Control.Monad +import qualified Data.Vector.Unboxed as V + -- * 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 Theta +randomVec n = V.fromList <$> replicateM n (randomRange (-1, 1)) + +randomTree :: Monad m => Int -> Int -> Int -> Rng m (Fix SRTree) -> Rng m (SRTree ()) -> Bool -> Rng m (Fix SRTree) +randomTree minDepth maxDepth maxSize genTerm genNonTerm grow + | 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.hs b/src/Numeric/Optimization/NLOPT.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Optimization/NLOPT.hs @@ -0,0 +1,976 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Module : Numeric.NLOPT +Copyright : (c) Matthew Peddie 2017 +License : BSD3 +Maintainer : Matthew Peddie +Stability : provisional +Portability : GHC + +This module provides a high-level, @hmatrix@-compatible interface to +the library by +Steven G. Johnson. + +NOTE: This is an adaptation from https://hackage.haskell.org/package/hmatrix-nlopt-0.2.0.0 +that removes the dependency to hmatrix and support any Vector Storage. + += Documentation + +Most non-numerical details are documented, but for specific +information on what the optimization methods do, how constraints are +handled, etc., you should consult: + + * The + + * The + + * The + += Example program + +The following interactive session example uses the Nelder-Mead simplex +algorithm, a derivative-free local optimizer, to minimize a trivial +function with a minimum of 22.0 at @(0, 0)@. + +>>> import Numeric.LinearAlgebra ( dot, fromList ) +>>> let objf x = x `dot` x + 22 -- define objective +>>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +>>> let algorithm = NELDERMEAD objf [] Nothing -- specify algorithm +>>> let problem = LocalProblem 2 stop algorithm -- specify problem +>>> let x0 = fromList [5, 10] -- specify initial guess +>>> minimizeLocal problem x0 +Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) + +-} + +module Numeric.Optimization.NLOPT ( + -- * Specifying the objective function + Objective + , ObjectiveD + , Preconditioner + -- * Specifying the constraints + -- ** Bound constraints + , Bounds(..) + -- ** Nonlinear constraints + -- + -- $nonlinearconstraints + + -- *** Constraint functions + , ScalarConstraint + , ScalarConstraintD + , VectorConstraint + , VectorConstraintD + -- *** Constraint types + , Constraint(..) + , EqualityConstraint(..) + , InequalityConstraint(..) + -- *** Collections of constraints + , EqualityConstraints + , EqualityConstraintsD + , InequalityConstraints + , InequalityConstraintsD + -- * Stopping conditions + -- + -- $nonempty + , StoppingCondition(..) + , NonEmpty(..) + -- * Additional configuration + , RandomSeed(..) + , Population(..) + , VectorStorage(..) + , InitialStep(..) + -- * Minimization problems + -- ** Local minimization + , LocalAlgorithm(..) + , LocalProblem(..) + , minimizeLocal + -- ** Global minimization + , GlobalAlgorithm(..) + , GlobalProblem(..) + , minimizeGlobal + -- ** Minimization by augmented Lagrangian + , AugLagAlgorithm(..) + , AugLagProblem(..) + , minimizeAugLag + -- ** Results + , Solution(..) + , N.Result(..) + ) where + +import qualified Numeric.Optimization.NLOPT.Bindings as N + +import Data.List.NonEmpty (NonEmpty(..)) + +import qualified Data.Vector.Storable as V +import Data.Vector.Storable ( Vector ) + +import Control.Exception ( Exception ) +import qualified Control.Exception as Ex +import Data.Typeable ( Typeable ) +import Data.Foldable ( traverse_ ) + +import System.IO.Unsafe ( unsafePerformIO ) + +-- each element i contains a row vec +type Matrix a = [Vector a] + +flatten :: V.Storable a => Matrix a -> Vector a +flatten = V.concat +{-# INLINE flatten #-} + +{- Function wrapping for the immutable HMatrix interface -} +wrapScalarFunction :: (Vector Double -> Double) -> N.ScalarFunction () +wrapScalarFunction f params _ _ = return $ f params + +wrapScalarFunctionD :: (Vector Double -> (Double, Vector Double)) + -> N.ScalarFunction () +wrapScalarFunctionD f params grad _ = do + case grad of + Nothing -> return () + Just g -> V.copy g usergrad + return result + where + (result, usergrad) = f params + +wrapVectorFunction :: (Vector Double -> Word -> Vector Double) + -> Word -> N.VectorFunction () +wrapVectorFunction f n params vout _ _ = V.copy vout $ f params n + +wrapVectorFunctionD :: (Vector Double -> Word -> (Vector Double, Matrix Double)) + -> Word -> N.VectorFunction () +wrapVectorFunctionD f n params vout jac _ = do + V.copy vout result + case jac of + Nothing -> return () + Just j -> V.copy j (flatten userjac) + where + (result, userjac) = f params n + +wrapPreconditionerFunction :: (Vector Double -> Vector Double -> Vector Double) + -> N.PreconditionerFunction () +wrapPreconditionerFunction f params v vpre _ = V.copy vpre (f params v) + +{- Objective functions -} +-- | An objective function that calculates the objective value at the +-- given parameter vector. +type Objective + = Vector Double -- ^ Parameter vector + -> Double -- ^ Objective function value + +-- | An objective function that calculates both the objective value +-- and the gradient of the objective with respect to the input +-- parameter vector, at the given parameter vector. +type ObjectiveD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Objective function value, gradient) + +-- | A preconditioner function, which computes @vpre = H(x) v@, where +-- @H@ is the Hessian matrix: the positive semi-definite second +-- derivative at the given parameter vector @x@, or an approximation +-- thereof. +type Preconditioner + = Vector Double -- ^ Parameter vector @x@ + -> Vector Double -- ^ Vector @v@ to precondition at @x@ + -> Vector Double -- ^ Preconditioned vector @vpre@ + +data ObjectiveFunction f + = MinimumObjective f + | PreconditionedMinimumObjective Preconditioner f + +applyObjective :: N.Opt -> ObjectiveFunction Objective -> IO N.Result +applyObjective opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunction f) () +applyObjective opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunction f) + (wrapPreconditionerFunction p) () + +applyObjectiveD :: N.Opt -> ObjectiveFunction ObjectiveD -> IO N.Result +applyObjectiveD opt (MinimumObjective f) = + N.set_min_objective opt (wrapScalarFunctionD f) () +applyObjectiveD opt (PreconditionedMinimumObjective p f) = + N.set_precond_min_objective opt (wrapScalarFunctionD f) + (wrapPreconditionerFunction p) () + +{- Constraint functions -} +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraint + = Vector Double -- ^ Parameter vector @x@ + -> Double -- ^ Constraint violation (deviation from 0) + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the gradient of @c(x)@ with respect to @x@ at +-- that point. The constraint will enforce that @c(x) == 0@ (equality +-- constraint) or @c(x) <= 0@ (inequality constraint). +type ScalarConstraintD + = Vector Double -- ^ Parameter vector + -> (Double, Vector Double) -- ^ (Constraint violation, constraint gradient) + +-- | A constraint function which returns a vector @c(x)@ given the +-- parameter vector @x@. The constraint will enforce that @c(x) == 0@ +-- (equality constraint) or @c(x) <= 0@ (inequality constraint). +type VectorConstraint + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> Vector Double -- ^ Constraint violation vector + +-- | A constraint function which returns @c(x)@ given the parameter +-- vector @x@ along with the Jacobian (first derivative) matrix of +-- @c(x)@ with respect to @x@ at that point. The constraint will +-- enforce that @c(x) == 0@ (equality constraint) or @c(x) <= 0@ +-- (inequality constraint). +type VectorConstraintD + = Vector Double -- ^ Parameter vector + -> Word -- ^ Constraint Vectorize + -> (Vector Double, Matrix Double) -- ^ (Constraint violation vector, + -- constraint Jacobian) + +-- $nonlinearconstraints +-- +-- Note that most NLOPT algorithms do not support nonlinear +-- constraints natively; if you need to enforce nonlinear constraints, +-- you may want to use the 'AugLagAlgorithm' family of solvers, which +-- can add nonlinear constraints to some algorithm that does not +-- support them by a principled modification of the objective +-- function. +-- +-- == Example program +-- +-- The following interactive session example enforces a scalar +-- constraint on the problem given in the beginning of the module: the +-- parameters must always sum to 1. The minimizer finds a constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = COBYLA objf [] [] [constraint] Nothing +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.500000000013028, solutionParams = [0.5000025521533521,0.49999744784664796], solutionResult = FTOL_REACHED}) + + +data Constraint s v + -- | A scalar constraint. + = Scalar s + -- | A vector constraint. + | Vector Word v + -- | A scalar constraint with an attached preconditioning function. + | Preconditioned Preconditioner s + +-- | An equality constraint, comprised of both the constraint function +-- (or functions, if a preconditioner is used) along with the desired +-- tolerance. +data EqualityConstraint s v = EqualityConstraint + { eqConstraintFunctions :: Constraint s v + , eqConstraintTolerance :: Double + } + +-- | An inequality constraint, comprised of both the constraint +-- function (or functions, if a preconditioner is used) along with the +-- desired tolerance. +data InequalityConstraint s v = InequalityConstraint + { ineqConstraintFunctions :: Constraint s v + , ineqConstraintTolerance :: Double + } + +-- | A collection of equality constraints that do not supply +-- constraint derivatives. +type EqualityConstraints = + [EqualityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of inequality constraints that do not supply +-- constraint derivatives. +type InequalityConstraints = + [InequalityConstraint ScalarConstraint VectorConstraint] + +-- | A collection of equality constraints that supply constraint +-- derivatives. +type EqualityConstraintsD = [EqualityConstraint ScalarConstraintD VectorConstraintD] + +-- | A collection of inequality constraints that supply constraint +-- derivatives. +type InequalityConstraintsD = [InequalityConstraint ScalarConstraintD VectorConstraintD] + +class ApplyConstraint constraint where + applyConstraint :: N.Opt -> constraint -> IO N.Result + +instance ApplyConstraint (EqualityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraint VectorConstraint) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunction s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunction v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunction s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (EqualityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (EqualityConstraint ty tol) = case ty of + Scalar s -> + N.add_equality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_equality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_equality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +instance ApplyConstraint (InequalityConstraint ScalarConstraintD VectorConstraintD) where + applyConstraint opt (InequalityConstraint ty tol) = case ty of + Scalar s -> + N.add_inequality_constraint opt (wrapScalarFunctionD s) () tol + Vector n v -> + N.add_inequality_mconstraint opt n (wrapVectorFunctionD v n) () tol + Preconditioned p s -> + N.add_precond_inequality_constraint opt (wrapScalarFunctionD s) + (wrapPreconditionerFunction p) () tol + +{- Bounds -} + +-- | Bound constraints are specified by vectors of the same dimension +-- as the parameter space. +-- +-- == Example program +-- +-- The following interactive session example enforces lower bounds on +-- the example from the beginning of the module. This prevents the +-- optimizer from locating the true minimum at @(0, 0)@; a slightly +-- higher constrained minimum at @(1, 1)@ is found. Note that the +-- optimizer returns 'N.XTOL_REACHED' rather than 'N.FTOL_REACHED', +-- because the bound constraint is active at the final minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion +-- >>> let lowerbound = LowerBounds $ fromList [1, 1] -- specify bounds +-- >>> let algorithm = NELDERMEAD objf [lowerbound] Nothing -- specify algorithm +-- >>> let problem = LocalProblem 2 stop algorithm -- specify problem +-- >>> let x0 = fromList [5, 10] -- specify initial guess +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 24.0, solutionParams = [1.0,1.0], solutionResult = XTOL_REACHED}) +data Bounds + -- | Lower bound vector @v@ means we want @x >= v@. + = LowerBounds (Vector Double) + -- | Upper bound vector @u@ means we want @x <= u@. + | UpperBounds (Vector Double) + deriving (Eq, Show, Read) + +applyBounds :: N.Opt -> Bounds -> IO N.Result +applyBounds opt (LowerBounds lbvec) = N.set_lower_bounds opt lbvec +applyBounds opt (UpperBounds ubvec) = N.set_upper_bounds opt ubvec + +{- Stopping conditions -} + +-- | A 'StoppingCondition' tells NLOPT when to stop working on a +-- minimization problem. When multiple 'StoppingCondition's are +-- provided, the problem will stop when any one condition is met. +data StoppingCondition + -- | Stop minimizing when an objective value @J@ less than or equal + -- to the provided value is found. + = MinimumValue Double + -- | Stop minimizing when an optimization step changes the objective + -- value @J@ by less than the provided tolerance multiplied by @|J|@. + | ObjectiveRelativeTolerance Double + -- | Stop minimizing when an optimization step changes the objective + -- value by less than the provided tolerance. + | ObjectiveAbsoluteTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than @x@ scaled by the provided + -- tolerance. + | ParameterRelativeTolerance Double + -- | Stop when an optimization step changes /every element/ of the + -- parameter vector @x@ by less than the corresponding element in + -- the provided vector of tolerances values. + | ParameterAbsoluteTolerance (Vector Double) + -- | Stop when the number of evaluations of the objective function + -- exceeds the provided count. + | MaximumEvaluations Word + -- | Stop when the optimization time exceeds the provided time (in + -- seconds). This is not a precise limit. + | MaximumTime Double + deriving (Eq, Show, Read) + +-- $nonempty +-- +-- The 'NonEmpty' data type from 'Data.List.NonEmpty' is re-exported +-- here, because it is used to ensure that you always specify at least +-- one stopping condition. + +applyStoppingCondition :: N.Opt -> StoppingCondition -> IO N.Result +applyStoppingCondition opt (MinimumValue x) = N.set_stopval opt x +applyStoppingCondition opt (ObjectiveRelativeTolerance x) = N.set_ftol_rel opt x +applyStoppingCondition opt (ObjectiveAbsoluteTolerance x) = N.set_ftol_abs opt x +applyStoppingCondition opt (ParameterRelativeTolerance x) = N.set_xtol_rel opt x +applyStoppingCondition opt (ParameterAbsoluteTolerance v) = N.set_xtol_abs opt v +applyStoppingCondition opt (MaximumEvaluations n) = N.set_maxeval opt n +applyStoppingCondition opt (MaximumTime deltat) = N.set_maxtime opt deltat + +{- Random seed control -} + +-- | This specifies how to initialize the random number generator for +-- stochastic algorithms. +data RandomSeed + -- | Seed the RNG with the provided value. + = SeedValue Word + -- | Seed the RNG using the system clock. + | SeedFromTime + -- | Don't perform any explicit initialization of the RNG. + | Don'tSeed + deriving (Eq, Show, Read) + +applyRandomSeed :: RandomSeed -> IO () +applyRandomSeed Don'tSeed = return () +applyRandomSeed (SeedValue n) = N.srand n +applyRandomSeed SeedFromTime = N.srand_time + +{- Random stuff -} + +-- | This specifies the population size for algorithms that use a pool +-- of solutions. +newtype Population = Population Word deriving (Eq, Show, Read) + +applyPopulation :: N.Opt -> Population -> IO N.Result +applyPopulation opt (Population n) = N.set_population opt n + +-- | This specifies the memory size to be used by algorithms like +-- 'LBFGS' which store approximate Hessian or Jacobian matrices. +newtype VectorStorage = VectorStorage Word deriving (Eq, Show, Read) + +applyVectorStorage :: N.Opt -> VectorStorage -> IO N.Result +applyVectorStorage opt (VectorStorage n) = N.set_vector_storage opt n + +-- | This vector with the same dimension as the parameter vector @x@ +-- specifies the initial step for the optimizer to take. (This +-- applies to local gradient-free algorithms, which cannot use +-- gradients to estimate how big a step to take.) +newtype InitialStep = InitialStep (Vector Double) deriving (Eq, Show, Read) + +applyInitialStep :: N.Opt -> InitialStep -> IO N.Result +applyInitialStep opt (InitialStep v) = N.set_initial_step opt v + +{- Algorithms -} + +data GlobalProblem = GlobalProblem + { lowerBounds :: Vector Double -- ^ Lower bounds for @x@ + , upperBounds :: Vector Double -- ^ Upper bounds for @x@ + , gstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , galgorithm :: GlobalAlgorithm -- ^ Algorithm specification + } + +-- | These are the global minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'Population', you can simply specify 'Nothing' to use +-- the default behavior. +data GlobalAlgorithm + -- | DIviding RECTangles + = DIRECT Objective + -- | DIviding RECTangles, locally-biased variant + | DIRECT_L Objective + -- | DIviding RECTangles, "slightly randomized" + | DIRECT_L_RAND Objective RandomSeed + -- | DIviding RECTangles, unscaled version + | DIRECT_NOSCAL Objective + -- | DIviding RECTangles, locally-biased and unscaled + | DIRECT_L_NOSCAL Objective + -- | DIviding RECTangles, locally-biased, unscaled and "slightly + -- randomized" + | DIRECT_L_RAND_NOSCAL Objective RandomSeed + -- | DIviding RECTangles, original FORTRAN implementation + | ORIG_DIRECT Objective InequalityConstraints + -- | DIviding RECTangles, locally-biased, original FORTRAN + -- implementation + | ORIG_DIRECT_L Objective InequalityConstraints + -- | Stochastic Global Optimization. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO ObjectiveD + -- | Stochastic Global Optimization, randomized variant. + -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__ + | STOGO_RAND ObjectiveD RandomSeed + -- | Controlled Random Search with Local Mutation + | CRS2_LM Objective RandomSeed (Maybe Population) + -- | Improved Stochastic Ranking Evolution Strategy + | ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population) + -- | Evolutionary Algorithm + | ESCH Objective + -- | Original Multi-Level Single-Linkage + | MLSL Objective LocalProblem (Maybe Population) + -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy + -- Sequence for starting points + | MLSL_LDS Objective LocalProblem (Maybe Population) + +algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm +algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT +algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L +algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND +algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL +algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL +algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT +algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L +algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO +algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND +algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM +algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES +algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH +algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL +algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS + +applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + + go (DIRECT o) = obj o + go (DIRECT_L o) = obj o + go (DIRECT_NOSCAL o) = obj o + go (DIRECT_L_NOSCAL o) = obj o + go (ESCH o) = obj o + go (STOGO o) = objD o + go (DIRECT_L_RAND o _) = obj o + go (DIRECT_L_RAND_NOSCAL o _) = obj o + go (ORIG_DIRECT o _) = obj o + go (ORIG_DIRECT_L o _) = obj o + go (STOGO_RAND o _) = objD o + go (CRS2_LM o _ _) = obj o + go (ISRES o _ _ _ _) = obj o + go (MLSL o _ _) = obj o + go (MLSL_LDS o _ _) = obj o + +applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO () +applyGlobalAlgorithm opt alg = do + applyGlobalObjective opt alg + go alg + where + seed = applyRandomSeed + pop = maybe (return ()) (tryTo . applyPopulation opt) + ic = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + + local lp = setupLocalProblem lp >>= N.set_local_optimizer opt + + go (DIRECT_L_RAND _ s) = seed s + go (DIRECT_L_RAND_NOSCAL _ s) = seed s + go (ORIG_DIRECT _ ineq) = ic ineq + go (ORIG_DIRECT_L _ ineq) = ic ineq + go (STOGO_RAND _ s) = seed s + go (CRS2_LM _ s p) = seed s *> pop p + go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p + go (MLSL _ lp p) = local lp *> pop p + go (MLSL_LDS _ lp p) = local lp *> pop p + go _ = return () + +tryTo :: IO N.Result -> IO () +tryTo act = do + result <- act + if (N.isSuccess result) + then return () + else Ex.throw $ NloptException result + +data NloptException = NloptException N.Result deriving (Show, Typeable) +instance Exception NloptException + +-- | Solve the specified global optimization problem. +-- +-- = Example program +-- +-- The following interactive session example uses the 'ISRES' +-- algorithm, a stochastic, derivative-free global optimizer, to +-- minimize a trivial function with a minimum of 22.0 at @(0, 0)@. +-- The search is conducted within a box from -10 to 10 in each +-- dimension. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList ) +-- >>> let objf x = x `dot` x + 22 -- define objective +-- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion +-- >>> let algorithm = ISRES objf [] [] (SeedValue 22) Nothing -- specify algorithm +-- >>> let lowerbounds = fromList [-10, -10] -- specify bounds +-- >>> let upperbounds = fromList [10, 10] -- specify bounds +-- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm +-- >>> let x0 = fromList [5, 8] -- specify initial guess +-- >>> minimizeGlobal problem x0 +-- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED}) +minimizeGlobal :: GlobalProblem -- ^ Problem specification + -> Vector Double -- ^ Initial parameter guess + -> Either N.Result Solution -- ^ Optimization results +minimizeGlobal prob x0 = + unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +applyGlobalProblem :: N.Opt -> GlobalProblem -> IO () +applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do + tryTo $ applyBounds opt (LowerBounds lb) + tryTo $ applyBounds opt (UpperBounds ub) + traverse_ (tryTo . applyStoppingCondition opt) stop + applyGlobalAlgorithm opt alg + +newOpt :: N.Algorithm -> Word -> IO N.Opt +newOpt alg sz = do + opt' <- N.create alg sz + case opt' of + Nothing -> Ex.throw $ NloptException N.FAILURE + Just opt -> return opt + +setupGlobalProblem :: GlobalProblem -> IO N.Opt +setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do + opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp) + applyGlobalProblem opt gp + return opt + +solveProblem :: N.Opt -> Vector Double -> IO Solution +solveProblem opt x0 = do + (N.Output outret outcost outx nevals) <- N.optimize opt x0 + if (N.isSuccess outret) + then return $ Solution outcost outx outret nevals + else Ex.throw $ NloptException outret + +minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution +minimizeGlobal' gp x0 = do + opt <- setupGlobalProblem gp + solveProblem opt x0 + +data LocalProblem = LocalProblem + { lsize :: Word -- ^ The dimension of the + -- parameter vector. + , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping + -- condition + , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification + } + +-- | These are the local minimization algorithms provided by NLOPT. Please see +-- +-- for more details on how the methods work and how they relate to one +-- another. Note that some local methods require you provide +-- derivatives (gradients or Jacobians) for your objective function +-- and constraint functions. +-- +-- Optional parameters are wrapped in a 'Maybe'; for example, if you +-- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to +-- use the default behavior. +data LocalAlgorithm + -- | Limited-memory BFGS + = LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage) + -- | Limited-memory BFGS + | LBFGS ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-2 + | VAR2 ObjectiveD (Maybe VectorStorage) + -- | Shifted limited-memory variable-metric, rank-1 + | VAR1 ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method + | TNEWTON ObjectiveD (Maybe VectorStorage) + -- | Truncated Newton's method with automatic restarting + | TNEWTON_RESTART ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method + | TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage) + -- | Preconditioned truncated Newton's method with automatic + -- restarting + | TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage) + -- | Method of moving averages + | MMA ObjectiveD InequalityConstraintsD + -- | Sequential Least-Squares Quadratic Programming + | SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD + -- | Conservative Convex Separable Approximation + | CCSAQ ObjectiveD Preconditioner + -- | PRincipal AXIS gradient-free local optimization + | PRAXIS Objective [Bounds] (Maybe InitialStep) + -- | Constrained Optimization BY Linear Approximations + | COBYLA Objective [Bounds] InequalityConstraints EqualityConstraints + (Maybe InitialStep) + -- | Powell's NEWUOA algorithm + | NEWUOA Objective (Maybe InitialStep) + -- | Powell's NEWUOA algorithm with bounds by SGJ + | NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep) + -- | Nelder-Mead Simplex gradient-free method + | NELDERMEAD Objective [Bounds] (Maybe InitialStep) + -- | NLOPT implementation of Rowan's Subplex algorithm + | SBPLX Objective [Bounds] (Maybe InitialStep) + -- | Bounded Optimization BY Quadratic Approximations + | BOBYQA Objective [Bounds] (Maybe InitialStep) + +algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm +algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL +algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS +algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2 +algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1 +algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON +algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART +algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND +algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART +algorithmEnumOfLocal (MMA _ _) = N.LD_MMA +algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP +algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ +algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS +algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA +algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA +algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD +algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX +algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA + +applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO () +applyLocalObjective opt alg = go alg + where + obj = tryTo . applyObjective opt . MinimumObjective + objD = tryTo . applyObjectiveD opt . MinimumObjective + precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p + + go (LBFGS_NOCEDAL o _) = objD o + go (LBFGS o _) = objD o + go (VAR2 o _) = objD o + go (VAR1 o _) = objD o + go (TNEWTON o _) = objD o + go (TNEWTON_RESTART o _) = objD o + go (TNEWTON_PRECOND o _) = objD o + go (TNEWTON_PRECOND_RESTART o _) = objD o + go (MMA o _) = objD o + go (SLSQP o _ _ _) = objD o + go (CCSAQ o prec) = precond prec o + go (PRAXIS o _ _) = obj o + go (COBYLA o _ _ _ _) = obj o + go (NEWUOA o _) = obj o + go (NEWUOA_BOUND o _ _) = obj o + go (NELDERMEAD o _ _) = obj o + go (SBPLX o _ _) = obj o + go (BOBYQA o _ _) = obj o + +applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO () +applyLocalAlgorithm opt alg = do + applyLocalObjective opt alg + go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + ec = traverse_ (tryTo . applyConstraint opt) + ecd = traverse_ (tryTo . applyConstraint opt) + store = maybe (return ()) (tryTo . applyVectorStorage opt) + bound = traverse_ (tryTo . applyBounds opt) + step0 = maybe (return ()) (tryTo . applyInitialStep opt) + + go (LBFGS_NOCEDAL _ vs) = store vs + go (LBFGS _ vs) = store vs + go (VAR2 _ vs) = store vs + go (VAR1 _ vs) = store vs + go (TNEWTON _ vs) = store vs + go (TNEWTON_RESTART _ vs) = store vs + go (TNEWTON_PRECOND _ vs) = store vs + go (TNEWTON_PRECOND_RESTART _ vs) = store vs + go (MMA _ ineqd) = icd ineqd + go (SLSQP _ b ineqd eqd) = + bound b *> icd ineqd *> ecd eqd + go (CCSAQ _ _ ) = return () + go (PRAXIS _ b s) = bound b *> step0 s + go (COBYLA _ b ineq eq s) = + bound b *> ic ineq *> ec eq *> step0 s + go (NEWUOA _ s) = step0 s + go (NEWUOA_BOUND _ b s) = bound b *> step0 s + go (NELDERMEAD _ b s) = bound b *> step0 s + go (SBPLX _ b s) = bound b *> step0 s + go (BOBYQA _ b s) = bound b *> step0 s + +applyLocalProblem :: N.Opt -> LocalProblem -> IO () +applyLocalProblem opt (LocalProblem _ stop alg) = do + traverse_ (tryTo . applyStoppingCondition opt) stop + applyLocalAlgorithm opt alg + +setupLocalProblem :: LocalProblem -> IO N.Opt +setupLocalProblem lp@(LocalProblem sz _ alg) = do + opt <- newOpt (algorithmEnumOfLocal alg) sz + applyLocalProblem opt lp + return opt + +minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution +minimizeLocal' lp x0 = do + opt <- setupLocalProblem lp + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the SLSQP solver to find the minimum. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale ) +-- >>> let objf x = (x `dot` x + 22, 2 `scale` x) +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1]) +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let algorithm = SLSQP objf [] [] [constraint] +-- >>> let problem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal problem x0 +-- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED}) +minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution +minimizeLocal prob x0 = + unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode + +class ProblemSize c where + problemSize :: c -> Word + +instance ProblemSize LocalProblem where + problemSize = lsize + +instance ProblemSize GlobalProblem where + problemSize = fromIntegral . V.length . lowerBounds + +instance ProblemSize AugLagProblem where + problemSize (AugLagProblem _ _ alg) = case alg of + AUGLAG_LOCAL lp _ _ -> problemSize lp + AUGLAG_EQ_LOCAL lp -> problemSize lp + AUGLAG_GLOBAL gp _ _ -> problemSize gp + AUGLAG_EQ_GLOBAL gp -> problemSize gp + + +-- | __IMPORTANT NOTE__ +-- +-- For augmented lagrangian problems, you, the user, are responsible +-- for providing the appropriate type of constraint. If the +-- subsidiary problem requires an `ObjectiveD`, then you should +-- provide constraint functions with derivatives. If the subsidiary +-- problem requires an `Objective`, you should provide constraint +-- functions without derivatives. If you don't do this, you may get a +-- runtime error. +data AugLagProblem = AugLagProblem + { alEquality :: EqualityConstraints -- ^ Possibly empty set of + -- equality constraints + , alEqualityD :: EqualityConstraintsD -- ^ Possibly empty set of + -- equality constraints with + -- derivatives + , alalgorithm :: AugLagAlgorithm -- ^ Algorithm specification. + } + +-- | The Augmented Lagrangian solvers allow you to enforce nonlinear +-- constraints while using local or global algorithms that don't +-- natively support them. The subsidiary problem is used to do the +-- minimization, but the @AUGLAG@ methods modify the objective to +-- enforce the constraints. Please see +-- +-- for more details on how the methods work and how they relate to one another. +-- +-- See the documentation for 'AugLagProblem' for an important note +-- about the constraint functions. +data AugLagAlgorithm + -- | AUGmented LAGrangian with a local subsidiary method + = AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a local subsidiary method and with + -- penalty functions only for equality constraints + | AUGLAG_EQ_LOCAL LocalProblem + -- | AUGmented LAGrangian with a global subsidiary method + | AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD + -- | AUGmented LAGrangian with a global subsidiary method and with + -- penalty functions only for equality constraints. + | AUGLAG_EQ_GLOBAL GlobalProblem + +algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm +algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ +algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG +algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ + +-- | This structure is returned in the event of a successful +-- optimization. +data Solution = Solution + { solutionCost :: Double -- ^ The objective function value + -- at the minimum + , solutionParams :: Vector Double -- ^ The parameter vector which + -- minimizes the objective + , solutionResult :: N.Result -- ^ Why the optimizer stopped + + , nEvals :: Int -- ^ Number of evaluations until stop + } deriving (Eq, Show, Read) + +applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO () +applyAugLagAlgorithm opt alg = go alg + where + ic = traverse_ (tryTo . applyConstraint opt) + icd = traverse_ (tryTo . applyConstraint opt) + -- AUGLAG won't work at all if you don't pass it the same + -- objective as the subproblem -- here we pull out the subproblem + -- objectives from the algorithm spec and set the same objective + -- function so the user can't mess it up. + local lp = tryTo $ do + localopt <- setupLocalProblem lp + applyLocalObjective opt (lalgorithm lp) + N.set_local_optimizer opt localopt + global gp = do + tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt + applyGlobalObjective opt (galgorithm gp) + + go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_LOCAL lp) = local lp + go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd + go (AUGLAG_EQ_GLOBAL gp) = global gp + +applyAugLagProblem :: N.Opt -> AugLagProblem -> IO () +applyAugLagProblem opt (AugLagProblem eq eqd alg) = do + traverse_ (tryTo . applyConstraint opt) eq + traverse_ (tryTo . applyConstraint opt) eqd + applyAugLagAlgorithm opt alg + +minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution +minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do + opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap) + applyAugLagProblem opt ap + solveProblem opt x0 + +-- | +-- == Example program +-- +-- The following interactive session example enforces the same scalar +-- constraint as the nonlinear constraint example, but this time it +-- uses the augmented Lagrangian method to enforce the constraint and +-- the 'SBPLX' algorithm, which does not support nonlinear constraints +-- itself, to perform the minimization. As before, the parameters +-- must always sum to 1, and the minimizer finds the same constrained +-- minimum of 22.5 at @(0.5, 0.5)@. +-- +-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList ) +-- >>> let objf x = x `dot` x + 22 +-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| [] +-- >>> let algorithm = SBPLX objf [] Nothing +-- >>> let subproblem = LocalProblem 2 stop algorithm +-- >>> let x0 = fromList [5, 10] +-- >>> minimizeLocal subproblem x0 +-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED}) +-- >>> -- define constraint function: +-- >>> let constraintf x = sum (toList x) - 1.0 +-- >>> -- define constraint object to pass to the algorithm: +-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6 +-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem) +-- >>> minimizeAugLag problem x0 +-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED}) + +minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution +minimizeAugLag prob x0 = + unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler + where + handler :: NloptException -> IO (Either N.Result a) + handler (NloptException retcode) = return $ Left retcode diff --git a/src/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,440 @@ +{-# 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, parseNonTerms, showOutput, SRAlgs(..), Output(..) ) -- parsePat, + 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 | NEOGP 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 +parseSR NEOGP header reparam = eitherResult . (`feed` "") . parse (parseNeoGP 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 NeoGP +parseNeoGP :: Bool -> Bool -> [(B.ByteString, Int)] -> ParseTree +parseNeoGP 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-1) + <|> do char 'p' + ix <- decimal + pure $ Fix $ Param (ix-1) + "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 = myParserFun -- 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,66 +1,235 @@ 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.39.6. -- -- see: https://github.com/sol/hpack -name: srtree -version: 1.0.0.5 -synopsis: A general framework to work with Symbolic Regression expression trees. -description: A Symbolic Regression Tree data structure to work with mathematical expressions with support to first order derivative and simplification; -category: Math, Data, Data Structures -homepage: https://github.com/folivetti/srtree#readme -bug-reports: https://github.com/folivetti/srtree/issues -author: Fabricio Olivetti de França -maintainer: fabricio.olivetti@gmail.com -copyright: 2023 Fabricio Olivetti de França -license: BSD3 -license-file: LICENSE -build-type: Simple +name: srtree +version: 3.0.0.2 +synopsis: A general library to work with Symbolic Regression expression trees. +description: A Symbolic Regression Tree data structure to work with mathematical expressions with support to first order derivative and simplification; +license: BSD3 +license-file: LICENSE +author: Fabricio Olivetti de França +maintainer: fabricio.olivetti@gmail.com +copyright: 2023 Fabricio Olivetti de França +category: Math, Data, Data Structures +homepage: https://github.com/folivetti/srtree#readme +bug-reports: https://github.com/folivetti/srtree/issues +build-type: Simple extra-source-files: - README.md - ChangeLog.md + README.md + ChangeLog.md source-repository head - type: git - location: https://github.com/folivetti/srtree + type: git + location: https://github.com/folivetti/srtree library - exposed-modules: - Data.SRTree - Data.SRTree.Internal - Data.SRTree.Print - Data.SRTree.Random - Data.SRTree.Recursion - other-modules: - Paths_srtree - hs-source-dirs: - src - build-depends: - base >=4.16 && <5 - , containers ==0.6.* - , dlist ==1.0.* - , mtl >=2.2 && <2.4 - , random ==1.2.* - , vector >=0.12 && <0.14 - default-language: Haskell2010 + exposed-modules: + Algorithm.EqSat + Algorithm.EqSat.Build + Algorithm.EqSat.DB + Algorithm.EqSat.Egraph + Algorithm.EqSat.Info + Algorithm.EqSat.Queries + Algorithm.EqSat.SearchSR + Algorithm.EqSat.Simplify + Algorithm.EqSat.Store + Algorithm.SRTree.AD + Algorithm.SRTree.AD.CompiledAD + Algorithm.SRTree.AD.Unboxed + Algorithm.SRTree.Compile + Algorithm.SRTree.ConfidenceIntervals + Algorithm.SRTree.Likelihoods + Algorithm.SRTree.ModelSelection + Algorithm.SRTree.NonlinearOpt + Algorithm.SRTree.Utils + Data.SRTree + Data.SRTree.Datasets + Data.SRTree.Derivative + Data.SRTree.Eval + Data.SRTree.Internal + Data.SRTree.Print + Data.SRTree.Random + Data.SRTree.Recursion + Numeric.Optimization.NLOPT + Numeric.Optimization.NLOPT.Bindings + Text.ParseSR + Text.ParseSR.IO + other-modules: + Paths_srtree + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , statistics >=0.15 && <0.17 + , time >=1.9 && <1.15 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + hs-source-dirs: + src + ghc-options: -O2 -fwarn-incomplete-patterns -fspec-constr + extra-libraries: + nlopt + default-language: Haskell2010 +executable bench + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/Bench + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 -fllvm -pgmlo opt-20 -pgmlc llc-20 -optlo-O3 -optlc-mcpu=native -mavx2 -mfma -fspec-constr -fmax-simplifier-iterations=20 -fexpose-all-unfoldings + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , criterion >=1.5 && <2 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + +executable bench-eqsat + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/BenchEqSat + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , criterion >=1.5 && <2 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + +executable srtree-report + main-is: Main.hs + other-modules: + Paths_srtree + hs-source-dirs: + apps/Report + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + build-depends: + async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , optparse-applicative >=0.16 && <0.20 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 + test-suite srtree-test - type: exitcode-stdio-1.0 - main-is: Spec.hs - other-modules: - Paths_srtree - hs-source-dirs: - test - ghc-options: -threaded -rtsopts -with-rtsopts=-N - build-depends: - HUnit - , ad - , base >=4.16 && <5 - , containers ==0.6.* - , dlist ==1.0.* - , mtl >=2.2 && <2.4 - , random ==1.2.* - , srtree - , vector >=0.12 && <0.14 - default-language: Haskell2010 + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + EqSatTests + StoreTests + Paths_srtree + hs-source-dirs: + test + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: + HUnit >=1.6 && <1.7 + , async >=2.2 && <2.3 + , attoparsec >=0.14.4 && <0.15 + , attoparsec-expr >=0.1.1.2 && <0.2 + , base >=4.19 && <5 + , binary >=0.8 && <0.9 + , bytestring >=0.11 && <0.13 + , containers >=0.6.7 && <0.9 + , deepseq >=1.4 && <1.6 + , directory >=1.3 && <1.4 + , exceptions >=0.10 && <0.11 + , filepath >=1.4.0.0 && <1.6 + , hashable >=1.4 && <1.6 + , ieee754 >=0.8 && <0.9 + , lens >=5.0 && <6 + , mtl >=2.2 && <2.4 + , parallel >=3.2 && <3.4 + , primitive >=0.8 && <0.10 + , random >=1.2 && <1.4 + , split >=0.2.5 && <0.3 + , srtree + , statistics >=0.15 && <0.17 + , unordered-containers >=0.2 && <0.3 + , vector >=0.12 && <0.14 + , zlib >=0.6.3 && <0.8 + default-language: Haskell2010 diff --git a/test/EqSatTests.hs b/test/EqSatTests.hs new file mode 100644 --- /dev/null +++ b/test/EqSatTests.hs @@ -0,0 +1,630 @@ +{-# LANGUAGE OverloadedStrings #-} + +module EqSatTests where + +import Test.HUnit +import Data.SRTree +import Data.SRTree.Print (showExpr) +import qualified Data.IntSet as IntSet +import qualified Data.IntMap as IntMap +import qualified Data.Map as Map +import qualified Data.HashSet as Set +import qualified Data.Vector.Unboxed as VU +import qualified Data.Set as RangeSet +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Simplify (simplifyEqSatDefault, rewrites, rewritesParams) +import Control.Monad.State.Strict +import Control.Monad (forM_) +import Control.Monad.Identity +import Data.List (nub, sort) + +eps :: Double +eps = 1e-9 + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin _ l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +runEG :: EGraphST Identity a -> (a, EGraph) +runEG m = runIdentity $ runStateT m emptyGraph + +evalEG :: EGraphST Identity a -> a +evalEG m = runIdentity $ evalStateT m emptyGraph + +-- | Test 1: fromTree with a leaf (variable) +test_fromTree_var :: Test +test_fromTree_var = TestCase $ do + let tree = var 0 + (eid, eg) = runEG $ fromTree myCost tree + assertBool "fromTree var: eid should be >= 0" (eid >= 0) + assertBool "fromTree var: eclass exists" (IntMap.member eid (_eClass eg)) + let ec = _eClass eg IntMap.! eid + assertBool "fromTree var: eclass has nodes" (not $ null (_eNodes ec)) + let bestNode = head $ Set.toList (_eNodes ec) + assertEqual "fromTree var: best is Var 0" (EVar 0) bestNode + +-- | Test 2: fromTree with a binary expression +test_fromTree_bin :: Test +test_fromTree_bin = TestCase $ do + let tree = var 0 + constv 1.0 + (eid, eg) = runEG $ fromTree myCost tree + assertBool "fromTree bin: eid >= 0" (eid >= 0) + let ec = _eClass eg IntMap.! eid + assertBool "fromTree bin: eclass has nodes" (not $ null (_eNodes ec)) + +-- | Test 3: Canonical identity (an e-class should be its own canonical) +test_canonical_identity :: Test +test_canonical_identity = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0) + (canId, _) = runIdentity $ runStateT (canonical eid) eg + assertEqual "canonical of fresh id is itself" eid canId + +-- | Test 4: canonize canonizes children +test_canonize :: Test +test_canonize = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0 + constv 1.0) + (canNode, _) = runIdentity $ runStateT (do + ec <- getEClass eid + let someNode = head $ Set.toList (_eNodes ec) + canonize someNode) eg + -- All children should be canonical now + let children = eChildren canNode + forM_ children $ \c -> do + let (canC, _) = runIdentity $ runStateT (canonical c) eg + assertEqual "canonize: child is canonical" c canC + +-- | Test 5: Adding duplicate e-node returns existing e-class +test_add_duplicate :: Test +test_add_duplicate = TestCase $ do + let tree = constv 2.0 + (eid1, eg1) = runEG $ fromTree myCost tree + (eid2, eg2) = runEG' eg1 $ add myCost (EConst 2.0) + assertEqual "add duplicate returns same eclass" eid1 eid2 + where + runEG' eg m = runIdentity $ runStateT m eg + +-- | Test 6: Merge two distinct e-classes +test_merge :: Test +test_merge = TestCase $ do + let (eid1, eg1) = runEG $ fromTree myCost (var 0) + (eid2, eg2) = runIdentity $ runStateT (fromTree myCost (var 1)) eg1 + assertBool "merge: eid1 and eid2 start different" (eid1 /= eid2) + let (mergedId, eg3) = runIdentity $ runStateT (merge myCost eid1 eid2) eg2 + can1 = _canonicalMap eg3 IntMap.! eid1 + can2 = _canonicalMap eg3 IntMap.! eid2 + assertEqual "merge: canonicals are equal" can1 can2 + assertEqual "merge: leader matches canonical" mergedId can1 + +-- | Test 7: Rebuild after add +test_rebuild :: Test +test_rebuild = TestCase $ do + let tree = var 0 + constv 1.0 + eg = snd $ runEG $ do + _ <- fromTree myCost tree + rebuild myCost + assertBool "rebuild: eNodeToEClass non-empty" (not $ null (_eNodeToEClass eg)) + assertBool "rebuild: worklist empty" (null (_worklist (_eDB eg))) + assertBool "rebuild: analysis empty" (null (_analysis (_eDB eg))) + +-- | Test 8: Basic pattern matching +test_match :: Test +test_match = TestCase $ do + let tree = var 0 + constv 1.0 + pat = Fixed (Bin Add (VarPat 'x') (VarPat 'y')) + (substs, _) = runEG $ do + _ <- fromTree myCost tree + match pat + assertBool "match: should have at least one substitution" (not $ null substs) + +-- | Test 9: Extraction (getBestExpr) +test_getBestExpr :: Test +test_getBestExpr = TestCase $ do + let tree = var 0 + constv 1.0 + (extracted, _) = runEG $ do + eid <- fromTree myCost tree + getBestExpr eid + assertEqual "getBestExpr preserves structure" (showExpr tree) (showExpr extracted) + +-- | Test 10: Equality saturation with x + 0 = x +test_eqsat_x_plus_0 :: Test +test_eqsat_x_plus_0 = TestCase $ do + let tree = var 0 + constv 0.0 + rule = "a" + 0 :=> "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x+0 = x" (showExpr (var 0)) (showExpr best) + +-- | Test 11: Equality saturation with x * 1 = x +test_eqsat_x_times_1 :: Test +test_eqsat_x_times_1 = TestCase $ do + let tree = var 0 * constv 1.0 + rule = "a" * 1 :=> "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x*1 = x" (showExpr (var 0)) (showExpr best) + +-- | Test 12: Fitness and theta storage round-trip +test_fitness_theta :: Test +test_fitness_theta = TestCase $ do + let theta = [VU.fromList [1.0, 2.0]] + (mf, _) = runEG $ do + eid <- fromTree myCost (var 0) + insertFitness eid 0.5 theta + getFitness eid + case mf of + Nothing -> assertFailure "getFitness returned Nothing" + Just f -> assertBool "fitness should be ~0.5" (abs (f - 0.5) < eps) + +-- | Test 13: Insert fitness and check range tree +test_fitness_range :: Test +test_fitness_range = TestCase $ do + let (eg, _) = runEG $ do + eid1 <- fromTree myCost (var 0) + eid2 <- fromTree myCost (constv 1.0) + insertFitness eid1 (-1.0) [] + insertFitness eid2 2.0 [] + gets _eDB + rt = _fitRangeDB eg + case getGreatest rt of + Just (bestFit, _) -> assertBool "fitness range: best is 2.0" (abs (bestFit - 2.0) < eps) + Nothing -> assertFailure "fitness range: non-empty" + +-- | Test 14: getTopFitEClassWithSize +test_top_fit_size :: Test +test_top_fit_size = TestCase $ do + let (eclasses, _) = runEG $ do + eid1 <- fromTree myCost (var 0) -- size 1 + eid2 <- fromTree myCost (constv 1.0) -- size 1 + eid3 <- fromTree myCost (var 0 + constv 1.0) -- size 3 + insertFitness eid1 0.5 [] + insertFitness eid2 1.0 [] + insertFitness eid3 2.0 [] + getTopFitEClassWithSize 1 1 + assertBool "top fit size 1: should have at least one" (not $ null eclasses) + assertEqual "top fit size 1: should be 1 result" 1 (length eclasses) + +-- | Test 15: Bidirectional rule (x + 0 == x) +test_eqsat_comm :: Test +test_eqsat_comm = TestCase $ do + let tree = var 0 + constv 0.0 + rule = "a" + 0 :==: "a" + (best, _) = runEG $ eqSat tree [rule] myCost 5 + assertEqual "eqSat: x+0 == x" (showExpr (var 0)) (showExpr best) + +-- | Test 16: Double negation elimination +test_eqsat_double_neg :: Test +test_eqsat_double_neg = TestCase $ do + -- var 0 - (var 0 - const 2) should simplify via x - (x - y) = y + -- but we don't have that rule. Instead use const folding: + -- (1 + 0) * x = x via x * 1 = x after const folding simplifies 1+0 to 1 + -- Actually let's use a simpler rule set + let tree = (constv 1.0 + constv 0.0) * var 0 -- (1+0)*x + rules = ["a" + 0 :=> "a", "a" * 1 :=> "a"] + (best, _) = runEG $ eqSat tree rules myCost 10 + assertEqual "eqSat: (1+0)*x = x" (showExpr (var 0)) (showExpr best) + +-- | Test 17: fromTrees builds multiple independent trees +test_fromTrees :: Test +test_fromTrees = TestCase $ do + let trees = [var 0, constv 1.0, var 0 + constv 1.0] + (eids, eg) = runEG $ fromTrees myCost trees + assertEqual "fromTrees: three trees" 3 (length eids) + -- each eid should be distinct and valid + let allDistinct = length eids == length (map (\x -> _canonicalMap eg IntMap.! x) eids) + assertBool "fromTrees: distinct eclasses" allDistinct + assertBool "fromTrees: each eid in eClass" (all (`IntMap.member` _eClass eg) eids) + +-- | Test 18: Cost function respects node types +test_cost :: Test +test_cost = TestCase $ do + let (eid, eg) = runEG $ fromTree myCost (var 0) + cost = _cost . _info $ (_eClass eg IntMap.! eid) + assertEqual "cost of Var is 1" 1 cost + +-- | Test 19: getAllExpressionsFrom +test_get_all_expr :: Test +test_get_all_expr = TestCase $ do + let (exprs, _) = runEG $ do + eid <- fromTree myCost (var 0 + constv 1.0) + getAllExpressionsFrom eid + assertBool "getAllExpressionsFrom: non-empty" (not $ null exprs) + assertEqual "getAllExpressionsFrom: includes original" (showExpr (var 0 + constv 1.0)) (showExpr (head exprs)) + +-- | Test 20: sizeFitDB has no stale entries after refit with lower fitness +test_sizeFitDB_no_stale :: Test +test_sizeFitDB_no_stale = TestCase $ do + let (eg, _) = runEG $ do + eid <- fromTree myCost (var 0) -- size = 1 + insertFitness eid 1.0 [] -- insert higher fitness + insertFitness eid 0.5 [] -- refit with lower fitness + gets _eDB + sfd = _sizeFitDB eg + -- size 1 should have exactly 1 entry (the new fitness 0.5) + size1Entries = case IntMap.lookup 1 sfd of + Nothing -> 0 + Just rt -> length (RangeSet.toList rt) + assertEqual "sizeFitDB: size 1 should have 1 entry after refit" 1 size1Entries + -- verify the entry is the new fitness, not the old one + case IntMap.lookup 1 sfd >>= RangeSet.lookupMax of + Nothing -> assertFailure "sizeFitDB: size 1 should have an entry" + Just (f, eId) -> assertBool "sizeFitDB: fitness should be 0.5" (abs (f - 0.5) < eps) + +-- | Test 21: trie paths are canonical after merge+rebuild +-- repair never calls addToDB, so stale non-canonical keys remain in the trie. +-- This test verifies that no stale (non-canonical) keys exist after a merge. +test_trie_no_stale_keys :: Test +test_trie_no_stale_keys = TestCase $ do + let (eg, _) = runEG $ do + eid_a <- fromTree myCost (var 0) -- eclass 0 + eid_0 <- fromTree myCost (constv 0.0) -- eclass 1 + eid_t <- fromTree myCost (addZero (var 0) (constv 0.0)) -- eclass 2 (a+0) + + -- Merge a+0 (2) with a (0), so 2 → canonical 0 + mergedId <- merge myCost eid_t eid_a + rebuild myCost + + -- Add a parent (a+0)*b after the merge + eid_b <- fromTree myCost (var 1) -- eclass 3 + eid_parent <- fromTree myCost (addZero (var 0) (constv 0.0) * var 1) -- (a+0)*b + rebuild myCost + + gets id + can = _canonicalMap eg + staleKeys = getAllStaleTrieKeys can (_patDB $ _eDB eg) + assertBool ("trie: expected exactly 1 stale key (2), got: " <> show staleKeys) (staleKeys == [2]) + +-- | Helper: construct a+0 bypassing Num instance optimization that rewrites +0 to identity +addZero :: Fix SRTree -> Fix SRTree -> Fix SRTree +addZero l r = Fix (Bin Add l r) + +-- | Helper: construct a binary tree bypassing Num instance simplifications +mkBin :: Op -> Fix SRTree -> Fix SRTree -> Fix SRTree +mkBin op l r = Fix (Bin op l r) + +-- | Test 22: multi-atom match works after merge (requires toCanon in intersectAtoms) +test_match_after_merge_multi_atom :: Test +test_match_after_merge_multi_atom = TestCase $ do + let pat = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'a') (Fixed (Const 0.0)))) (VarPat 'b')) + ((substs, _, _, _, _), _) = runEG $ do + eid_a <- fromTree myCost (var 0) + eid_0 <- fromTree myCost (constv 0.0) + eid_t <- fromTree myCost (addZero (var 0) (constv 0.0)) + mergedId <- merge myCost eid_t eid_a + rebuild myCost + eid_b <- fromTree myCost (var 1) + eid_parent <- fromTree myCost (addZero (var 0) (constv 0.0) * var 1) + rebuild myCost + substs <- match pat + pure (substs, (), (), (), ()) + assertBool "match: multi-atom should work after merge" (not $ null substs) + +-- | Test 23: flattened ENAry multiset for a right-nested Add +test_enary_flatten :: Test +test_enary_flatten = TestCase $ do + let tree = mkBin Add (var 0) (mkBin Add (var 1) (var 2)) + (eid, eg) = runEG $ fromTree myCost tree + ec = _eClass eg IntMap.! eid + case _best . _info $ ec of + ENAry EAdd xs -> do + let children = expandedList xs + assertEqual "enary: 3 children" 3 (length children) + assertBool "enary: distinct children" (length (nub children) == length children) + assertBool "enary: sorted children" (children == sort children) + _ -> assertFailure "enary: best should be a 3-ary ENAry EAdd" + +-- | Test 24: commutativity is structural (a+b ≡ b+a, no rules needed) +test_enary_comm :: Test +test_enary_comm = TestCase $ do + let ((c1, c2), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (var 0) (var 1)) + eid2 <- fromTree myCost (mkBin Add (var 1) (var 0)) + a <- canonical eid1 + b <- canonical eid2 + pure (a, b) + assertEqual "comm: a+b == b+a" c1 c2 + +-- | Test 25: associativity flattens (a+b)+c ≡ a+(b+c) ≡ a+(c+b) +test_enary_assoc :: Test +test_enary_assoc = TestCase $ do + let ((c1, c2, c3), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (mkBin Add (var 0) (var 1)) (var 2)) + eid2 <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 1) (var 2))) + eid3 <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 2) (var 1))) + a <- canonical eid1 + b <- canonical eid2 + c <- canonical eid3 + pure (a, b, c) + assertEqual "assoc: (a+b)+c == a+(b+c)" c1 c2 + assertEqual "assoc: (a+b)+c == a+(c+b)" c1 c3 + +-- | Test 26: multiset semantics (x+x is distinct from x) +test_enary_multiset :: Test +test_enary_multiset = TestCase $ do + let ((cX, cXX), _) = runEG $ do + eidX <- fromTree myCost (var 0) + eidXX <- fromTree myCost (mkBin Add (var 0) (var 0)) + a <- canonical eidX + b <- canonical eidXX + pure (a, b) + assertBool "multiset: x+x /= x" (cX /= cXX) + +-- | Test 27: constants fold inside flattened nodes (2+3+x ≡ 5+x) +test_enary_fold_const :: Test +test_enary_fold_const = TestCase $ do + let ((c1, c2), _) = runEG $ do + eid1 <- fromTree myCost (mkBin Add (mkBin Add (constv 2.0) (constv 3.0)) (var 0)) + eid2 <- fromTree myCost (mkBin Add (constv 5.0) (var 0)) + a <- canonical eid1 + b <- canonical eid2 + pure (a, b) + assertEqual "fold-const: 2+3+x == 5+x" c1 c2 + +-- | Test 28: direct add of an unsorted ENAry canonicalizes and folds consts +test_enary_direct_add :: Test +test_enary_direct_add = TestCase $ do + let ((c1, c2), _) = runEG $ do + e2 <- fromTree myCost (constv 2.0) + e3 <- fromTree myCost (constv 3.0) + ex <- fromTree myCost (var 0) + eid <- add myCost (ENAry EAdd (imFromList [e3, ex, e2])) + eid5x <- fromTree myCost (mkBin Add (constv 5.0) (var 0)) + a <- canonical eid + b <- canonical eid5x + pure (a, b) + assertEqual "direct add: ENAry [3,x,2] sorts and folds to 5+x" c1 c2 + +-- | Test 29: extraction of a flattened class right-folds to a binary tree +test_enary_extract :: Test +test_enary_extract = TestCase $ do + let t1 = mkBin Add (var 0) (mkBin Add (var 1) (var 2)) + (extracted, _) = runEG $ do + eid <- fromTree myCost t1 + getBestExpr eid + assertEqual "extract: flattened a+b+c == a+(b+c)" (showExpr t1) (showExpr extracted) + +-- | Test 30: merge cascade propagates through ENAry parents (a≡b -> a+c ≡ b+c) +test_enary_merge_cascade :: Test +test_enary_merge_cascade = TestCase $ do + let ((c1, c2), _) = runEG $ do + ea <- fromTree myCost (var 0) + eb <- fromTree myCost (var 1) + _ <- fromTree myCost (var 2) + eac <- fromTree myCost (mkBin Add (var 0) (var 2)) + ebc <- fromTree myCost (mkBin Add (var 1) (var 2)) + merge myCost ea eb + rebuild myCost + a <- canonical eac + b <- canonical ebc + pure (a, b) + assertEqual "cascade: after a==b, a+c == b+c" c1 c2 + +-- | Soundness: a closed 2-ary pattern (a+b) does NOT match a 3-ary multiset. +test_match_closed2_not_3ary :: Test +test_match_closed2_not_3ary = TestCase $ do + let pat = "a" + "b" + (substs, _) = runEG $ do + x <- fromTree myCost (var 0) + y <- fromTree myCost (var 1) + z <- fromTree myCost (var 2) + _ <- add myCost (ENAry EAdd (imFromList [x, y, z])) + match pat + assertBool "closed2: a+b does not match x+y+z" (null substs) + +-- | Soundness: a+a does NOT match x+x+y (only exact multisets match). +test_match_aa_not_3ary :: Test +test_match_aa_not_3ary = TestCase $ do + let pat = "a" + "a" + (substs, _) = runEG $ do + _ <- fromTree myCost (mkBin Add (var 0) (mkBin Add (var 0) (var 1))) + match pat + assertBool "aa: a+a does not match x+x+y" (null substs) + +-- | B3: 0 + x + y = x + y (n-ary open-rest rule). +test_eqsat_zero_plus_rest :: Test +test_eqsat_zero_plus_rest = TestCase $ do + let tree = addZero (constv 0.0) (addZero (var 0) (var 1)) + assertEqual "0+x+y = x+y" + (showExpr (var 0 + var 1)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B7: xy + xz + w = x(y+z) + w (n-ary factoring with a rest variable). +test_eqsat_factoring :: Test +test_eqsat_factoring = TestCase $ do + let tree = ((var 0 * var 1) + (var 0 * var 2)) + var 3 + assertEqual "xy+xz+w = x(y+z)+w" + (showExpr ((var 0 * (var 1 + var 2)) + var 3)) + (showExpr (simplifyEqSatDefault tree)) + +-- | C9 is a closed 2-ary rule: (x+y+z)^2 is NOT expanded to a binomial. +test_eqsat_binomial_closed2 :: Test +test_eqsat_binomial_closed2 = TestCase $ do + let tree = ((var 0 + var 1) + var 2) ** constv 2.0 + assertEqual "(x+y+z)^2 not expanded" + (showExpr ((var 0 + (var 1 + var 2)) ** constv 2.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | C14: sqrt(x*x) = abs x (closed 2-ary multiset). +test_eqsat_sqrt_square :: Test +test_eqsat_sqrt_square = TestCase $ do + let rule = sqrt (NAry EMul [Ch "x", Ch "x"]) :=> abs "x" + (best, _) = runEG $ eqSat (sqrt (var 0 * var 0)) [rule] myCost 5 + assertEqual "sqrt(x*x) = abs x" (showExpr (abs (var 0))) (showExpr best) + +-- | x/x = 1 and x-x = 0 (constant identities). +test_eqsat_identities :: Test +test_eqsat_identities = TestCase $ do + assertEqual "x/x = 1" (showExpr (constv 1.0)) (showExpr (simplifyEqSatDefault (var 0 / var 0))) + assertEqual "x-x = 0" (showExpr (constv 0.0)) (showExpr (simplifyEqSatDefault (var 0 - var 0))) + +-- | helper: run eqSat with the full rule set and collect every expression +-- in the root eclass (used to assert that a rule "fires" even if a cheaper +-- representative is extracted). +allExprsOf :: Fix SRTree -> [Fix SRTree] +allExprsOf t = fst $ runEG $ do + root <- fromTree myCost t + _ <- runEqSat myCost rewrites 20 + getAllExpressionsFrom root + +-- | C11 fires: log(x*y) expands to log x + log y inside the root eclass. +test_eqsat_log_distributes :: Test +test_eqsat_log_distributes = TestCase $ do + let exprs = allExprsOf (log (var 0 * var 1)) + target = showExpr (log (var 0) + log (var 1)) + assertBool "log(x*y) contains log x + log y" + (any (\e -> showExpr e == target) exprs) + +-- | C12 fires: abs(x*y) expands to abs x * abs y inside the root eclass. +test_eqsat_abs_distributes :: Test +test_eqsat_abs_distributes = TestCase $ do + let exprs = allExprsOf (abs (var 0 * var 1)) + target = showExpr (abs (var 0) * abs (var 1)) + assertBool "abs(x*y) contains abs x * abs y" + (any (\e -> showExpr e == target) exprs) + +-- | C13 fires: (x*y)^z expands to x^z * y^z inside the root eclass. +test_eqsat_pow_distributes :: Test +test_eqsat_pow_distributes = TestCase $ do + let exprs = allExprsOf ((var 0 * var 1) ** constv 2.0) + target = showExpr ((var 0 ** constv 2.0) * (var 1 ** constv 2.0)) + assertBool "(x*y)^2 contains x^2 * y^2" + (any (\e -> showExpr e == target) exprs) + +-- | B9 (a :==: rule): x^2 * x^3 = x^5. +test_eqsat_pow_mul :: Test +test_eqsat_pow_mul = TestCase $ do + let tree = (var 0 ** constv 2.0) * (var 0 ** constv 3.0) + assertEqual "x^2*x^3 = x^5" (showExpr (var 0 ** constv 5.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B11 (a :==: rule): (x^2)^3 = x^6. +test_eqsat_pow_pow :: Test +test_eqsat_pow_pow = TestCase $ do + let tree = (var 0 ** constv 2.0) ** constv 3.0 + assertEqual "(x^2)^3 = x^6" (showExpr (var 0 ** constv 6.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | x^y * x = x^(y+1): x^2 * x = x^3. +test_eqsat_pow_mul_x :: Test +test_eqsat_pow_mul_x = TestCase $ do + let tree = (var 0 ** constv 2.0) * var 0 + assertEqual "x^2*x = x^3" (showExpr (var 0 ** constv 3.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B4: (0*x)*y = 0. +test_eqsat_zero_mul :: Test +test_eqsat_zero_mul = TestCase $ do + let tree = mkBin Mul (mkBin Mul (constv 0.0) (var 0)) (var 1) + assertEqual "(0*x)*y = 0" (showExpr (constv 0.0)) + (showExpr (simplifyEqSatDefault tree)) + +-- | B4 guard: (0*NaN)*x is NOT folded to 0 (NaN invalidates the rest). +test_eqsat_zero_mul_nan :: Test +test_eqsat_zero_mul_nan = TestCase $ do + let tree = mkBin Mul (mkBin Mul (constv 0.0) (constv (0/0))) (var 0) + best = simplifyEqSatDefault tree + assertBool "(0*NaN)*x /= 0" (showExpr best /= showExpr (constv 0.0)) + +-- | rewritesParams: x-x and x/x become Param 0. +test_eqsat_params :: Test +test_eqsat_params = TestCase $ do + let (b1, _) = runEG $ eqSat (var 0 - var 0) rewritesParams myCost 10 + (b2, _) = runEG $ eqSat (var 0 / var 0) rewritesParams myCost 10 + assertEqual "x-x = Param 0 (param mode)" (showExpr (param 0)) (showExpr b1) + assertEqual "x/x = Param 0 (param mode)" (showExpr (param 0)) (showExpr b2) + +-- | Soundness: x*x*y stays as a right-folded Mul, NOT x^2 (B1 is 2-ary only). +test_eqsat_xxy_sound :: Test +test_eqsat_xxy_sound = TestCase $ do + let tree = mkBin Mul (mkBin Mul (var 0) (var 0)) (var 1) + assertEqual "x*x*y stays right-folded" + (showExpr (var 0 * (var 0 * var 1))) + (showExpr (simplifyEqSatDefault tree)) + +-- | Completeness: a*b matches every Mul node inside a merged class. +test_match_complete_multinode :: Test +test_match_complete_multinode = TestCase $ do + let pat = "a" * "b" + (n, _) = runEG $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 2) + _ <- fromTree myCost (var 3) + m1 <- fromTree myCost (var 0 * var 1) + m2 <- fromTree myCost (var 2 * var 3) + _ <- merge myCost m1 m2 + rebuild myCost + s <- match pat + pure (length s) + assertBool "complete: a*b yields all substs in a merged class" (n >= 2) + +-- | helper: find all non-canonical eclass ids in the trie +getAllStaleTrieKeys :: IntMap.IntMap Int -> DB -> [EClassId] +getAllStaleTrieKeys can = concatMap goIntTrie . Map.elems + where + goIntTrie (IntTrie m) = + [k | k <- IntMap.keys m, not (isCanon k)] + ++ concatMap goIntTrie (IntMap.elems m) + isCanon eid = case IntMap.lookup eid can of + Just v -> v == eid + Nothing -> False + +prependLabel :: String -> Test -> Test +prependLabel label t = TestLabel label t + +tests :: Test +tests = TestList + [ prependLabel "fromTree-var" test_fromTree_var + , prependLabel "fromTree-bin" test_fromTree_bin + , prependLabel "canonical-identity" test_canonical_identity + , prependLabel "canonize" test_canonize + , prependLabel "add-duplicate" test_add_duplicate + , prependLabel "merge" test_merge + , prependLabel "rebuild" test_rebuild + , prependLabel "match" test_match + , prependLabel "getBestExpr" test_getBestExpr + , prependLabel "eqsat-x+0" test_eqsat_x_plus_0 + , prependLabel "eqsat-x*1" test_eqsat_x_times_1 + , prependLabel "fitness-theta" test_fitness_theta + , prependLabel "fitness-range" test_fitness_range + , prependLabel "top-fit-size" test_top_fit_size + , prependLabel "eqsat-comm" test_eqsat_comm + , prependLabel "eqsat-double-neg" test_eqsat_double_neg + , prependLabel "fromTrees" test_fromTrees + , prependLabel "cost" test_cost + , prependLabel "getAllExpressions" test_get_all_expr + , prependLabel "sizeFitDB-no-stale" test_sizeFitDB_no_stale + , prependLabel "trie-no-stale-keys" test_trie_no_stale_keys + , prependLabel "match-after-merge" test_match_after_merge_multi_atom + , prependLabel "enary-flatten" test_enary_flatten + , prependLabel "enary-comm" test_enary_comm + , prependLabel "enary-assoc" test_enary_assoc + , prependLabel "enary-multiset" test_enary_multiset + , prependLabel "enary-fold-const" test_enary_fold_const + , prependLabel "enary-direct-add" test_enary_direct_add + , prependLabel "enary-extract" test_enary_extract + , prependLabel "enary-merge-cascade" test_enary_merge_cascade + , prependLabel "match-closed2-3ary" test_match_closed2_not_3ary + , prependLabel "match-aa-not-3ary" test_match_aa_not_3ary + , prependLabel "eqsat-0+rest" test_eqsat_zero_plus_rest + , prependLabel "eqsat-factoring" test_eqsat_factoring + , prependLabel "eqsat-binomial-2ary" test_eqsat_binomial_closed2 + , prependLabel "eqsat-sqrt-square" test_eqsat_sqrt_square + , prependLabel "eqsat-identities" test_eqsat_identities + , prependLabel "eqsat-log-dist" test_eqsat_log_distributes + , prependLabel "eqsat-abs-dist" test_eqsat_abs_distributes + , prependLabel "eqsat-pow-dist" test_eqsat_pow_distributes + , prependLabel "eqsat-pow-mul" test_eqsat_pow_mul + , prependLabel "eqsat-pow-pow" test_eqsat_pow_pow + , prependLabel "eqsat-pow-mul-x" test_eqsat_pow_mul_x + , prependLabel "eqsat-0*mul" test_eqsat_zero_mul + , prependLabel "eqsat-0*mul-NaN" test_eqsat_zero_mul_nan + , prependLabel "eqsat-params" test_eqsat_params + , prependLabel "eqsat-x*x*y-sound" test_eqsat_xxy_sound + , prependLabel "match-complete" test_match_complete_multinode + ] diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -1,76 +1,115 @@ -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 +import Test.HUnit +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as VS +import Data.SRTree.Internal +import Data.SRTree.Recursion (Fix) +import Data.SRTree.Eval (compile) +import Algorithm.SRTree.AD.Unboxed (CompiledTree, compileTree, compileTreeMulti, evalGrad, evalGradVec, evalGradMulti) +import qualified EqSatTests +import qualified StoreTests +import Data.SRTree.Random (randomTree, tossBiased, randomFrom) +import System.Random (mkStdGen) +import Control.Monad.State.Strict (evalStateT) +import Data.SRTree.Datasets (loadDataset) +import Control.Monad (forM_) --- 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 +-- Small epsilon compare for Doubles +eps :: Double +eps = 1e-9 --- values of the evaluated expressions -exprVals :: [Double] -exprVals = map (evalTree xs thetaSingle id . relabelParams) exprs +approxEqual :: [Double] -> [Double] -> Bool +approxEqual a b = and $ zipWith (\x y -> abs (x - y) < eps) a b -refGrad :: [(Double, [Double])] -refGrad = zip exprVals autoDiffSingle +test_compile :: Test +test_compile = TestCase $ do + let xss = [VU.fromList [1.0, 2.0, 3.0]] + tree = var 0 * param 0 + param 1 + theta = VU.fromList [2.0, 0.5] + yhat = compile xss tree theta + got = VU.toList yhat + expected = [2.5, 4.5, 6.5] + assertBool ("compile produced " ++ show got ++ " expected " ++ show expected) (approxEqual got expected) -testDiff :: (Eq a, Show a) => String -> String -> a -> a -> Test -testDiff lbl name a b = TestLabel lbl $ TestCase (assertEqual name a b) +-- Gradient correctness: the compact ctStatic layout must agree with finite +-- differences (objective) and with the row-fused `evalGrad` backend across +-- the vectorized `evalGradVec` and chunked `evalGradMulti` paths. +test_grad :: Test +test_grad = TestCase $ do + let xss = [ VU.fromList [1.0, 2.0, 3.0, 4.0] + , VU.fromList [0.5, 1.5, 2.5, 3.5] + , VU.fromList [2.0, 1.0, 0.5, 0.25] ] + y = VU.fromList [3.1, 5.2, 7.3, 9.4] + -- ((x0 + t0) * exp(x1)) / (x2 + t1) -- mixes static and dynamic subtrees + tree = (var 0 + param 0) * exp (var 1) / (var 2 + param 1) + theta = VS.fromList [1.0, 0.5] + ct = compileTree xss y Nothing tree + cts = compileTreeMulti xss y Nothing tree + (f0, g0) = evalGrad ct theta + (f1, g1) = evalGradVec ct theta + (f2, g2) = evalGradMulti cts theta + -- finite-difference gradient + h = 1e-6 + gfd = VS.toList $ VS.generate (VS.length theta) $ \i -> + let e = VS.fromList (map (\j -> if j == i then h else 0) [0 .. VS.length theta - 1]) + (fp, _) = evalGradVec ct (VS.zipWith (+) theta e) + (fm, _) = evalGradVec ct (VS.zipWith (-) theta e) + in (fp - fm) / (2 * h) + assertBool "evalGradVec objective != evalGrad" (abs (f1 - f0) < 1e-6) + assertBool "evalGradMulti objective != evalGrad" (abs (f2 - f0) < 1e-6) + assertBool "evalGradVec gradient != finite diff" + (and (zipWith (\a b -> abs (a - b) < 1e-4) (VS.toList g1) gfd)) + assertBool "evalGrad gradient != finite diff" + (and (zipWith (\a b -> abs (a - b) < 1e-4) (VS.toList g0) gfd)) -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) +test_benchgrad :: Test +test_benchgrad = TestCase $ do + let genTerm = do coin <- tossBiased 0.4 + if coin then randomFrom [Fix $ Var ix | ix <- [0..8]] else randomFrom [Fix $ Param ix | ix <- [0..9]] + genNonTerm = randomFrom [Bin Add () (), Bin Sub () (), Bin Mul () (), Uni LogAbs (), Uni SqrtAbs ()] + genMultipleTrees 0 = pure [] + genMultipleTrees n = do + t <- randomTree 5 10 150 genTerm genNonTerm False + ts <- genMultipleTrees (n-1) + pure (t:ts) + g = mkStdGen 42 + trees' <- evalStateT (genMultipleTrees 5) g + ((dataset, y, _, _), _, _, _) <- loadDataset "data.tsv" True + let thetaU = VU.fromList [1.0, 0.5, 0.2, 0.3, 0.1, 0.5, 0.9, 0.3, 0.2, 0.4] + thetaS = VS.convert thetaU + trees = map relabelParamsOrder $ filter (\t -> let v = VU.sum (compile dataset t thetaU) in not (isInfinite v || isNaN v)) trees' + h = 1e-6 + gfd :: CompiledTree -> VS.Vector Double + gfd ct = VS.generate (VS.length thetaS) $ \i -> + let e = VS.fromList (map (\j -> if j == i then h else 0) [0 .. VS.length thetaS - 1]) + (fp, _) = evalGradVec ct (VS.zipWith (+) thetaS e) + (fm, _) = evalGradVec ct (VS.zipWith (-) thetaS e) + in (fp - fm) / (2 * h) + forM_ (zip [0..] trees) $ \(i, t) -> do + let ct = compileTree dataset y Nothing t + cts = compileTreeMulti dataset y Nothing t + (f1, g1) = evalGradVec ct thetaS + (f0, g0) = evalGrad ct thetaS + (f2, g2) = evalGradMulti cts thetaS + fd = gfd ct + putStrLn ("benchgrad tree " ++ show i ++ " obj=" ++ show f1) + assertBool ("tree " ++ show i ++ " evalGradVec objective != evalGrad") (abs (f1 - f0) < 1e-6 * max 1 (abs f0)) + assertBool ("tree " ++ show i ++ " evalGradMulti objective != evalGrad") (abs (f2 - f0) < 1e-6 * max 1 (abs f0)) + assertBool ("tree " ++ show i ++ " evalGradVec gradient mismatch") (and (zipWith (\a b -> abs (a - b) < 1e-3 * max 1 (abs a)) (VS.toList g1) (VS.toList fd))) + assertBool ("tree " ++ show i ++ " evalGrad gradient mismatch") (and (zipWith (\a b -> abs (a - b) < 1e-3 * max 1 (abs a)) (VS.toList g0) (VS.toList fd))) + assertBool ("tree " ++ show i ++ " evalGradMulti gradient != evalGrad") (and (zipWith (\a b -> abs (a - b) < 1e-9 * max 1 (abs a)) (VS.toList g0) (VS.toList g2))) main :: IO () main = do - result <- runTestTT tests - putStrLn $ showCounts result + let t1 = TestLabel "compile" test_compile + t2 = TestLabel "grad" test_grad + + counts <- runTestTT $ TestList + [ t1 + , t2 + , TestLabel "benchgrad" test_benchgrad + , TestLabel "eqsat" EqSatTests.tests + , TestLabel "store" StoreTests.tests + ] + if failures counts /= 0 || errors counts /= 0 + then error "Some tests failed" + else pure () diff --git a/test/StoreTests.hs b/test/StoreTests.hs new file mode 100644 --- /dev/null +++ b/test/StoreTests.hs @@ -0,0 +1,169 @@ +{-# LANGUAGE TupleSections #-} + +module StoreTests where + +import Test.HUnit +import Data.SRTree +import qualified Data.IntMap as IntMap +import qualified Data.HashMap.Strict as HashMap +import Algorithm.EqSat +import Algorithm.EqSat.Egraph +import Algorithm.EqSat.Build +import Algorithm.EqSat.DB +import Algorithm.EqSat.Info +import Algorithm.EqSat.Queries +import Algorithm.EqSat.Store +import Control.Monad.State.Strict +import Control.Monad.Identity + +myCost :: SRTree Int -> Int +myCost (Var _) = 1 +myCost (Const _) = 1 +myCost (Param _) = 1 +myCost (Bin _ l r) = 2 + l + r +myCost (Uni _ t) = 3 + t + +-- | run a stateful computation on a specific graph +runIn :: EGraph -> EGraphST Identity a -> (a, EGraph) +runIn g m = runIdentity $ runStateT m g + +evalIn :: EGraph -> EGraphST Identity a -> a +evalIn g m = runIdentity $ evalStateT m g + +-- | graph A: x0, x1, x0+x1 (with fitness on the sum) +buildA :: (EClassId, EGraph) +buildA = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + eidSum <- fromTree myCost (var 0 + var 1) + insertFitness eidSum 0.5 [] + pure eidSum + +-- | graph B: x1, x0+x1, (x0+x1)*x2 (shares x1 and x0+x1 with A) +buildB :: EGraph +buildB = snd $ runIn emptyGraph $ do + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 0 + var 1) + _ <- fromTree myCost ((var 0 + var 1) * var 2) + pure () + +-- | pattern (x0+x1)*x2 = (A + B) * C +prodPattern :: Pattern +prodPattern = Fixed (Bin Mul (Fixed (Bin Add (VarPat 'A') (VarPat 'B'))) (VarPat 'C')) + +-- | Test 1: export/import round-trip preserves the rows exactly +test_roundtrip :: Test +test_roundtrip = TestCase $ do + let (_, g) = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + _ <- fromTree myCost (var 1) + _ <- fromTree myCost (var 0 + var 1) + _ <- fromTree myCost ((var 0 + var 1) * var 2) + pure () + rows = exportEGraph g + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right g' -> do + let rows' = exportEGraph g' + assertBool "round-trip: rows differ" (rows == rows') + assertBool "round-trip: class count" (IntMap.size (_grEClasses rows) == IntMap.size (_grEClasses rows')) + assertBool "round-trip: node count" (HashMap.size (_grENodeToEClass rows) == HashMap.size (_grENodeToEClass rows')) + +-- | Test 2: round-trip preserves fitness and rebuilds the range DB +test_roundtrip_fitness :: Test +test_roundtrip_fitness = TestCase $ do + let (sumEid, g) = runIn emptyGraph $ do + eidSum <- fromTree myCost (var 0 + var 1) + insertFitness eidSum 0.42 [] + pure eidSum + rows = exportEGraph g + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right g' -> do + let fit = evalIn g' (getFitness sumEid) + assertEqual "round-trip: fitness" (Just 0.42) fit + let mx = getGreatest (_fitRangeDB (_eDB g')) + assertEqual "round-trip: fitRangeDB max" (Just (0.42, sumEid)) mx + -- a node added *after* import dedups against the loaded graph (no dup class) + let (eidNew, g'') = runIn g' $ fromTree myCost (var 0 + var 1) + nClasses = IntMap.size (_eClass g'') + assertBool "post-import dedup adds no class" (eidNew == sumEid && nClasses == IntMap.size (_eClass g')) + +-- | Test 3: import rejects inconsistent rows +test_import_invalid :: Test +test_import_invalid = TestCase $ do + let (_, g) = runIn emptyGraph $ do + _ <- fromTree myCost (var 0) + pure () + rows = exportEGraph g + bad = rows { _grENodeToEClass = HashMap.insert (EVar 0) 999 (_grENodeToEClass rows) } -- 999 not in canonical map + case importEGraph bad of + Left _ -> pure () + Right _ -> assertFailure "invalid rows should have been rejected" + +-- | Test 4: merge dedups shared structure and adds only new classes +test_merge :: Test +test_merge = TestCase $ do + let (sumEidA, gA) = buildA + gM = case mergeEGraph myCost gA buildB of + Left err -> error ("merge failed: " ++ err) + Right g -> g + nA = IntMap.size (_eClass gA) + nM = IntMap.size (_eClass gM) + assertEqual "merge: adds only classes absent from A (x2, product)" (nA + 2) nM + -- B's unique expression (x0+x1)*x2 is present and matchable + let nMatch = length $ evalIn gM (match prodPattern) + assertBool "merge: B's unique expression present" (nMatch > 0) + -- A's fitness on the shared sum class is preserved (same canonical id) + assertEqual "merge: A fitness preserved" (Just 0.5) (evalIn gM (getFitness sumEidA)) + +-- | Test 5: merge preserves round-trip +test_merge_roundtrip :: Test +test_merge_roundtrip = TestCase $ do + let (_, gA) = buildA + gM = case mergeEGraph myCost gA buildB of + Left err -> error ("merge failed: " ++ err) + Right g -> g + rows = exportEGraph gM + case importEGraph rows of + Left err -> assertFailure ("import failed: " ++ err) + Right gM' -> assertBool "merge round-trip: rows differ" (exportEGraph gM' == rows) + +-- | Test 6: stale node->class entries (a node pointing at a class whose +-- canonical representative is another class) are canonicalized on import +test_import_stale_canonicalizes :: Test +test_import_stale_canonicalizes = TestCase $ do + let (keep, g) = buildA -- keep = x0+x1, a root class, has fitness + rows0 = exportEGraph g + dead = _grNextId rows0 -- a fresh id not yet in the graph + rows = rows0 { _grCanonical = IntMap.insert dead keep (_grCanonical rows0) + , _grEClasses = IntMap.insert dead + (IntMap.findWithDefault (error "keep missing") keep (_grEClasses rows0)) + (_grEClasses rows0) + , _grENodeToEClass = HashMap.insert (EBin Add 2 3) dead (_grENodeToEClass rows0) + , _grNextId = dead + 1 } + case importEGraph rows of + Left err -> assertFailure ("import of stale rows failed: " ++ err) + Right g' -> do + let canon = _grCanonical (exportEGraph g') + posts = exportEGraph g' + deadNext = IntMap.lookup dead (_grEClasses posts) + -- the dead class is gone and every node points at a canonical class + assertEqual "dead class dropped" Nothing deadNext + assertBool "all node->class values canonical" + (all (\eid -> IntMap.lookup eid canon == Just eid) (HashMap.elems (_grENodeToEClass posts))) + -- the kept class is still there with its fitness (via the fit range db) + assertEqual "kept fitness preserved" (Just 0.5) (evalIn g' (getFitness keep)) + +prependLabel :: String -> Test -> Test +prependLabel label t = TestLabel label t + +tests :: Test +tests = TestList + [ prependLabel "store-roundtrip" test_roundtrip + , prependLabel "store-roundtrip-fit" test_roundtrip_fitness + , prependLabel "store-import-invalid" test_import_invalid + , prependLabel "store-merge" test_merge + , prependLabel "store-merge-roundtrip" test_merge_roundtrip + , prependLabel "store-stale-canon" test_import_stale_canonicalizes + ]