srtree 1.0.0.5 → 3.0.0.2
raw patch · 40 files changed
Files
- ChangeLog.md +92/−0
- LICENSE +3/−4
- README.md +269/−14
- apps/Bench/Main.hs +127/−0
- apps/BenchEqSat/Main.hs +259/−0
- apps/Report/Main.hs +247/−0
- src/Algorithm/EqSat.hs +380/−0
- src/Algorithm/EqSat/Build.hs +743/−0
- src/Algorithm/EqSat/DB.hs +695/−0
- src/Algorithm/EqSat/Egraph.hs +925/−0
- src/Algorithm/EqSat/Info.hs +222/−0
- src/Algorithm/EqSat/Queries.hs +213/−0
- src/Algorithm/EqSat/SearchSR.hs +277/−0
- src/Algorithm/EqSat/Simplify.hs +284/−0
- src/Algorithm/EqSat/Store.hs +243/−0
- src/Algorithm/SRTree/AD.hs +32/−0
- src/Algorithm/SRTree/AD/CompiledAD.hs +39/−0
- src/Algorithm/SRTree/AD/Unboxed.hs +965/−0
- src/Algorithm/SRTree/Compile.hs +106/−0
- src/Algorithm/SRTree/ConfidenceIntervals.hs +422/−0
- src/Algorithm/SRTree/Likelihoods.hs +329/−0
- src/Algorithm/SRTree/ModelSelection.hs +198/−0
- src/Algorithm/SRTree/NonlinearOpt.hs +98/−0
- src/Algorithm/SRTree/Utils.hs +320/−0
- src/Data/SRTree.hs +25/−23
- src/Data/SRTree/Datasets.hs +416/−0
- src/Data/SRTree/Derivative.hs +140/−0
- src/Data/SRTree/Eval.hs +393/−0
- src/Data/SRTree/Internal.hs +310/−299
- src/Data/SRTree/Print.hs +120/−35
- src/Data/SRTree/Random.hs +81/−27
- src/Data/SRTree/Recursion.hs +12/−0
- src/Numeric/Optimization/NLOPT.hs +976/−0
- src/Numeric/Optimization/NLOPT/Bindings.hs +1075/−0
- src/Text/ParseSR.hs +440/−0
- src/Text/ParseSR/IO.hs +73/−0
- srtree.cabal +223/−54
- test/EqSatTests.hs +630/−0
- test/Spec.hs +107/−68
- test/StoreTests.hs +169/−0
ChangeLog.md view
@@ -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
LICENSE view
@@ -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
README.md view
@@ -1,28 +1,283 @@-# srtree: A symbolic regression expression tree structure.+# srtree: A supporting library for tree-based symbolic regression -`srtree` is a Haskell library with a data structure and supporting functions to manipulate expression trees for symbolic regression.+`srtree` is a Haskell library that implements a tree-based structure for expressions and supporting functions to be used in the context of **symbolic regression** (SR). -The tree-like structure is defined as a fixed-point of an n-ary tree. The variables and parameters of the regression model are indexed as `Int`type and the constant values are `Double`.+This repository is also the home for different algorithm implementations for SR and software tools to support the post-processing of SR models (please refer to their corresponding README files): -The tree supports leaf nodes containing a variable, a free parameter, or a constant value; internal nodes that represents binary operators such as the four basic math operations, logarithm with custom base, and the power of two expressions; and unary functions specified by `Function` data type.+- [srsimplify](apps/srsimplify/README.md): a parser and simplification tool supporting the output of many popular SR algorithms.+- [srtools](apps/srtools/README.md): a tool that can be used to evaluate symbolic regression expressions and create nice reports with confidence intervals. +- [tinygp](apps/tinygp/README.md): a simple GP implementation based on tinyGP.+- [rEGGression](https://github.com/folivetti/reggression/blob/main/README.md): nonlinear regression models exploration and query system with e-graphs (egg).+- [symregg](https://github.com/folivetti/symregg/blob/main/README.md): Equality graph Assisted Search Technique for Equation Recovery.+- [eggp](https://github.com/folivetti/eggp/blob/main/README.md): E-graph Genetic Programming. -The `SRTree` structure has instances for `Num, Fractional, Floating` which allows to create an expression as a valid Haskell expression such as:+## SRTree +The expression structure is defined as a fixed-point of a mix of unary and binary tree. This makes it easier to implement supporting functions that requires the traversal of the trees. Also, since it is a parameterized structure, we can creating partial trees to pattern math structures of interest.+This structure may contain four types of nodes:++- `Bin Op l r` that represents a binary operator `Op` with two children.+- `Uni Function t` that represents an unary function `Function` with a single child.+- `Var Int` representing the index of a variable (i.e., x0, x1, etc.).+- `Param Int` representing the index of a adjustable parameter (i.e., theta0, theta1, etc.).+- `Const Double` representing a constant value.+++The `SRTree` structure has instances for `Num, Fractional, Floating, IsString` which allows to create an expression as a valid Haskell expression such as (remember to turn on OverloadedStrings extension):+ ```haskell-x = var 0)-y = var 1-expr = x * 2 + sin(y * pi + x) :: Fix SRTree+expr = "x0" * 2 + sin("x1" * pi + "x0") :: Fix SRTree ``` -## Other features:+This library comes with support to many quality of life functions to handle this data structure. Such as: -- derivative w.r.t. a variable (`deriveByVar`) and w.r.t. a parameter (`deriveByParam`)-- evaluation (`evalTree`)-- relabel free parameters sequentially (`relabelParams`)-- gradient calculation with `forwardMode`, or optimized with `gradParams` if there is only a single occurrence of each parameter (most of the cases).+- getting the arity of a node+- getting the children of a node as a list+- count the number of nodes +- number of nodes of a specific type+- counting unique tokens +- number of variables and parameters +- relabeling the parameters from 0 to p +- converting floating point constants to parameters +Additionally, the library provides supporting function to work with datasets, evaluating the expressions, +calculating the derivatives, printing, generating random trees, simplifying the expression, calculating overall statistics,+optimizing parameters, and model selection metrics. +++## Organization++The library is organized as `Data`, `Algorithm`, and `Text` modules where the `Data` modules implement functions directly tied to the data structure and the `Algorithm` modules implement algorithms related to symbolic regression, finally, the `Text` modules parse string expressions from different formats and apply simplification, when requested.++### `Data` modules++The `Data` modules is split into $5$ submodules:++- `Data.SRTree` contains the data strucuture and basic supporting functions.+- `Data.SRTree.Datasets` contains functions supporting loading datasets into Massiv.Arrays (aka numpy arrays).+- `Data.SRTree.Derivative` contains the symbolic derivatives of the functions and operators.+- `Data.SRTree.Eval` contains functions to evaluate the tree given a dataset and parameters.+- `Data.SRTree.Print` contains supporting functions for converting trees to different string representation.+- `Data.SRTree.Random` contains functions to generate random trees.++#### `Data.SRTree`++The `SRTree val` data structure is a sum type structure that can be either a variable index, a parameter index, a constant value (of type `Double`), an univariate function or a binary operator. The data type is implemented as a fixed point so all the algorithms act on `Fix SRTree`:++```haskell +t = "x0" + "t0" * sin("x1" + "t1"**2) :: Fix SRTree +```++When creating the expression in a more natural notation, the variables and parameters are `String` composed of the first letter either `x`, for variables, or `t` for parameters (as in theta), and an integer corresponding to the index of the variable or parameter. The fixed point notation, allows us to implment recursive processing of a tree without many of the common boilerplate:++```haskell+countNodes = + \case + Var _ = 1+ Const _ = 1+ Param _ = 1+ Uni _ t = 1 + t + Bin _ l r = 1 + l + r+```++The children are parameterized by the `val` type parameter. This allows us to create convenient partial structures, such as:++```haskell+-- + operator pointing to some structure+-- with index 1 and 2+Bin Add 1 2 ++-- canonical representation of + operator +Bin Add () ()+```++The main functions of this module are:++- `arity`: returns the arity of an operator.+- `getChildren`: returns the children of a `Fix SRTree` as a list +- `countNodes`: returns the number of nodes +- `countOccurrences`: counts the occurence of a given variable +- `countVars`: returns the number of unique variables appearing the expression +- `relabelParams`: relabels the parameters from the left leaves to the right +- `constsToParams`: replace `Const` nodes with `Param` nodes.++#### `Data.SRTree.Datasets` module ++This module exports only the `loadDataset` function which takes a filename and +returns the training and test sets together with the column labels.+The filename must follow the format:++`filename.ext:start_row:end_row:target:features`++where each ':' field is optional. The fields are:++- **start_row:end_row** is the range of the training rows (default 0:nrows-1).+ every other row not included in this range will be used as validation+- **target** is either the name of the PVector (if the datafile has headers) or the index+ of the target variable+- **features** is a comma separated list of SRMatrix names or indices to be used as+ input variables of the regression model.++Example of valid names: `dataset.csv`, `mydata.tsv`, `dataset.csv:20:100`, `dataset.tsv:20:100:price:m2,rooms,neighborhood`, `dataset.csv:::5:0,1,2`.++#### `Data.SRTree.Derivative` module ++Calculates symbolic derivatives of the expression w.r.t. the variables or the parameters.+The main functions of this module are:++- `deriveBy`: returns the symbolic derivative w.r.t. a certain variable or a certain parameter.+- `deriveByVar`: shortcut to `deriveBy` to derive by a variable.+- `deriveByParam`: shortcut to `deriveBy` to derive by a parameter.++#### `Data.SRTree.Eval` module ++Evaluates an expression given a dataset.+The main functions of this module are:++- `evalTree`: given a data matrix and a vector of parameters, evaluates the expression tree.+- `evalInverse`: evaluates the inverse of a function. +- `invright`: evaluates the right inverse of an operator.+- `invleft`: evaluates the left inverse of an operator.++#### `Data.SRTree.Print` module ++Support functions to convert an expression tree into a `String`.+The main functions of this module are: ++- `showExpr` and `printExpr`: converts/print the expression into math notation .+- `showPython` and `printPython`: converts/print to a numpy notation.+- `showLatex` and `printLatex`: converts/print to a LaTeX notation.+- `showTikz` and `printTikz`: converts/print to a TikZ notation.++#### `Data.SRTree.Random` module ++Auxiliary functions to create random trees. +The main functions of this module are:++- `randomTree`: creates a random tree with a certain number of nodes.+- `randomTreeBalanced`: creates a (almost) balanced random tree with a certain number of nodes.++### `Text` modules ++The `Text` module is split into $2$ modules:++- `Text.ParseSR`: contains the main parsers for different SR algorithms output.+- `Text.ParseSR.IO`: auxiliary functions to handle files containing many expressions.++#### `Text.ParseSR` module ++The only important function of this module is `parseSR` that parses an string expression from a given algorithm to a certain output. It also converts variable names to x0, x1,...++#### `Text.ParseSR.IO` module ++The two main functions of this module are: ++- `withInput`: that reads the stdin or a text file and parse all expressions +- `withOutput`: that writes the parsed expression into stdout or a file with one of the choices of output format. ++These functions handle any errors with an `Either` type and they can be safely pipelined together. Any invalid expression will be printed as "invalid expression <error message>".++### `Algorithm` modules ++The `Algorithm` modules are split into $5$ submodules:++- `Algorithm.SRTree.AD` contains automatic differentiation functions.+- `Algorithm.SRTree.ConfidenceIntervals` contains functions to calculate the confidence intervals of parameters and predictions of a symbolic expression using Laplace approximation or profile likelihood.+- `Algorithm.SRTree.Likelihood` contains functions support different likelihood functions and their derivatives (gradient and hessian).+- `Algorithm.SRTree.ModelSelection` implements different model selection criteria such as AIC, BIC, MDL.+- `Algorithm.SRTree.Opt` implements functions to optimize the parameters of an expression supporting different likelihood functions.++#### `Algorithm.SRTree.AD` module ++The main functions of this module are:++- `forwardMode`: returns the prediction errors vector multiplied by the Jacobian matrix using forward mode AD.+- `forwardModeUnique`: same as above, but assuming each parameter index appear only once in the tree. +- `reverseModeUnique`: same as above, but using reverse mode +- `forwardModeUniqueJac`: same as `forwardModeUnique` but returns the Jacobian (does not mutiply by the error).++#### `Algorithm.SRTree.Likelihood` module ++The main functions of this module are: ++- `sse, mse, rmse`: calculates the sum-of-square, mean squared, root of mean squared errors.+- `nll`: returns the negative log-likelihood given a distribution and the associated error (`Nothing` if unknown)+- `gradNLL`: returns the gradient of the negative log-likelihood. +- `gradNLLNonUnique`: same as above but assumes non-unique parameters +- `hessianNLL`: returns the hessian of the neg log-likelihood.++#### `Algorithm.SRTree.Opt` module ++The main functions of this module are:++- `minimizeNLL`: minimizes the negative log-likelihood of a distribution.+- `minimizeNLLNonUnique`: same as above but assumes repeated occurrences of parameters.+- `minimizeNLLWithFixedParam`: minimizes the neg log-likelihood but fixing the value of a single parameter.+- `minimizeGaussian`, `minimizePoisson`, `minimizeBinomial`: shortcut to minimize these three distributions.++#### `Algorithm.SRTree.ModelSelection` module ++The main functions of this module are:++- `bic`: Bayesian Information Criteria +- `aic`: Akaike Information Criteria +- `mdl`: Minimum Description Length as described in Bartlett, Deaglan J., Harry Desmond, and Pedro G. Ferreira. "Exhaustive symbolic regression." IEEE Transactions on Evolutionary Computation (2023)+- `mdlLattice`: as described in Bartlett, Deaglan, Harry Desmond, and Pedro Ferreira. "Priors for symbolic regression." Proceedings of the Companion Conference on Genetic and Evolutionary Computation. 2023.+- `mdlFreq` : MDL weighted by the frequency of occurrence of functions ++#### `Algorithm.SRTree.ConfidenceIntervals` module ++The main functions of this module are: ++- `paramCI`: calculates the parameters confidence intervals. +- `predictionCI`: calculates the predictions confidence intervals ++### `EqSat` modules ++The `EqSat` modules are split into $4$ submodules:++- `Algorithm.EqSat.Simplify` contains function supporting algebraic simplification with equality saturation.+- `Algorithm.EqSat` contains the main equality saturation function.+- `Algorithm.EqSat.EGraph` contains the e-graph data structure and supporting functions.+- `Algorithm.EqSat.EqSatDB` contains supporting functions to pattern matching and insert equivalent expressions into an e-graph. ++#### `Algorithm.EqSat` module++The main functions of this module are: ++- `eqSat` : runs equality saturation over a single expression. +- `getBest` : returns the best expression given the cost function used to generate the e-graph +- `recalculateBest` : recalculates the cost of each e-class using a new cost function +- `runEqSat` : runs equality saturation inside `EGraphST` monad. Use this if you want to return the e-graph. ++#### `Algorithm.EqSat.EGraph` module++The main functions of this module are: ++- `fromTree` : creates an e-graph from an expression tree.+- `fromTrees` : creates an e-graph from multiple expressions +- `fromTreeWith` : inserts a new expression into the e-graph +- `findRootClasses` : returns the roots of the e-graph, if any .+- `getExpressionFrom` : returns a single expression from a given e-class always picking the first e-node as the path +- `getAllExpressionsFrom` : returns all expressions from the given e-class +- `getRndExpressionFrom` : returns a random expression from this e-class +++#### `Algorithm.EqSat.EqSatDB` module++The main functions of this module are: ++- TODO: create auxiliary functions to apply substution rules inside an EGraphST monad . ++#### `Algorithm.EqSat.Simplify` module++The main functions of this module are: ++- `simplifyEqSatDefault` : simplifies an expression using the default parameters +- `simplifyEqSat` : simplifies with custom parameters+ ## TODO: - support more advanced functions - support conditional branching (`IF-THEN-ELSE`)-+- document egraph-search and ieeexplore
+ apps/Bench/Main.hs view
@@ -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++ ]+ ]
+ apps/BenchEqSat/Main.hs view
@@ -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)
+ apps/Report/Main.hs view
@@ -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 ""
+ src/Algorithm/EqSat.hs view
@@ -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
+ src/Algorithm/EqSat/Build.hs view
@@ -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 #-}
+ src/Algorithm/EqSat/DB.hs view
@@ -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 <> " | <cond>"++-- 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 #-}
+ src/Algorithm/EqSat/Egraph.hs view
@@ -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 #-}
+ src/Algorithm/EqSat/Info.hs view
@@ -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)++
+ src/Algorithm/EqSat/Queries.hs view
@@ -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)+
+ src/Algorithm/EqSat/SearchSR.hs view
@@ -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''
+ src/Algorithm/EqSat/Simplify.hs view
@@ -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
+ src/Algorithm/EqSat/Store.hs view
@@ -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
+ src/Algorithm/SRTree/AD.hs view
@@ -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
+ src/Algorithm/SRTree/AD/CompiledAD.hs view
@@ -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)+ }
+ src/Algorithm/SRTree/AD/Unboxed.hs view
@@ -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)
+ src/Algorithm/SRTree/Compile.hs view
@@ -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 #-}
+ src/Algorithm/SRTree/ConfidenceIntervals.hs view
@@ -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+
+ src/Algorithm/SRTree/Likelihoods.hs view
@@ -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)+
+ src/Algorithm/SRTree/ModelSelection.hs view
@@ -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 #-}
+ src/Algorithm/SRTree/NonlinearOpt.hs view
@@ -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+
+ src/Algorithm/SRTree/Utils.hs view
@@ -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
src/Data/SRTree.hs view
@@ -1,7 +1,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -16,28 +16,29 @@ , Op(..) , param , var+ , constv , arity , getChildren+ , childrenOf+ , replaceChildren+ , getOperator , countNodes , countVarNodes , countConsts , countParams+ , countParamsUniq , countOccurrences- , deriveBy- , deriveByVar- , deriveByParam- , derivative- , forwardMode- , gradParamsFwd- , gradParamsRev- , evalFun- , evalOp- , inverseFunc- , evalTree+ , countUniqueTokens+ , numberOfVars+ , getIntConsts , relabelParams+ , relabelParamsOrder+ , relabelVars , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) ) where@@ -48,27 +49,28 @@ , Op(..) , param , var+ , constv , arity , getChildren+ , childrenOf+ , replaceChildren+ , getOperator , countNodes , countVarNodes , countConsts , countParams+ , countParamsUniq , countOccurrences- , deriveBy- , deriveByVar- , deriveByParam- , derivative- , forwardMode- , gradParamsFwd- , gradParamsRev- , evalFun- , evalOp- , inverseFunc- , evalTree+ , countUniqueTokens+ , numberOfVars+ , getIntConsts , relabelParams+ , relabelParamsOrder+ , relabelVars , constsToParam , floatConstsToParam , paramsToConst+ , removeProtectedOps+ , convertProtectedOps , Fix (..) )
+ src/Data/SRTree/Datasets.hs view
@@ -0,0 +1,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 #-}
+ src/Data/SRTree/Derivative.hs view
@@ -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 #-}
+ src/Data/SRTree/Eval.hs view
@@ -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 #-}+ +
src/Data/SRTree/Internal.hs view
@@ -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
src/Data/SRTree/Print.hs view
@@ -1,7 +1,9 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE LambdaCase #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Print --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -12,52 +14,103 @@ ----------------------------------------------------------------------------- module Data.SRTree.Print ( showExpr+ , showExprWithVars , printExpr+ , printExprWithVars , showTikz , printTikz , showPython , printPython , showLatex+ , showLatexWithVars , printLatex+ , showOp ) where -import Control.Monad.Reader ( asks, runReader, Reader )-import Data.Char ( toLower )-+import Control.Monad.Reader (Reader, asks, runReader)+import Data.Char (toLower) import Data.SRTree.Internal-import Data.SRTree.Recursion+import Data.SRTree.Recursion (cata) +-- | converts a tree with protected operators to+-- a conventional math tree+removeProtection :: Fix SRTree -> Fix SRTree+removeProtection = cata $+ \case+ Var ix -> Fix (Var ix)+ Param ix -> Fix (Param ix)+ Const x -> Fix (Const x)+ Uni SqrtAbs t -> sqrt (abs t)+ Uni LogAbs t -> log (abs t)+ Uni Cube t -> t ** 3+ Uni f t -> Fix (Uni f t)+ Bin AQ l r -> l / sqrt (1 + r*r)+ Bin PowerAbs l r -> abs l ** r+ Bin op l r -> Fix (Bin op l r)++-- | convert a tree into a string in math notation +--+-- >>> showExpr $ "x0" + sin ( tanh ("t0" + 2) )+-- "(x0 + Sin(Tanh((t0 + 2.0))))" showExpr :: Fix SRTree -> String-showExpr = cata alg- where- alg (Var ix) = 'x' : show ix- alg (Param ix) = 't' : show ix- alg (Const c) = show c- alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"]- alg (Uni f t) = concat [show f, "(", t, ")"]+showExpr = cata alg . removeProtection+ where alg = \case+ Var ix -> 'x' : show ix+ Param ix -> 't' : show ix+ Const c -> show c+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [show f, "(", t, ")"] +-- | convert a tree into a string in math notation+-- given named vars.+--+-- >>> showExprWithVar ["mu", "eps"] $ "x0" + sin ( "x1" * tanh ("t0" + 2) )+-- "(mu + Sin(Tanh(eps * (t0 + 2.0))))"+showExprWithVars :: [String] -> Fix SRTree -> String+showExprWithVars varnames = cata alg . removeProtection+ where alg = \case+ Var ix -> varnames !! ix+ Param ix -> 't' : show ix+ Const c -> show c+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [show f, "(", t, ")"]++-- | prints the expression printExpr :: Fix SRTree -> IO () printExpr = putStrLn . showExpr +-- | prints the expression+printExprWithVars :: [String] -> Fix SRTree -> IO ()+printExprWithVars varnames = putStrLn . showExprWithVars varnames++-- how to display an operator +showOp :: Op -> String showOp Add = "+" showOp Sub = "-" showOp Mul = "*" showOp Div = "/" showOp Power = "^"+showOp AQ = "aq"+showOp PowerAbs = "|^|" {-# INLINE showOp #-} -- | Displays a tree as a numpy compatible expression.+--+-- >>> showPython $ "x0" + sin ( tanh ("t0" + 2) )+-- "(x[:, 0] + np.sin(np.tanh((t[:, 0] + 2.0))))" showPython :: Fix SRTree -> String-showPython = cata alg+showPython = cata alg . removeProtection where- alg (Var ix) = concat ["x[:, ", show ix, "]"]- alg (Param ix) = concat ["t[:, ", show ix, "]"]- alg (Const c) = show c- alg (Bin Power l r) = concat [l, " ** ", r]- alg (Bin op l r) = concat ["(", l, " ", showOp op, " ", r, ")"]- alg (Uni f t) = concat [pyFun f, "(", t, ")"]+ alg = \case+ Var ix -> concat ["x[:, ", show ix, "]"]+ Param ix -> concat ["t[", show ix, "]"]+ Const c -> show c+ Bin Power l r -> concat [l, " ** ", r]+ Bin op l r -> concat ["(", l, " ", showOp op, " ", r, ")"]+ Uni f t -> concat [pyFun f, "(", t, ")"] + pyFun Id = "" pyFun Abs = "np.abs" pyFun Sin = "np.sin"@@ -76,39 +129,69 @@ pyFun Square = "np.square" pyFun Log = "np.log" pyFun Exp = "np.exp"+ pyFun Cbrt = "np.cbrt"+ pyFun Recip = "np.reciprocal" +-- | print the expression in numpy notation printPython :: Fix SRTree -> IO () printPython = putStrLn . showPython --- | Displays a tree as a sympy compatible expression.+-- | Displays a tree as a LaTeX compatible expression.+--+-- >>> showLatex $ "x0" + sin ( tanh ("t0" + 2) )+-- "\\left(x_{, 0} + \\operatorname{sin}(\\operatorname{tanh}(\\left(\\theta_{, 0} + 2.0\\right)))\\right)" showLatex :: Fix SRTree -> String-showLatex = cata alg+showLatex = cata alg . removeProtection where- alg (Var ix) = concat ["x_{, ", show ix, "}"]- alg (Param ix) = concat ["\\theta_{, ", show ix, "}"]- alg (Const c) = show c- alg (Bin Power l r) = concat [l, "^{", r, "}"]- alg (Bin op l r) = concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]- alg (Uni Abs t) = concat ["\\left |", t, "\\right |"]- alg (Uni f t) = concat [showLatexFun f, "(", t, ")"]-+ alg = \case+ Var ix -> concat ["x_{", show ix, "}"]+ Param ix -> concat ["\\theta_{", show ix, "}"]+ Const c -> show c+ Bin Power l r -> concat ["{", l, "^{", r, "}}"]+ Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"]+ Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"]+ Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"]+ Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]+ Uni Abs t -> concat ["\\left |", t, "\\right |"]+ Uni Recip t -> concat ["\\frac{1}{", t, "}"]+ Uni f t -> concat [showLatexFun f, "(", t, ")"]+ +showLatexWithVars :: [String] -> Fix SRTree -> String+showLatexWithVars varnames = cata alg . removeProtection+ where + alg = \case+ Var ix -> concat ["\\operatorname{", varnames !! ix, "}"]+ Param ix -> concat ["\\theta_{", show ix, "}"]+ Const c -> show c+ Bin Power l r -> concat ["{", l, "^{", r, "}}"]+ Bin PowerAbs l r -> concat ["{\\left|", l, "\\right|^{", r, "}}"]+ Bin Mul l r -> concat ["\\left(", l, " \\cdot ", r, "\\right)"]+ Bin Div l r -> concat ["\\frac{", l, "}{", r, "}"]+ Bin op l r -> concat ["\\left(", l, " ", showOp op, " ", r, "\\right)"]+ Uni Abs t -> concat ["\\left |", t, "\\right |"]+ Uni Recip t -> concat ["\\frac{1}{", t, "}"]+ Uni f t -> concat [showLatexFun f, "(", t, ")"]+ showLatexFun :: Function -> String showLatexFun f = mconcat ["\\operatorname{", map toLower $ show f, "}"] {-# INLINE showLatexFun #-} +-- | prints expression in LaTeX notation. printLatex :: Fix SRTree -> IO () printLatex = putStrLn . showLatex -- | Displays a tree in Tikz format showTikz :: Fix SRTree -> String-showTikz = cata alg+showTikz = cata alg . removeProtection where+ alg = \case+ Var ix -> concat ["[$x_{, ", show ix, "}$]\n"]+ Param ix -> concat ["[$\\theta_{, ", show ix, "}$]\n"]+ Const c -> concat ["[$", show (roundN 2 c), "$]\n"]+ Bin op l r -> concat ["[", showOpTikz op, l, r, "]\n"]+ Uni f t -> concat ["[", map toLower $ show f, t, "]\n"]+ roundN n x = let ten = 10^n in (/ ten) . fromIntegral . round $ x*ten- alg (Var ix) = concat ["[$x_{, ", show ix, "}$]\n"]- alg (Param ix) = concat ["[$\\theta_{, ", show ix, "}$]\n"]- alg (Const c) = concat ["[$", show (roundN 2 c), "$]\n"]- alg (Bin op l r) = concat ["[", showOpTikz op, l, r, "]\n"]- alg (Uni f t) = concat ["[", map toLower $ show f, t, "]\n"] showOpTikz Add = "+\n" showOpTikz Sub = "-\n"@@ -116,4 +199,6 @@ showOpTikz Div = "÷\n" showOpTikz Power = "\\^{}\n" +-- | prints the tree in TikZ format +printTikz :: Fix SRTree -> IO () printTikz = putStrLn . showTikz
src/Data/SRTree/Random.hs view
@@ -2,7 +2,7 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Random --- Copyright : (c) Fabricio Olivetti 2021 - 2021+-- Copyright : (c) Fabricio Olivetti 2021 - 2024 -- License : BSD3 -- Maintainer : fabricio.olivetti@gmail.com -- Stability : experimental@@ -18,25 +18,35 @@ , HasEverything , FullParams(..) , RndTree+ , Rng(..) , randomVar , randomConst , randomPow , randomFunction , randomNode , randomNonTerminal+ , randomRange+ , randomTreeTemplate , randomTree , randomTreeBalanced+ , toss+ , tossBiased+ , randomVal+ , randomVec+ , randomFrom ) where -import System.Random -import Control.Monad.State -import Control.Monad.Reader +import Control.Monad.Reader (ReaderT, asks, runReaderT)+import Control.Monad.State.Strict ( MonadState(state), MonadTrans(lift), StateT ) import Data.Maybe (fromJust)- import Data.SRTree.Internal-import Data.SRTree.Recursion+import System.Random (Random (random, randomR), StdGen, mkStdGen)+import Data.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 #-}
src/Data/SRTree/Recursion.hs view
@@ -1,5 +1,17 @@ {-# language RankNTypes #-} {-# language DeriveFunctor #-}+-----------------------------------------------------------------------------+-- |+-- Module : Data.SRTree.Recursion +-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : FlexibleInstances, DeriveFunctor, ScopedTypeVariables+--+-- Recursion schemes+--+----------------------------------------------------------------------------- module Data.SRTree.Recursion where import Control.Monad ( (>=>) )
+ src/Numeric/Optimization/NLOPT.hs view
@@ -0,0 +1,976 @@+{-# OPTIONS_GHC -Wall #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TypeApplications #-}++{- |+Module : Numeric.NLOPT+Copyright : (c) Matthew Peddie 2017+License : BSD3+Maintainer : Matthew Peddie <mpeddie@gmail.com>+Stability : provisional+Portability : GHC++This module provides a high-level, @hmatrix@-compatible interface to+the <http://ab-initio.mit.edu/wiki/index.php/NLopt NLOPT> library by+Steven G. Johnson.++NOTE: This is an adaptation from https://hackage.haskell.org/package/hmatrix-nlopt-0.2.0.0+that removes the dependency to hmatrix and support any Vector Storage.++= Documentation++Most non-numerical details are documented, but for specific+information on what the optimization methods do, how constraints are+handled, etc., you should consult:++ * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Introduction NLOPT introduction>++ * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Reference NLOPT reference manual>++ * The <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms NLOPT algorithm manual>++= Example program++The following interactive session example uses the Nelder-Mead simplex+algorithm, a derivative-free local optimizer, to minimize a trivial+function with a minimum of 22.0 at @(0, 0)@.++>>> import Numeric.LinearAlgebra ( dot, fromList )+>>> let objf x = x `dot` x + 22 -- define objective+>>> let stop = ObjectiveRelativeTolerance 1e-6 :| [] -- define stopping criterion+>>> let algorithm = NELDERMEAD objf [] Nothing -- specify algorithm+>>> let problem = LocalProblem 2 stop algorithm -- specify problem+>>> let x0 = fromList [5, 10] -- specify initial guess+>>> minimizeLocal problem x0+Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED})++-}++module 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+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- for more details on how the methods work and how they relate to one another.+--+-- Optional parameters are wrapped in a 'Maybe'; for example, if you+-- see 'Maybe' 'Population', you can simply specify 'Nothing' to use+-- the default behavior.+data GlobalAlgorithm+ -- | DIviding RECTangles+ = DIRECT Objective+ -- | DIviding RECTangles, locally-biased variant+ | DIRECT_L Objective+ -- | DIviding RECTangles, "slightly randomized"+ | DIRECT_L_RAND Objective RandomSeed+ -- | DIviding RECTangles, unscaled version+ | DIRECT_NOSCAL Objective+ -- | DIviding RECTangles, locally-biased and unscaled+ | DIRECT_L_NOSCAL Objective+ -- | DIviding RECTangles, locally-biased, unscaled and "slightly+ -- randomized"+ | DIRECT_L_RAND_NOSCAL Objective RandomSeed+ -- | DIviding RECTangles, original FORTRAN implementation+ | ORIG_DIRECT Objective InequalityConstraints+ -- | DIviding RECTangles, locally-biased, original FORTRAN+ -- implementation+ | ORIG_DIRECT_L Objective InequalityConstraints+ -- | Stochastic Global Optimization.+ -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__+ | STOGO ObjectiveD+ -- | Stochastic Global Optimization, randomized variant.+ -- __This algorithm is only available if you have linked with @libnlopt_cxx@.__+ | STOGO_RAND ObjectiveD RandomSeed+ -- | Controlled Random Search with Local Mutation+ | CRS2_LM Objective RandomSeed (Maybe Population)+ -- | Improved Stochastic Ranking Evolution Strategy+ | ISRES Objective InequalityConstraints EqualityConstraints RandomSeed (Maybe Population)+ -- | Evolutionary Algorithm+ | ESCH Objective+ -- | Original Multi-Level Single-Linkage+ | MLSL Objective LocalProblem (Maybe Population)+ -- | Multi-Level Single-Linkage with Sobol Low-Discrepancy+ -- Sequence for starting points+ | MLSL_LDS Objective LocalProblem (Maybe Population)++algorithmEnumOfGlobal :: GlobalAlgorithm -> N.Algorithm+algorithmEnumOfGlobal (DIRECT _) = N.GN_DIRECT+algorithmEnumOfGlobal (DIRECT_L _) = N.GN_DIRECT_L+algorithmEnumOfGlobal (DIRECT_L_RAND _ _) = N.GN_DIRECT_L_RAND+algorithmEnumOfGlobal (DIRECT_NOSCAL _) = N.GN_DIRECT_NOSCAL+algorithmEnumOfGlobal (DIRECT_L_NOSCAL _) = N.GN_DIRECT_L_NOSCAL+algorithmEnumOfGlobal (DIRECT_L_RAND_NOSCAL _ _) = N.GN_DIRECT_L_RAND_NOSCAL+algorithmEnumOfGlobal (ORIG_DIRECT _ _) = N.GN_ORIG_DIRECT+algorithmEnumOfGlobal (ORIG_DIRECT_L _ _) = N.GN_ORIG_DIRECT_L+algorithmEnumOfGlobal (STOGO _) = N.GD_STOGO+algorithmEnumOfGlobal (STOGO_RAND _ _) = N.GD_STOGO_RAND+algorithmEnumOfGlobal (CRS2_LM _ _ _) = N.GN_CRS2_LM+algorithmEnumOfGlobal (ISRES _ _ _ _ _) = N.GN_ISRES+algorithmEnumOfGlobal (ESCH _) = N.GN_ESCH+algorithmEnumOfGlobal (MLSL _ _ _) = N.G_MLSL+algorithmEnumOfGlobal (MLSL_LDS _ _ _) = N.G_MLSL_LDS++applyGlobalObjective :: N.Opt -> GlobalAlgorithm -> IO ()+applyGlobalObjective opt alg = go alg+ where+ obj = tryTo . applyObjective opt . MinimumObjective+ objD = tryTo . applyObjectiveD opt . MinimumObjective++ go (DIRECT o) = obj o+ go (DIRECT_L o) = obj o+ go (DIRECT_NOSCAL o) = obj o+ go (DIRECT_L_NOSCAL o) = obj o+ go (ESCH o) = obj o+ go (STOGO o) = objD o+ go (DIRECT_L_RAND o _) = obj o+ go (DIRECT_L_RAND_NOSCAL o _) = obj o+ go (ORIG_DIRECT o _) = obj o+ go (ORIG_DIRECT_L o _) = obj o+ go (STOGO_RAND o _) = objD o+ go (CRS2_LM o _ _) = obj o+ go (ISRES o _ _ _ _) = obj o+ go (MLSL o _ _) = obj o+ go (MLSL_LDS o _ _) = obj o++applyGlobalAlgorithm :: N.Opt -> GlobalAlgorithm -> IO ()+applyGlobalAlgorithm opt alg = do+ applyGlobalObjective opt alg+ go alg+ where+ seed = applyRandomSeed+ pop = maybe (return ()) (tryTo . applyPopulation opt)+ ic = traverse_ (tryTo . applyConstraint opt)+ ec = traverse_ (tryTo . applyConstraint opt)++ local lp = setupLocalProblem lp >>= N.set_local_optimizer opt++ go (DIRECT_L_RAND _ s) = seed s+ go (DIRECT_L_RAND_NOSCAL _ s) = seed s+ go (ORIG_DIRECT _ ineq) = ic ineq+ go (ORIG_DIRECT_L _ ineq) = ic ineq+ go (STOGO_RAND _ s) = seed s+ go (CRS2_LM _ s p) = seed s *> pop p+ go (ISRES _ ineq eq s p) = ic ineq *> ec eq *> seed s *> pop p+ go (MLSL _ lp p) = local lp *> pop p+ go (MLSL_LDS _ lp p) = local lp *> pop p+ go _ = return ()++tryTo :: IO N.Result -> IO ()+tryTo act = do+ result <- act+ if (N.isSuccess result)+ then return ()+ else Ex.throw $ NloptException result++data NloptException = NloptException N.Result deriving (Show, Typeable)+instance Exception NloptException++-- | Solve the specified global optimization problem.+--+-- = Example program+--+-- The following interactive session example uses the 'ISRES'+-- algorithm, a stochastic, derivative-free global optimizer, to+-- minimize a trivial function with a minimum of 22.0 at @(0, 0)@.+-- The search is conducted within a box from -10 to 10 in each+-- dimension.+--+-- >>> import Numeric.LinearAlgebra ( dot, fromList )+-- >>> let objf x = x `dot` x + 22 -- define objective+-- >>> let stop = ObjectiveRelativeTolerance 1e-12 :| [] -- define stopping criterion+-- >>> let algorithm = ISRES objf [] [] (SeedValue 22) Nothing -- specify algorithm+-- >>> let lowerbounds = fromList [-10, -10] -- specify bounds+-- >>> let upperbounds = fromList [10, 10] -- specify bounds+-- >>> let problem = GlobalProblem lowerbounds upperbounds stop algorithm+-- >>> let x0 = fromList [5, 8] -- specify initial guess+-- >>> minimizeGlobal problem x0+-- Right (Solution {solutionCost = 22.000000000002807, solutionParams = [-1.660591102367038e-6,2.2407062393213684e-7], solutionResult = FTOL_REACHED})+minimizeGlobal :: GlobalProblem -- ^ Problem specification+ -> Vector Double -- ^ Initial parameter guess+ -> Either N.Result Solution -- ^ Optimization results+minimizeGlobal prob x0 =+ unsafePerformIO $ (Right <$> minimizeGlobal' prob x0) `Ex.catch` handler+ where+ handler :: NloptException -> IO (Either N.Result a)+ handler (NloptException retcode) = return $ Left retcode++applyGlobalProblem :: N.Opt -> GlobalProblem -> IO ()+applyGlobalProblem opt (GlobalProblem lb ub stop alg) = do+ tryTo $ applyBounds opt (LowerBounds lb)+ tryTo $ applyBounds opt (UpperBounds ub)+ traverse_ (tryTo . applyStoppingCondition opt) stop+ applyGlobalAlgorithm opt alg++newOpt :: N.Algorithm -> Word -> IO N.Opt+newOpt alg sz = do+ opt' <- N.create alg sz+ case opt' of+ Nothing -> Ex.throw $ NloptException N.FAILURE+ Just opt -> return opt++setupGlobalProblem :: GlobalProblem -> IO N.Opt+setupGlobalProblem gp@(GlobalProblem _ _ _ alg) = do+ opt <- newOpt (algorithmEnumOfGlobal alg) (problemSize gp)+ applyGlobalProblem opt gp+ return opt++solveProblem :: N.Opt -> Vector Double -> IO Solution+solveProblem opt x0 = do+ (N.Output outret outcost outx nevals) <- N.optimize opt x0+ if (N.isSuccess outret)+ then return $ Solution outcost outx outret nevals+ else Ex.throw $ NloptException outret++minimizeGlobal' :: GlobalProblem -> Vector Double -> IO Solution+minimizeGlobal' gp x0 = do+ opt <- setupGlobalProblem gp+ solveProblem opt x0++data LocalProblem = LocalProblem+ { lsize :: Word -- ^ The dimension of the+ -- parameter vector.+ , lstop :: NonEmpty StoppingCondition -- ^ At least one stopping+ -- condition+ , lalgorithm :: LocalAlgorithm -- ^ Algorithm specification+ }++-- | These are the local minimization algorithms provided by NLOPT. Please see+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- for more details on how the methods work and how they relate to one+-- another. Note that some local methods require you provide+-- derivatives (gradients or Jacobians) for your objective function+-- and constraint functions.+--+-- Optional parameters are wrapped in a 'Maybe'; for example, if you+-- see 'Maybe' 'VectorStorage', you can simply specify 'Nothing' to+-- use the default behavior.+data LocalAlgorithm+ -- | Limited-memory BFGS+ = LBFGS_NOCEDAL ObjectiveD (Maybe VectorStorage)+ -- | Limited-memory BFGS+ | LBFGS ObjectiveD (Maybe VectorStorage)+ -- | Shifted limited-memory variable-metric, rank-2+ | VAR2 ObjectiveD (Maybe VectorStorage)+ -- | Shifted limited-memory variable-metric, rank-1+ | VAR1 ObjectiveD (Maybe VectorStorage)+ -- | Truncated Newton's method+ | TNEWTON ObjectiveD (Maybe VectorStorage)+ -- | Truncated Newton's method with automatic restarting+ | TNEWTON_RESTART ObjectiveD (Maybe VectorStorage)+ -- | Preconditioned truncated Newton's method+ | TNEWTON_PRECOND ObjectiveD (Maybe VectorStorage)+ -- | Preconditioned truncated Newton's method with automatic+ -- restarting+ | TNEWTON_PRECOND_RESTART ObjectiveD (Maybe VectorStorage)+ -- | Method of moving averages+ | MMA ObjectiveD InequalityConstraintsD+ -- | Sequential Least-Squares Quadratic Programming+ | SLSQP ObjectiveD [Bounds] InequalityConstraintsD EqualityConstraintsD+ -- | Conservative Convex Separable Approximation+ | CCSAQ ObjectiveD Preconditioner+ -- | PRincipal AXIS gradient-free local optimization+ | PRAXIS Objective [Bounds] (Maybe InitialStep)+ -- | Constrained Optimization BY Linear Approximations+ | COBYLA Objective [Bounds] InequalityConstraints EqualityConstraints+ (Maybe InitialStep)+ -- | Powell's NEWUOA algorithm+ | NEWUOA Objective (Maybe InitialStep)+ -- | Powell's NEWUOA algorithm with bounds by SGJ+ | NEWUOA_BOUND Objective [Bounds] (Maybe InitialStep)+ -- | Nelder-Mead Simplex gradient-free method+ | NELDERMEAD Objective [Bounds] (Maybe InitialStep)+ -- | NLOPT implementation of Rowan's Subplex algorithm+ | SBPLX Objective [Bounds] (Maybe InitialStep)+ -- | Bounded Optimization BY Quadratic Approximations+ | BOBYQA Objective [Bounds] (Maybe InitialStep)++algorithmEnumOfLocal :: LocalAlgorithm -> N.Algorithm+algorithmEnumOfLocal (LBFGS_NOCEDAL _ _) = N.LD_LBFGS_NOCEDAL+algorithmEnumOfLocal (LBFGS _ _) = N.LD_LBFGS+algorithmEnumOfLocal (VAR2 _ _) = N.LD_VAR2+algorithmEnumOfLocal (VAR1 _ _) = N.LD_VAR1+algorithmEnumOfLocal (TNEWTON _ _) = N.LD_TNEWTON+algorithmEnumOfLocal (TNEWTON_RESTART _ _) = N.LD_TNEWTON_RESTART+algorithmEnumOfLocal (TNEWTON_PRECOND _ _) = N.LD_TNEWTON_PRECOND+algorithmEnumOfLocal (TNEWTON_PRECOND_RESTART _ _) = N.LD_TNEWTON_PRECOND_RESTART+algorithmEnumOfLocal (MMA _ _) = N.LD_MMA+algorithmEnumOfLocal (SLSQP _ _ _ _) = N.LD_SLSQP+algorithmEnumOfLocal (CCSAQ _ _) = N.LD_CCSAQ+algorithmEnumOfLocal (PRAXIS _ _ _) = N.LN_PRAXIS+algorithmEnumOfLocal (COBYLA _ _ _ _ _) = N.LN_COBYLA+algorithmEnumOfLocal (NEWUOA _ _) = N.LN_NEWUOA+algorithmEnumOfLocal (NEWUOA_BOUND _ _ _) = N.LN_NEWUOA+algorithmEnumOfLocal (NELDERMEAD _ _ _) = N.LN_NELDERMEAD+algorithmEnumOfLocal (SBPLX _ _ _) = N.LN_SBPLX+algorithmEnumOfLocal (BOBYQA _ _ _) = N.LN_BOBYQA++applyLocalObjective :: N.Opt -> LocalAlgorithm -> IO ()+applyLocalObjective opt alg = go alg+ where+ obj = tryTo . applyObjective opt . MinimumObjective+ objD = tryTo . applyObjectiveD opt . MinimumObjective+ precond p = tryTo . applyObjectiveD opt . PreconditionedMinimumObjective p++ go (LBFGS_NOCEDAL o _) = objD o+ go (LBFGS o _) = objD o+ go (VAR2 o _) = objD o+ go (VAR1 o _) = objD o+ go (TNEWTON o _) = objD o+ go (TNEWTON_RESTART o _) = objD o+ go (TNEWTON_PRECOND o _) = objD o+ go (TNEWTON_PRECOND_RESTART o _) = objD o+ go (MMA o _) = objD o+ go (SLSQP o _ _ _) = objD o+ go (CCSAQ o prec) = precond prec o+ go (PRAXIS o _ _) = obj o+ go (COBYLA o _ _ _ _) = obj o+ go (NEWUOA o _) = obj o+ go (NEWUOA_BOUND o _ _) = obj o+ go (NELDERMEAD o _ _) = obj o+ go (SBPLX o _ _) = obj o+ go (BOBYQA o _ _) = obj o++applyLocalAlgorithm :: N.Opt -> LocalAlgorithm -> IO ()+applyLocalAlgorithm opt alg = do+ applyLocalObjective opt alg+ go alg+ where+ ic = traverse_ (tryTo . applyConstraint opt)+ icd = traverse_ (tryTo . applyConstraint opt)+ ec = traverse_ (tryTo . applyConstraint opt)+ ecd = traverse_ (tryTo . applyConstraint opt)+ store = maybe (return ()) (tryTo . applyVectorStorage opt)+ bound = traverse_ (tryTo . applyBounds opt)+ step0 = maybe (return ()) (tryTo . applyInitialStep opt)++ go (LBFGS_NOCEDAL _ vs) = store vs+ go (LBFGS _ vs) = store vs+ go (VAR2 _ vs) = store vs+ go (VAR1 _ vs) = store vs+ go (TNEWTON _ vs) = store vs+ go (TNEWTON_RESTART _ vs) = store vs+ go (TNEWTON_PRECOND _ vs) = store vs+ go (TNEWTON_PRECOND_RESTART _ vs) = store vs+ go (MMA _ ineqd) = icd ineqd+ go (SLSQP _ b ineqd eqd) =+ bound b *> icd ineqd *> ecd eqd+ go (CCSAQ _ _ ) = return ()+ go (PRAXIS _ b s) = bound b *> step0 s+ go (COBYLA _ b ineq eq s) =+ bound b *> ic ineq *> ec eq *> step0 s+ go (NEWUOA _ s) = step0 s+ go (NEWUOA_BOUND _ b s) = bound b *> step0 s+ go (NELDERMEAD _ b s) = bound b *> step0 s+ go (SBPLX _ b s) = bound b *> step0 s+ go (BOBYQA _ b s) = bound b *> step0 s++applyLocalProblem :: N.Opt -> LocalProblem -> IO ()+applyLocalProblem opt (LocalProblem _ stop alg) = do+ traverse_ (tryTo . applyStoppingCondition opt) stop+ applyLocalAlgorithm opt alg++setupLocalProblem :: LocalProblem -> IO N.Opt+setupLocalProblem lp@(LocalProblem sz _ alg) = do+ opt <- newOpt (algorithmEnumOfLocal alg) sz+ applyLocalProblem opt lp+ return opt++minimizeLocal' :: LocalProblem -> Vector Double -> IO Solution+minimizeLocal' lp x0 = do+ opt <- setupLocalProblem lp+ solveProblem opt x0++-- |+-- == Example program+--+-- The following interactive session example enforces the same scalar+-- constraint as the nonlinear constraint example, but this time it+-- uses the SLSQP solver to find the minimum.+--+-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList, scale )+-- >>> let objf x = (x `dot` x + 22, 2 `scale` x)+-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []+-- >>> let constraintf x = (sum (toList x) - 1.0, fromList [1, 1])+-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6+-- >>> let algorithm = SLSQP objf [] [] [constraint]+-- >>> let problem = LocalProblem 2 stop algorithm+-- >>> let x0 = fromList [5, 10]+-- >>> minimizeLocal problem x0+-- Right (Solution {solutionCost = 22.5, solutionParams = [0.4999999999999998,0.5000000000000002], solutionResult = FTOL_REACHED})+minimizeLocal :: LocalProblem -> Vector Double -> Either N.Result Solution+minimizeLocal prob x0 =+ unsafePerformIO $ (Right <$> minimizeLocal' prob x0) `Ex.catch` handler+ where+ handler :: NloptException -> IO (Either N.Result a)+ handler (NloptException retcode) = return $ Left retcode++class ProblemSize c where+ problemSize :: c -> Word++instance ProblemSize LocalProblem where+ problemSize = lsize++instance ProblemSize GlobalProblem where+ problemSize = fromIntegral . V.length . lowerBounds++instance ProblemSize AugLagProblem where+ problemSize (AugLagProblem _ _ alg) = case alg of+ AUGLAG_LOCAL lp _ _ -> problemSize lp+ AUGLAG_EQ_LOCAL lp -> problemSize lp+ AUGLAG_GLOBAL gp _ _ -> problemSize gp+ AUGLAG_EQ_GLOBAL gp -> problemSize gp+++-- | __IMPORTANT NOTE__+--+-- For augmented lagrangian problems, you, the user, are responsible+-- for providing the appropriate type of constraint. If the+-- subsidiary problem requires an `ObjectiveD`, then you should+-- provide constraint functions with derivatives. If the subsidiary+-- problem requires an `Objective`, you should provide constraint+-- functions without derivatives. If you don't do this, you may get a+-- runtime error.+data AugLagProblem = AugLagProblem+ { alEquality :: EqualityConstraints -- ^ Possibly empty set of+ -- equality constraints+ , alEqualityD :: EqualityConstraintsD -- ^ Possibly empty set of+ -- equality constraints with+ -- derivatives+ , alalgorithm :: AugLagAlgorithm -- ^ Algorithm specification.+ }++-- | The Augmented Lagrangian solvers allow you to enforce nonlinear+-- constraints while using local or global algorithms that don't+-- natively support them. The subsidiary problem is used to do the+-- minimization, but the @AUGLAG@ methods modify the objective to+-- enforce the constraints. Please see+-- <http://ab-initio.mit.edu/wiki/index.php/NLopt_Algorithms the NLOPT algorithm manual>+-- for more details on how the methods work and how they relate to one another.+--+-- See the documentation for 'AugLagProblem' for an important note+-- about the constraint functions.+data AugLagAlgorithm+ -- | AUGmented LAGrangian with a local subsidiary method+ = AUGLAG_LOCAL LocalProblem InequalityConstraints InequalityConstraintsD+ -- | AUGmented LAGrangian with a local subsidiary method and with+ -- penalty functions only for equality constraints+ | AUGLAG_EQ_LOCAL LocalProblem+ -- | AUGmented LAGrangian with a global subsidiary method+ | AUGLAG_GLOBAL GlobalProblem InequalityConstraints InequalityConstraintsD+ -- | AUGmented LAGrangian with a global subsidiary method and with+ -- penalty functions only for equality constraints.+ | AUGLAG_EQ_GLOBAL GlobalProblem++algorithmEnumOfAugLag :: AugLagAlgorithm -> N.Algorithm+algorithmEnumOfAugLag (AUGLAG_LOCAL _ _ _) = N.AUGLAG+algorithmEnumOfAugLag (AUGLAG_EQ_LOCAL _) = N.AUGLAG_EQ+algorithmEnumOfAugLag (AUGLAG_GLOBAL _ _ _) = N.AUGLAG+algorithmEnumOfAugLag (AUGLAG_EQ_GLOBAL _) = N.AUGLAG_EQ++-- | This structure is returned in the event of a successful+-- optimization.+data Solution = Solution+ { solutionCost :: Double -- ^ The objective function value+ -- at the minimum+ , solutionParams :: Vector Double -- ^ The parameter vector which+ -- minimizes the objective+ , solutionResult :: N.Result -- ^ Why the optimizer stopped++ , nEvals :: Int -- ^ Number of evaluations until stop+ } deriving (Eq, Show, Read)++applyAugLagAlgorithm :: N.Opt -> AugLagAlgorithm -> IO ()+applyAugLagAlgorithm opt alg = go alg+ where+ ic = traverse_ (tryTo . applyConstraint opt)+ icd = traverse_ (tryTo . applyConstraint opt)+ -- AUGLAG won't work at all if you don't pass it the same+ -- objective as the subproblem -- here we pull out the subproblem+ -- objectives from the algorithm spec and set the same objective+ -- function so the user can't mess it up.+ local lp = tryTo $ do+ localopt <- setupLocalProblem lp+ applyLocalObjective opt (lalgorithm lp)+ N.set_local_optimizer opt localopt+ global gp = do+ tryTo $ setupGlobalProblem gp >>= N.set_local_optimizer opt+ applyGlobalObjective opt (galgorithm gp)++ go (AUGLAG_LOCAL lp ineq ineqd) = local lp *> ic ineq *> icd ineqd+ go (AUGLAG_EQ_LOCAL lp) = local lp+ go (AUGLAG_GLOBAL gp ineq ineqd) = global gp *> ic ineq *> icd ineqd+ go (AUGLAG_EQ_GLOBAL gp) = global gp++applyAugLagProblem :: N.Opt -> AugLagProblem -> IO ()+applyAugLagProblem opt (AugLagProblem eq eqd alg) = do+ traverse_ (tryTo . applyConstraint opt) eq+ traverse_ (tryTo . applyConstraint opt) eqd+ applyAugLagAlgorithm opt alg++minimizeAugLag' :: AugLagProblem -> Vector Double -> IO Solution+minimizeAugLag' ap@(AugLagProblem _ _ alg) x0 = do+ opt <- newOpt (algorithmEnumOfAugLag alg) (problemSize ap)+ applyAugLagProblem opt ap+ solveProblem opt x0++-- |+-- == Example program+--+-- The following interactive session example enforces the same scalar+-- constraint as the nonlinear constraint example, but this time it+-- uses the augmented Lagrangian method to enforce the constraint and+-- the 'SBPLX' algorithm, which does not support nonlinear constraints+-- itself, to perform the minimization. As before, the parameters+-- must always sum to 1, and the minimizer finds the same constrained+-- minimum of 22.5 at @(0.5, 0.5)@.+--+-- >>> import Numeric.LinearAlgebra ( dot, fromList, toList )+-- >>> let objf x = x `dot` x + 22+-- >>> let stop = ObjectiveRelativeTolerance 1e-9 :| []+-- >>> let algorithm = SBPLX objf [] Nothing+-- >>> let subproblem = LocalProblem 2 stop algorithm+-- >>> let x0 = fromList [5, 10]+-- >>> minimizeLocal subproblem x0+-- Right (Solution {solutionCost = 22.0, solutionParams = [0.0,0.0], solutionResult = FTOL_REACHED})+-- >>> -- define constraint function:+-- >>> let constraintf x = sum (toList x) - 1.0+-- >>> -- define constraint object to pass to the algorithm:+-- >>> let constraint = EqualityConstraint (Scalar constraintf) 1e-6+-- >>> let problem = AugLagProblem [constraint] [] (AUGLAG_EQ_LOCAL subproblem)+-- >>> minimizeAugLag problem x0+-- Right (Solution {solutionCost = 22.500000015505844, solutionParams = [0.5000880506776678,0.4999119493223323], solutionResult = FTOL_REACHED})++minimizeAugLag :: AugLagProblem -> Vector Double -> Either N.Result Solution+minimizeAugLag prob x0 =+ unsafePerformIO $ (Right <$> minimizeAugLag' prob x0) `Ex.catch` handler+ where+ handler :: NloptException -> IO (Either N.Result a)+ handler (NloptException retcode) = return $ Left retcode
+ src/Numeric/Optimization/NLOPT/Bindings.hs view
@@ -0,0 +1,1075 @@+{-# OPTIONS_GHC -Wall #-}+{-# LANGUAGE ForeignFunctionInterface #-}+{-# LANGUAGE NoMonomorphismRestriction #-}++{- |+Module : Numeric.Optimization.NLOPT.Bindings+Copyright : (c) Matthew Peddie 2017+License : BSD3+Maintainer : Matthew Peddie <mpeddie@gmail.com>+Stability : provisional+Portability : GHC++Low-level interface to the NLOPT library. Please see+<http://ab-initio.mit.edu/wiki/index.php/NLopt_Reference the NLOPT reference manual>+for detailed information; the Haskell functions in this module closely+follow the interface to the C library in @nlopt.h@.++Differences between this module and the C interface are documented+here; functions with identical interfaces are not. In general:++ ['Opt'] corresponds to an @nlopt_opt@ object++ ['Result'] corresponds to @nlopt_result@++ ['V.Vector' 'Double'] corresponds to a @const double *@ input or a+ @double *@ output++ ['ScalarFunction'] corresponds to @nlopt_func@++ ['VectorFunction'] corresponds to @nlopt_mfunc@++ ['PreconditionerFunction'] corresponds to @nlopt_precond@++User data that is handled by @void *@ in the C bindings can be any+Haskell value.++-}++module Numeric.Optimization.NLOPT.Bindings (+ -- * C enums+ Algorithm(..)+ , algorithm_name+ , Result(..)+ , isSuccess+ -- * Optimizer object+ , Opt+ , create+ , destroy+ , copy+ -- * Random number generator seeding+ , srand+ , srand_time+ -- * Metadata+ , Version(..)+ , version+ , get_algorithm+ , get_dimension+ -- * Callbacks+ , ScalarFunction+ , VectorFunction+ , PreconditionerFunction+ -- * Running the optimizer+ , Output(..)+ , optimize+ -- * Objective function configuration+ , set_min_objective+ , set_max_objective+ , set_precond_min_objective+ , set_precond_max_objective+ -- * Bound configuration+ , set_lower_bounds+ , set_lower_bounds1+ , get_lower_bounds+ , set_upper_bounds+ , set_upper_bounds1+ , get_upper_bounds+ -- * Constraint configuration+ , remove_inequality_constraints+ , add_inequality_constraint+ , add_precond_inequality_constraint+ , add_inequality_mconstraint+ , remove_equality_constraints+ , add_equality_constraint+ , add_precond_equality_constraint+ , add_equality_mconstraint+ -- * Stopping criterion configuration+ , set_stopval+ , get_stopval+ , set_ftol_rel+ , get_ftol_rel+ , set_ftol_abs+ , get_ftol_abs+ , set_xtol_rel+ , get_xtol_rel+ , set_xtol_abs1+ , set_xtol_abs+ , get_xtol_abs+ , set_maxeval+ , get_maxeval+ , set_maxtime+ , get_maxtime+ , force_stop+ , set_force_stop+ , get_force_stop+ -- * Algorithm-specific configuration+ , set_local_optimizer+ , set_population+ , get_population+ , set_vector_storage+ , get_vector_storage+ , set_default_initial_step+ , set_initial_step+ , set_initial_step1+ , get_initial_step+ ) where++import Foreign hiding (void)+import Foreign.C.String+import Foreign.C.Types+import qualified Foreign.Concurrent as CFP++import qualified Data.Vector.Storable.Mutable as MV+import qualified Data.Vector.Storable as V++{- C enums -}++-- | The NLOPT algorithm names, apart from the names of the actual+-- optimization methods, follow this scheme:+--+-- [@G@] means a global method+-- [@L@] means a local method+-- [@D@] means a method that requires the derivative+-- [@N@] means a method that does not require the derivative+-- [@*_RAND@] means the algorithm involves some randomization.+-- [@*_NOSCAL@] means the algorithm is *not* scaled to a unit+-- hypercube (i.e. it is sensitive to the units of x)+data Algorithm+ = GN_DIRECT -- ^ DIviding RECTangles+ | GN_DIRECT_L -- ^ DIviding RECTangles,+ -- locally-biased variant+ | GN_DIRECT_L_RAND -- ^ DIviding RECTangles, "slightly+ -- randomized"+ | GN_DIRECT_NOSCAL -- ^ DIviding RECTangles, unscaled version+ | GN_DIRECT_L_NOSCAL -- ^ DIviding RECTangles,+ -- locally-biased and unscaled+ | GN_DIRECT_L_RAND_NOSCAL -- ^ DIviding RECTangles, locally-biased,+ -- unscaled and "slightly randomized"+ | GN_ORIG_DIRECT -- ^ DIviding RECTangles, original FORTRAN+ -- implementation+ | GN_ORIG_DIRECT_L -- ^ DIviding RECTangles,+ -- locally-biased, original FORTRAN+ -- implementation+ | GD_STOGO -- ^ Stochastic Global Optimization+ | GD_STOGO_RAND -- ^ Stochastic Global Optimization,+ -- randomized variant+ | LD_LBFGS_NOCEDAL -- ^ Limited-memory BFGS+ | LD_LBFGS -- ^ Limited-memory BFGS+ | LN_PRAXIS -- ^ PRincipal AXIS gradient-free local+ -- optimization+ | LD_VAR2 -- ^ Shifted limited-memory+ -- variable-metric, rank-2+ | LD_VAR1 -- ^ Shifted limited-memory+ -- variable-metric, rank-1+ | LD_TNEWTON -- ^ Truncated Newton's method+ | LD_TNEWTON_RESTART -- ^ Truncated Newton's method with+ -- automatic restarting+ | LD_TNEWTON_PRECOND -- ^ Preconditioned truncated Newton's+ -- method+ | LD_TNEWTON_PRECOND_RESTART -- ^ Preconditioned truncated Newton's+ -- method with automatic restarting+ | GN_CRS2_LM -- ^ Controlled Random Search with+ -- Local Mutation+ | GN_MLSL -- ^ Original Multi-Level+ -- Single-Linkage+ | GD_MLSL -- ^ Original Multi-Level+ -- Single-Linkage, user-provided+ -- derivative+ | GN_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points+ | GD_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points, user-provided+ -- derivative+ | LD_MMA -- ^ Method of moving averages+ | LN_COBYLA -- ^ Constrained Optimization BY Linear+ -- Approximations+ | LN_NEWUOA -- ^ Powell's NEWUOA algorithm+ | LN_NEWUOA_BOUND -- ^ Powell's NEWUOA algorithm with+ -- bounds by SGJ+ | LN_NELDERMEAD -- ^ Nelder-Mead Simplex gradient-free+ -- method+ | LN_SBPLX -- ^ NLOPT implementation of Rowan's+ -- Subplex algorithm+ | LN_AUGLAG -- ^ AUGmented LAGrangian+ | LD_AUGLAG -- ^ AUGmented LAGrangian,+ -- user-provided derivative+ | LN_AUGLAG_EQ -- ^ AUGmented LAGrangian with penalty+ -- functions only for equality+ -- constraints+ | LD_AUGLAG_EQ -- ^ AUGmented LAGrangian with+ -- penalty functions only for equality+ -- constraints, user-provided+ -- derivative+ | LN_BOBYQA -- ^ Bounded Optimization BY Quadratic+ -- Approximations+ | GN_ISRES -- ^ Improved Stochastic Ranking+ -- Evolution Strategy++ | AUGLAG -- ^ AUGmented LAGrangian, requires+ -- local_optimizer to be set+ | AUGLAG_EQ -- ^ AUGmented LAGrangian with penalty+ -- functions only for equality+ -- constraints, requires+ -- local_optimizer to be set+ | G_MLSL -- ^ Original Multi-Level+ -- Single-Linkage, user-provided+ -- derivative, requires local_optimizer+ -- to be set+ | G_MLSL_LDS -- ^ Multi-Level Single-Linkage with+ -- Sobol Low-Discrepancy Sequence for+ -- starting points, requires+ -- local_optimizer to be set+ | LD_SLSQP -- ^ Sequential Least-SQuares Programming+ | LD_CCSAQ -- ^ Conservative Convex Separable+ -- Approximation+ | GN_ESCH -- ^ Evolutionary Algorithm+ deriving (Eq, Show, Read, Bounded)++instance Enum Algorithm where+ fromEnum GN_DIRECT = 0+ fromEnum GN_DIRECT_L = 1+ fromEnum GN_DIRECT_L_RAND = 2+ fromEnum GN_DIRECT_NOSCAL = 3+ fromEnum GN_DIRECT_L_NOSCAL = 4+ fromEnum GN_DIRECT_L_RAND_NOSCAL = 5+ fromEnum GN_ORIG_DIRECT = 6+ fromEnum GN_ORIG_DIRECT_L = 7+ fromEnum GD_STOGO = 8+ fromEnum GD_STOGO_RAND = 9+ fromEnum LD_LBFGS_NOCEDAL = 10+ fromEnum LD_LBFGS = 11+ fromEnum LN_PRAXIS = 12+ fromEnum LD_VAR2 = 13+ fromEnum LD_VAR1 = 14+ fromEnum LD_TNEWTON = 15+ fromEnum LD_TNEWTON_RESTART = 16+ fromEnum LD_TNEWTON_PRECOND = 17+ fromEnum LD_TNEWTON_PRECOND_RESTART = 18+ fromEnum GN_CRS2_LM = 19+ fromEnum GN_MLSL = 20+ fromEnum GD_MLSL = 21+ fromEnum GN_MLSL_LDS = 22+ fromEnum GD_MLSL_LDS = 23+ fromEnum LD_MMA = 24+ fromEnum LN_COBYLA = 25+ fromEnum LN_NEWUOA = 26+ fromEnum LN_NEWUOA_BOUND = 27+ fromEnum LN_NELDERMEAD = 28+ fromEnum LN_SBPLX = 29+ fromEnum LN_AUGLAG = 30+ fromEnum LD_AUGLAG = 31+ fromEnum LN_AUGLAG_EQ = 32+ fromEnum LD_AUGLAG_EQ = 33+ fromEnum LN_BOBYQA = 34+ fromEnum GN_ISRES = 35+ fromEnum AUGLAG = 36+ fromEnum AUGLAG_EQ = 37+ fromEnum G_MLSL = 38+ fromEnum G_MLSL_LDS = 39+ fromEnum LD_SLSQP = 40+ fromEnum LD_CCSAQ = 41+ fromEnum GN_ESCH = 42+ toEnum 0 = GN_DIRECT+ toEnum 1 = GN_DIRECT_L+ toEnum 2 = GN_DIRECT_L_RAND+ toEnum 3 = GN_DIRECT_NOSCAL+ toEnum 4 = GN_DIRECT_L_NOSCAL+ toEnum 5 = GN_DIRECT_L_RAND_NOSCAL+ toEnum 6 = GN_ORIG_DIRECT+ toEnum 7 = GN_ORIG_DIRECT_L+ toEnum 8 = GD_STOGO+ toEnum 9 = GD_STOGO_RAND+ toEnum 10 = LD_LBFGS_NOCEDAL+ toEnum 11 = LD_LBFGS+ toEnum 12 = LN_PRAXIS+ toEnum 13 = LD_VAR2+ toEnum 14 = LD_VAR1+ toEnum 15 = LD_TNEWTON+ toEnum 16 = LD_TNEWTON_RESTART+ toEnum 17 = LD_TNEWTON_PRECOND+ toEnum 18 = LD_TNEWTON_PRECOND_RESTART+ toEnum 19 = GN_CRS2_LM+ toEnum 20 = GN_MLSL+ toEnum 21 = GD_MLSL+ toEnum 22 = GN_MLSL_LDS+ toEnum 23 = GD_MLSL_LDS+ toEnum 24 = LD_MMA+ toEnum 25 = LN_COBYLA+ toEnum 26 = LN_NEWUOA+ toEnum 27 = LN_NEWUOA_BOUND+ toEnum 28 = LN_NELDERMEAD+ toEnum 29 = LN_SBPLX+ toEnum 30 = LN_AUGLAG+ toEnum 31 = LD_AUGLAG+ toEnum 32 = LN_AUGLAG_EQ+ toEnum 33 = LD_AUGLAG_EQ+ toEnum 34 = LN_BOBYQA+ toEnum 35 = GN_ISRES+ toEnum 36 = AUGLAG+ toEnum 37 = AUGLAG_EQ+ toEnum 38 = G_MLSL+ toEnum 39 = G_MLSL_LDS+ toEnum 40 = LD_SLSQP+ toEnum 41 = LD_CCSAQ+ toEnum 42 = GN_ESCH+ toEnum e = error $+ "Algorithm.toEnum: invalid C value '" ++ show e ++ "' received."++foreign import ccall "nlopt.h nlopt_algorithm_name"+ nlopt_algorithm_name :: CInt -> CString++algorithm_name :: Algorithm -> IO String+algorithm_name = peekCString . nlopt_algorithm_name . fromIntegral . fromEnum++-- | Mostly self-explanatory.+data Result+ = FAILURE -- ^ Generic failure code+ | INVALID_ARGS+ | OUT_OF_MEMORY+ | ROUNDOFF_LIMITED+ | FORCED_STOP+ | SUCCESS -- ^ Generic success code+ | STOPVAL_REACHED+ | FTOL_REACHED+ | XTOL_REACHED+ | MAXEVAL_REACHED+ | MAXTIME_REACHED+ deriving (Eq, Read, Show, Bounded)++instance Enum Result where+ fromEnum FAILURE = -1+ fromEnum INVALID_ARGS = -2+ fromEnum OUT_OF_MEMORY = -3+ fromEnum ROUNDOFF_LIMITED = -4+ fromEnum FORCED_STOP = -5+ fromEnum SUCCESS = 1+ fromEnum STOPVAL_REACHED = 2+ fromEnum FTOL_REACHED = 3+ fromEnum XTOL_REACHED = 4+ fromEnum MAXEVAL_REACHED = 5+ fromEnum MAXTIME_REACHED = 6+ toEnum (-1) = FAILURE+ toEnum (-2) = INVALID_ARGS+ toEnum (-3) = OUT_OF_MEMORY+ toEnum (-4) = ROUNDOFF_LIMITED+ toEnum (-5) = FORCED_STOP+ toEnum 1 = SUCCESS+ toEnum 2 = STOPVAL_REACHED+ toEnum 3 = FTOL_REACHED+ toEnum 4 = XTOL_REACHED+ toEnum 5 = MAXEVAL_REACHED+ toEnum 6 = MAXTIME_REACHED+ toEnum e = error $+ "Result.toEnum: invalid C value '" ++ show e ++ "' received."++isSuccess :: Result -> Bool+isSuccess SUCCESS = True+isSuccess STOPVAL_REACHED = True+isSuccess FTOL_REACHED = True+isSuccess XTOL_REACHED = True+isSuccess MAXEVAL_REACHED = True+isSuccess MAXTIME_REACHED = True+isSuccess _ = False++parseEnum :: (Integral a, Enum b) => a -> b+parseEnum = toEnum . fromIntegral++{- NLOPT optimizer object -}++type NloptOpt = Ptr ()++-- | An optimizer object which must be created, configured and then+-- passed to 'optimize' to solve a problem+newtype Opt = Opt { pointerFromOpt :: ForeignPtr () }++withOpt :: Opt -> (NloptOpt -> IO a) -> IO a+withOpt (Opt p) f = do+ ret <- withForeignPtr p f+ touchForeignPtr p -- This is critical! Otherwise the GC might+ -- think it's done with everything in the middle+ -- of the problem.+ return ret++useOpt :: (NloptOpt -> IO a) -> Opt -> IO a+useOpt = flip withOpt++-- Every time we make a "wrapper" call, the runtime allocates a new+-- function pointer and won't release it until we explicitly tell it+-- to. This doesn't mesh well with NLOPT's "object-oriented" design,+-- wherein we have to allocate an object and make a bunch of setup+-- calls before we run the problem, so what we do is add a finalizer+-- to the 'Opt' object's 'ForeignPtr' every time we need to create a+-- function pointer for C to use.+addFunPtrFinalizer :: Opt -> FunPtr a -> IO ()+addFunPtrFinalizer (Opt p) funptr =+ CFP.addForeignPtrFinalizer p (freeHaskellFunPtr funptr)++foreign import ccall "nlopt.h nlopt_create"+ nlopt_create :: CInt -> CUInt -> IO (NloptOpt)++-- | Create a new 'Opt' object+create :: Algorithm -- ^ Choice of algorithm+ -> Word -- ^ Parameter vector dimension+ -> IO (Maybe Opt) -- ^ Optimizer object+create alg dimension = do+ outp <- nlopt_create (fromIntegral $ fromEnum alg) (fromIntegral dimension)+ if (outp == nullPtr)+ then return Nothing+ else Just . Opt <$> CFP.newForeignPtr outp (nlopt_destroy outp)++foreign import ccall "nlopt.h nlopt_destroy"+ nlopt_destroy :: NloptOpt -> IO ()++-- It shouldn't be strictly necessary to call this by hand since we've+-- already put a call to 'nlopt_destroy' into the 'ForeignPtr', but+-- it's available in the C interface.+destroy :: Opt -> IO ()+destroy = finalizeForeignPtr . pointerFromOpt++foreign import ccall "nlopt.h nlopt_copy"+ nlopt_copy :: NloptOpt -> IO (NloptOpt)++copy :: Opt -> IO Opt+copy = useOpt $ \inp -> do+ outp <- nlopt_copy inp+ Opt <$> CFP.newForeignPtr outp (nlopt_destroy outp)++{- Random seeding functions -}++foreign import ccall "nlopt.h nlopt_srand"+ nlopt_srand :: CUInt -> IO ()++srand :: Integral a => a -> IO ()+srand = nlopt_srand . fromIntegral++foreign import ccall "nlopt.h nlopt_srand_time"+ nlopt_srand_time :: IO ()++srand_time :: IO ()+srand_time = nlopt_srand_time++{- Metadata -}++foreign import ccall "nlopt.h nlopt_version"+ nlopt_version :: Ptr CInt -> Ptr CInt -> Ptr CInt -> IO ()++-- | NLOPT library version, e.g. @2.4.2@+data Version = Version+ { major :: Int+ , minor :: Int+ , bugfix :: Int+ } deriving (Eq, Ord, Read, Show)++version :: IO Version+version =+ alloca $ \majptr ->+ alloca $ \minptr ->+ alloca $ \bfptr -> do+ nlopt_version majptr minptr bfptr+ Version <$> pk majptr <*> pk minptr <*> pk bfptr+ where+ pk = fmap fromIntegral . peek++foreign import ccall "nlopt.h nlopt_get_algorithm"+ nlopt_get_algorithm :: NloptOpt -> IO CInt++get_algorithm :: Opt -> IO Algorithm+get_algorithm = useOpt $ fmap parseEnum . nlopt_get_algorithm++foreign import ccall "nlopt.h nlopt_get_dimension"+ nlopt_get_dimension :: NloptOpt -> IO CUInt++get_dimension :: Opt -> IO Word+get_dimension = useOpt $ fmap fromIntegral . nlopt_get_dimension++{- Callback functions -}++asMVector :: CUInt -> Ptr CDouble -> IO (MV.IOVector Double)+asMVector dim ptr =+ MV.unsafeCast . flip MV.unsafeFromForeignPtr0 (fromIntegral dim) <$>+ newForeignPtr_ ptr++asVector :: CUInt -> Ptr CDouble -> IO (V.Vector Double)+asVector dim ptr =+ V.unsafeCast . flip V.unsafeFromForeignPtr0 (fromIntegral dim) <$>+ newForeignPtr_ ptr++type CFunc a = CUInt -> Ptr CDouble -> Ptr CDouble -> StablePtr a -> IO CDouble++-- | This function type corresponds to @nlopt_func@ in C and is used+-- for scalar functions of the parameter vector. You may pass data of+-- any type @a@ to the functions in this module that take a+-- 'ScalarFunction' as an argument; this data will be supplied to your+-- your function when it is called.+type ScalarFunction a+ = V.Vector Double -- ^ Parameter vector+ -> Maybe (MV.IOVector Double) -- ^ Gradient vector to be filled in+ -> a -- ^ User data+ -> IO Double -- ^ Scalar result++-- | This function type corresponds to @nlopt_mfunc@ in C and is used+-- for vector functions of the parameter vector. You may pass data of+-- any type @a@ to the functions in this module that take a+-- 'VectorFunction' as an argument; this data will be supplied to your+-- function when it is called.+type VectorFunction a+ = V.Vector Double -- ^ Parameter vector+ -> MV.IOVector Double -- ^ Output vector to be filled in+ -> Maybe (MV.IOVector Double) -- ^ Gradient vector to be filled in+ -> a -- ^ User data+ -> IO ()++-- | This function type corresponds to @nlopt_precond@ in C and is+-- used for functions that precondition a vector at a given point in+-- the parameter space. You may pass data of any type @a@ to the+-- functions in this module that take a 'PreconditionerFunction' as an+-- argument; this data will be supplied to your function when it is+-- called.+type PreconditionerFunction a+ = V.Vector Double -- ^ Parameter vector+ -> V.Vector Double -- ^ Vector @v@ to precondition+ -> MV.IOVector Double -- ^ Output vector @vpre@ to be filled in+ -> a -- ^ User data+ -> IO ()++wrapCFunction :: ScalarFunction a -> CFunc a+wrapCFunction cfunc dim stateptr gradientptr userptr = do+ nloptgradient <- asMVector dim gradientptr+ statevec <- asVector dim stateptr+ userdata <- deRefStablePtr userptr+ let+ gradptr = if gradientptr /= nullPtr+ then Just nloptgradient+ else Nothing+ realToFrac <$> cfunc statevec gradptr userdata++foreign import ccall safe "wrapper"+ mkCFunction :: CFunc a -> IO (FunPtr (CFunc a))++type CMFunc a = CUInt -> Ptr CDouble -> CUInt -> Ptr CDouble+ -> Ptr CDouble -> StablePtr a -> IO ()++wrapMFunction :: VectorFunction a -> CMFunc a+wrapMFunction mfunc constrdim constrptr dim stateptr gradientptr userptr+ = do+ nloptgradient <- asMVector (dim * constrdim) gradientptr+ nloptconstraint <- asMVector constrdim constrptr+ statevec <- asVector dim stateptr+ userdata <- deRefStablePtr userptr+ let+ gradptr = if gradientptr /= nullPtr+ then Just nloptgradient+ else Nothing+ mfunc statevec nloptconstraint gradptr userdata++foreign import ccall safe "wrapper"+ mkMFunction :: CMFunc a -> IO (FunPtr (CMFunc a))++type CPrecond a = CUInt -> Ptr CDouble -> Ptr CDouble+ -> Ptr CDouble -> StablePtr a -> IO ()++wrapPreconditioner :: PreconditionerFunction a -> CPrecond a+wrapPreconditioner prec dim stateptr vptr preptr userptr = do+ nloptpre <- asMVector dim preptr+ statevec <- asVector dim stateptr+ vvec <- asVector dim vptr+ userdata <- deRefStablePtr userptr+ prec statevec vvec nloptpre userdata++foreign import ccall safe "wrapper"+ mkPreconditionerFunction :: CPrecond a -> IO (FunPtr (CPrecond a))++-- We have to do the same silly dance with our user-data 'StablePtr's+-- as we do with function pointer wrappers: because NLOPT expects+-- these pointers before the actual optimization run, we have to+-- attach finalizers for them to the 'Opt' object so that they get+-- cleaned up properly.+addStablePtrFinalizer :: Opt -> StablePtr a -> IO ()+addStablePtrFinalizer (Opt p) sp =+ CFP.addForeignPtrFinalizer p (freeStablePtr sp)++getStablePtr :: Opt -> a -> IO (StablePtr a)+getStablePtr opt a = do+ aptr <- newStablePtr a+ addStablePtrFinalizer opt aptr+ return aptr++exportFunPtr :: (t1 -> IO (FunPtr a)) -> (t -> t1) -> t -> Opt -> IO (FunPtr a)+exportFunPtr mk wrap fun opt = do+ funptr <- mk $ wrap fun+ addFunPtrFinalizer opt funptr+ return funptr++{- Invoking the optimizer -}++-- | The output of an NLOPT optimizer run.+data Output = Output+ { resultCode :: Result -- ^ Return code+ , resultCost :: Double -- ^ Minimum of the objective+ -- function if optimization+ -- succeeded+ , resultParameters :: V.Vector Double -- ^ Parameters corresponding+ -- to the minimum if+ -- optimization succeeded++ , nEvals :: Int -- ^ number of evaluations+ }++foreign import ccall "nlopt.h nlopt_optimize"+ nlopt_optimize :: NloptOpt -> Ptr CDouble -> Ptr CDouble -> IO CInt++-- | This function is very similar to the C function @nlopt_optimize@,+-- but it does not use mutable vectors and returns an 'Output'+-- structure.+optimize :: Opt -- ^ Optimizer object set up to solve the problem+ -> V.Vector Double -- ^ Initial-guess parameter vector+ -> IO Output -- ^ Results of the optimization run+optimize optimizer x0 = withOpt optimizer $ \opt -> do+ vmut <- V.thaw $ V.unsafeCast x0+ (result, outputCost, iceout) <- alloca $ \costPtr -> do+ result <- MV.unsafeWith vmut $ \xptr ->+ parseEnum <$> nlopt_optimize opt xptr costPtr+ outputCost <- peek . castPtr $ costPtr+ iceout <- V.unsafeFreeze (MV.unsafeCast vmut)+ return (result, outputCost, iceout)+ nEvals <- fromIntegral <$> get_numevals optimizer+ return $ Output result outputCost iceout nEvals++{- Objective function setup -}++foreign import ccall "nlopt.h nlopt_set_min_objective"+ nlopt_set_min_objective :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> IO CInt++foreign import ccall "nlopt.h nlopt_set_max_objective"+ nlopt_set_max_objective :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> IO CInt++set_min_objective :: Opt -> ScalarFunction a -> a -> IO Result+set_min_objective opt objf userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_set_min_objective o objfunptr userptr++set_max_objective :: Opt -> ScalarFunction a -> a -> IO Result+set_max_objective opt objf userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_max_objective o objfunptr userptr++foreign import ccall "nlopt.h nlopt_set_precond_min_objective"+ nlopt_set_precond_min_objective :: NloptOpt+ -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a)+ -> StablePtr a+ -> IO CInt++foreign import ccall "nlopt.h nlopt_set_precond_max_objective"+ nlopt_set_precond_max_objective :: NloptOpt+ -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a)+ -> StablePtr a+ -> IO CInt++set_precond_min_objective :: Opt+ -> ScalarFunction a+ -> PreconditionerFunction a+ -> a+ -> IO Result+set_precond_min_objective opt objf pref userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ prefunptr <- exportFunPtr mkPreconditionerFunction wrapPreconditioner pref opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o -> parseEnum <$>+ nlopt_set_precond_min_objective o objfunptr prefunptr userptr++set_precond_max_objective :: Opt+ -> ScalarFunction a+ -> PreconditionerFunction a+ -> a+ -> IO Result+set_precond_max_objective opt objf pref userdata = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objf opt+ prefunptr <- exportFunPtr mkPreconditionerFunction wrapPreconditioner pref opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o -> parseEnum <$>+ nlopt_set_precond_max_objective o objfunptr prefunptr userptr++{- Working with bounds -}++foreign import ccall "nlopt.h nlopt_set_lower_bounds"+ nlopt_set_lower_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_lower_bounds1"+ nlopt_set_lower_bounds1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_lower_bounds"+ nlopt_get_lower_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_upper_bounds"+ nlopt_set_upper_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_upper_bounds1"+ nlopt_set_upper_bounds1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_upper_bounds"+ nlopt_get_upper_bounds :: NloptOpt -> Ptr CDouble -> IO CInt++set_lower_bounds :: Opt -> V.Vector Double -> IO Result+set_lower_bounds opt bounds =+ withForeignPtr (fst . V.unsafeToForeignPtr0 . V.unsafeCast $ bounds) $+ \bptr -> withOpt opt $ \o ->+ parseEnum <$> nlopt_set_lower_bounds o bptr++set_lower_bounds1 :: Opt -> Double -> IO Result+set_lower_bounds1 opt bound =+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_lower_bounds1 o (realToFrac bound)++get_lower_bounds :: Opt -> IO (V.Vector Double, Result)+get_lower_bounds opt = do+ v <- get_dimension opt >>= MV.new . fromIntegral+ MV.unsafeWith (MV.unsafeCast v) $ \vptr -> withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_lower_bounds o vptr+ retv <- V.unsafeFreeze v+ return (retv, result)++set_upper_bounds :: Opt -> V.Vector Double -> IO Result+set_upper_bounds opt bounds =+ withForeignPtr (fst . V.unsafeToForeignPtr0 . V.unsafeCast $ bounds) $+ \bptr -> withOpt opt $ \o ->+ parseEnum <$> nlopt_set_upper_bounds o bptr++set_upper_bounds1 :: Opt -> Double -> IO Result+set_upper_bounds1 opt bound =+ withOpt opt $ \o ->+ parseEnum <$> nlopt_set_upper_bounds1 o (realToFrac bound)++get_upper_bounds :: Opt -> IO (V.Vector Double, Result)+get_upper_bounds opt = do+ v <- get_dimension opt >>= MV.new . fromIntegral+ MV.unsafeWith (MV.unsafeCast v) $ \vptr -> withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_upper_bounds o vptr+ retv <- V.unsafeFreeze v+ return (retv, result)++{- Working with constraints -}++foreign import ccall "nlopt.h nlopt_remove_inequality_constraints"+ nlopt_remove_inequality_constraints :: NloptOpt -> IO CInt++foreign import ccall "nlopt.h nlopt_add_inequality_constraint"+ nlopt_add_inequality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_precond_inequality_constraint"+ nlopt_add_precond_inequality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a) -> StablePtr a+ -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_inequality_mconstraint"+ nlopt_add_inequality_mconstraint :: NloptOpt -> CUInt -> FunPtr (CMFunc a)+ -> StablePtr a -> CDouble -> IO CInt++remove_inequality_constraints :: Opt -> IO Result+remove_inequality_constraints =+ useOpt $ fmap parseEnum . nlopt_remove_inequality_constraints++add_inequality_constraint :: Opt -> ScalarFunction a+ -> a -> Double -> IO Result+add_inequality_constraint opt objfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_inequality_constraint o objfunptr userptr (realToFrac tol)++add_precond_inequality_constraint :: Opt -> ScalarFunction a+ -> PreconditionerFunction a -> a -> Double+ -> IO Result+add_precond_inequality_constraint opt objfun precfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ precfunptr <-+ exportFunPtr mkPreconditionerFunction wrapPreconditioner precfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_precond_inequality_constraint o objfunptr+ precfunptr userptr (realToFrac tol)++add_inequality_mconstraint :: Opt -> Word -> VectorFunction a -> a+ -> Double -> IO Result+add_inequality_mconstraint opt constraintsize constrfun userdata tol = do+ constrfunptr <- exportFunPtr mkMFunction wrapMFunction constrfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_inequality_mconstraint o (fromIntegral constraintsize)+ constrfunptr userptr (realToFrac tol)++foreign import ccall "nlopt.h nlopt_remove_equality_constraints"+ nlopt_remove_equality_constraints :: NloptOpt -> IO CInt++foreign import ccall "nlopt.h nlopt_add_equality_constraint"+ nlopt_add_equality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> StablePtr a -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_precond_equality_constraint"+ nlopt_add_precond_equality_constraint :: NloptOpt -> FunPtr (CFunc a)+ -> FunPtr (CPrecond a) -> StablePtr a+ -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_add_equality_mconstraint"+ nlopt_add_equality_mconstraint :: NloptOpt -> CUInt -> FunPtr (CMFunc a)+ -> StablePtr a -> CDouble -> IO CInt++remove_equality_constraints :: Opt -> IO Result+remove_equality_constraints =+ useOpt $ fmap parseEnum . nlopt_remove_equality_constraints++add_equality_constraint :: Opt -> ScalarFunction a+ -> a -> Double -> IO Result+add_equality_constraint opt objfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_equality_constraint o objfunptr userptr (realToFrac tol)++add_precond_equality_constraint :: Opt -> ScalarFunction a+ -> PreconditionerFunction a -> a -> Double+ -> IO Result+add_precond_equality_constraint opt objfun precfun userdata tol = do+ objfunptr <- exportFunPtr mkCFunction wrapCFunction objfun opt+ precfunptr <-+ exportFunPtr mkPreconditionerFunction wrapPreconditioner precfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_precond_equality_constraint o objfunptr+ precfunptr userptr (realToFrac tol)++add_equality_mconstraint :: Opt -> Word -> VectorFunction a -> a+ -> Double -> IO Result+add_equality_mconstraint opt constraintsize constrfun userdata tol = do+ constrfunptr <- exportFunPtr mkMFunction wrapMFunction constrfun opt+ userptr <- getStablePtr opt userdata+ withOpt opt $ \o ->+ parseEnum <$>+ nlopt_add_equality_mconstraint o (fromIntegral constraintsize)+ constrfunptr userptr (realToFrac tol)++{- Stopping criteria -}++withInputVector :: (Storable c, Storable a)+ => V.Vector c -> (Ptr a -> IO b) -> IO b+withInputVector = withForeignPtr . fst . V.unsafeToForeignPtr0 . V.unsafeCast+withOutputVector :: (Storable c, Storable a)+ => V.MVector s c -> (Ptr a -> IO b) -> IO b+withOutputVector = withForeignPtr . fst . MV.unsafeToForeignPtr0 . MV.unsafeCast++setScalar :: (Enum a, Integral b) => (NloptOpt -> t1 -> IO b)+ -> (t -> t1) -> Opt -> t -> IO a+setScalar setter conv opt val = withOpt opt $ \o ->+ parseEnum <$> setter o (conv val)++getScalar :: (NloptOpt -> IO b) -> (b -> a) -> Opt -> IO a+getScalar getter conv = useOpt $ fmap conv . getter++foreign import ccall "nlopt.h nlopt_set_stopval"+ nlopt_set_stopval :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_stopval"+ nlopt_get_stopval :: NloptOpt -> IO CDouble++set_stopval :: Opt -> Double -> IO Result+set_stopval = setScalar nlopt_set_stopval realToFrac++get_stopval :: Opt -> IO Double+get_stopval = getScalar nlopt_get_stopval realToFrac++foreign import ccall "nlopt.h nlopt_set_ftol_rel"+ nlopt_set_ftol_rel :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_ftol_rel"+ nlopt_get_ftol_rel :: NloptOpt -> IO CDouble++set_ftol_rel :: Opt -> Double -> IO Result+set_ftol_rel = setScalar nlopt_set_ftol_rel realToFrac++get_ftol_rel :: Opt -> IO Double+get_ftol_rel = getScalar nlopt_get_ftol_rel realToFrac++foreign import ccall "nlopt.h nlopt_set_ftol_abs"+ nlopt_set_ftol_abs :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_ftol_abs"+ nlopt_get_ftol_abs :: NloptOpt -> IO CDouble++set_ftol_abs :: Opt -> Double -> IO Result+set_ftol_abs = setScalar nlopt_set_ftol_abs realToFrac++get_ftol_abs :: Opt -> IO Double+get_ftol_abs = getScalar nlopt_get_ftol_abs realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_rel"+ nlopt_set_xtol_rel :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_xtol_rel"+ nlopt_get_xtol_rel :: NloptOpt -> IO CDouble++set_xtol_rel :: Opt -> Double -> IO Result+set_xtol_rel = setScalar nlopt_set_xtol_rel realToFrac++get_xtol_rel :: Opt -> IO Double+get_xtol_rel = getScalar nlopt_get_xtol_rel realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_abs1"+ nlopt_set_xtol_abs1 :: NloptOpt -> CDouble -> IO CInt++set_xtol_abs1 :: Opt -> Double -> IO Result+set_xtol_abs1 = setScalar nlopt_set_xtol_abs1 realToFrac++foreign import ccall "nlopt.h nlopt_set_xtol_abs"+ nlopt_set_xtol_abs :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_xtol_abs"+ nlopt_get_xtol_abs :: NloptOpt -> Ptr CDouble -> IO CInt++set_xtol_abs :: Opt -> V.Vector Double -> IO Result+set_xtol_abs opt tolvec =+ withInputVector tolvec $ \tolptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_xtol_abs o tolptr++get_xtol_abs :: Opt -> IO (Result, V.Vector Double)+get_xtol_abs opt = do+ mutv <- get_dimension opt >>= MV.new . fromIntegral+ withOutputVector mutv $ \vecptr ->+ withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_xtol_abs o vecptr+ outvec <- V.unsafeFreeze mutv+ return (result, outvec)++foreign import ccall "nlopt.h nlopt_set_maxeval"+ nlopt_set_maxeval :: NloptOpt -> CInt -> IO CInt++foreign import ccall "nlopt.h nlopt_get_maxeval"+ nlopt_get_maxeval :: NloptOpt -> IO CInt++set_maxeval :: Opt -> Word -> IO Result+set_maxeval = setScalar nlopt_set_maxeval fromIntegral++get_maxeval :: Opt -> IO Word+get_maxeval = getScalar nlopt_get_maxeval fromIntegral++foreign import ccall "nlopt.h nlopt_get_numevals"+ nlopt_get_numevals :: NloptOpt -> IO CInt++get_numevals :: Opt -> IO CInt+get_numevals = getScalar nlopt_get_numevals fromIntegral++foreign import ccall "nlopt.h nlopt_set_maxtime"+ nlopt_set_maxtime :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_maxtime"+ nlopt_get_maxtime :: NloptOpt -> IO CDouble++set_maxtime :: Opt -> Double -> IO Result+set_maxtime = setScalar nlopt_set_maxtime realToFrac++get_maxtime :: Opt -> IO Double+get_maxtime = getScalar nlopt_get_maxtime realToFrac++foreign import ccall "nlopt.h nlopt_force_stop"+ nlopt_force_stop :: NloptOpt -> IO CInt++force_stop :: Opt -> IO Result+force_stop = useOpt $ fmap parseEnum . nlopt_force_stop++foreign import ccall "nlopt.h nlopt_set_force_stop"+ nlopt_set_force_stop :: NloptOpt -> CInt -> IO CInt++foreign import ccall "nlopt.h nlopt_get_force_stop"+ nlopt_get_force_stop :: NloptOpt -> IO CInt++set_force_stop :: Opt -> Word -> IO Result+set_force_stop = setScalar nlopt_set_force_stop fromIntegral++get_force_stop :: Opt -> IO Word+get_force_stop = getScalar nlopt_get_force_stop fromIntegral++{- Algorithm-specific configuration -}++foreign import ccall "nlopt.h nlopt_set_local_optimizer"+ nlopt_set_local_optimizer :: NloptOpt -> NloptOpt -> IO CInt++set_local_optimizer :: Opt -- ^ Primary optimizer+ -> Opt -- ^ Subsidiary (local) optimizer+ -> IO Result+set_local_optimizer p s =+ withOpt p $ \primary -> withOpt s $ \secondary ->+ parseEnum <$> nlopt_set_local_optimizer primary secondary++foreign import ccall "nlopt.h nlopt_set_population"+ nlopt_set_population :: NloptOpt -> Word -> IO CInt++foreign import ccall "nlopt.h nlopt_get_population"+ nlopt_get_population :: NloptOpt -> IO Word++set_population :: Opt -> Word -> IO Result+set_population = setScalar nlopt_set_population fromIntegral++get_population :: Opt -> IO Word+get_population = getScalar nlopt_get_population fromIntegral++foreign import ccall "nlopt.h nlopt_set_vector_storage"+ nlopt_set_vector_storage :: NloptOpt -> Word -> IO CInt++foreign import ccall "nlopt.h nlopt_get_vector_storage"+ nlopt_get_vector_storage :: NloptOpt -> IO Word++set_vector_storage :: Opt -> Word -> IO Result+set_vector_storage = setScalar nlopt_set_vector_storage fromIntegral++get_vector_storage :: Opt -> IO Word+get_vector_storage = getScalar nlopt_get_vector_storage fromIntegral++foreign import ccall "nlopt.h nlopt_set_default_initial_step"+ nlopt_set_default_initial_step :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_initial_step"+ nlopt_set_initial_step :: NloptOpt -> Ptr CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_set_initial_step1"+ nlopt_set_initial_step1 :: NloptOpt -> CDouble -> IO CInt++foreign import ccall "nlopt.h nlopt_get_initial_step"+ nlopt_get_initial_step :: NloptOpt -> Ptr CDouble -> Ptr CDouble -> IO CInt++set_default_initial_step :: Opt -> V.Vector Double -> IO Result+set_default_initial_step opt stepvec =+ withInputVector stepvec $ \stepptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_default_initial_step o stepptr++set_initial_step :: Opt -> V.Vector Double -> IO Result+set_initial_step opt stepvec =+ withInputVector stepvec $ \stepptr ->+ withOpt opt $ \o -> parseEnum <$> nlopt_set_initial_step o stepptr++set_initial_step1 :: Opt -> Double -> IO Result+set_initial_step1 = setScalar nlopt_set_initial_step1 realToFrac++get_initial_step :: Opt -> V.Vector Double -> IO (Result, V.Vector Double)+get_initial_step opt xvec = do+ mutv <- get_dimension opt >>= MV.new . fromIntegral+ withOutputVector mutv $ \outptr ->+ withInputVector xvec $ \inptr ->+ withOpt opt $ \o -> do+ result <- parseEnum <$> nlopt_get_initial_step o inptr outptr+ outvec <- V.unsafeFreeze mutv+ return (result, outvec)
+ src/Text/ParseSR.hs view
@@ -0,0 +1,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
+ src/Text/ParseSR/IO.hs view
@@ -0,0 +1,73 @@+{-# language LambdaCase #-}+-----------------------------------------------------------------------------+-- |+-- Module : Text.ParseSR.IO+-- Copyright : (c) Fabricio Olivetti 2021 - 2024+-- License : BSD3+-- Maintainer : fabricio.olivetti@gmail.com+-- Stability : experimental+-- Portability : ConstraintKinds+--+-- Functions to parse multiple expressions from stdin or a text file.+--+-----------------------------------------------------------------------------+module Text.ParseSR.IO ( withInput, withOutput, withOutputDebug )+ where++-- import Data.SRTree.EqSat1+--import Algorithm.EqSat.Simplify ( simplifyEqSatDefault )+import Control.Monad (forM_, unless)+import qualified Data.ByteString.Char8 as B+import Data.SRTree+import Data.SRTree.Recursion (Fix (..))+import System.IO+import Text.ParseSR (Output, SRAlgs, parseSR, showOutput)++-- | given a filename, the symbolic regression algorithm, a string of variables name, +-- and two booleans indicating whether to convert float values to parameters and +-- whether to simplify the expression or not, it will read the file and parse everything +-- returning a list of either an error message or a tree.+--+-- empty filename defaults to stdin +withInput :: String -> SRAlgs -> String -> Bool -> Bool -> IO [Either String (Fix SRTree)]+withInput fname sr hd param simpl = do+ h <- if null fname then pure stdin else openFile fname ReadMode+ contents <- hGetLines h + let myParserFun = parseSR sr (B.pack hd) param . B.pack+ -- myParser = if simpl then fmap simplifyEqSat . myParserFun else myParserFun+ myParser = 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
srtree.cabal view
@@ -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
+ test/EqSatTests.hs view
@@ -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+ ]
test/Spec.hs view
@@ -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 ()
+ test/StoreTests.hs view
@@ -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+ ]