diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,10 @@ # Changelog for srtree +## 1.0.0.2 + +- Export `Fix` from `Data.SRTRee` +- `paramsToConst` function + ## 1.0.0.1 - Fix `vector` version bounds diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -19,10 +19,10 @@ alg (Bin op l r) = \xs -> evalOp op (l xs) (r xs) alg (Uni f t) = \xs -> evalFun f (t xs) -xs = fromList [1.0, 2.0] -ps = fromList [0.1, 0.2 .. 30.0] +xs = fromList [1.0, 2.0 .. 3000.0] +ps = fromList [0.1, 0.2 .. 110.0] -params = P [0,1] (-1.0, 1.0) (-2, 2) [Id ..] +params = P [0,1] (-1.0, 1.0) (-2, 2) [Id] -- , Sin, Cos, Log, Exp] runRnd g ns p = flip evalStateT g $ traverse (\n -> runReaderT (randomTree n) p) ns runRndBalance g ns p = flip evalStateT g $ traverse (\n -> runReaderT (randomTreeBalanced n) p) ns @@ -38,16 +38,59 @@ benchAutodiff f = do ts <- f g lens params let ps' = Data.Vector.toList ps pure $ map ((`autograd` ps') . relabelParams . fst . constsToParam) ts +{- main :: IO () main = defaultMain [ bgroup "unbalanced" - [ bench "forwardMode" $ nfIO (benchTree runRnd forwardMode) - , bench "grad" $ nfIO (benchTree runRnd gradParams) + [ -- bench "forwardMode" $ nfIO (benchTree runRnd forwardMode) + bench "grad" $ nfIO (benchTree runRnd gradParams) + , bench "grad2" $ nfIO (benchTree runRnd gradParams2) , bench "autodiff" $ nfIO (benchAutodiff runRnd) ] , bgroup "balanced" - [ bench "forwardMode" $ nfIO (benchTree runRndBalance forwardMode) - , bench "grad" $ nfIO (benchTree runRndBalance gradParams) + [ -- bench "forwardMode" $ nfIO (benchTree runRndBalance forwardMode) + bench "grad" $ nfIO (benchTree runRndBalance gradParams) + , bench "grad2" $ nfIO (benchTree runRndBalance gradParams2) , bench "autodiff" $ nfIO (benchAutodiff runRndBalance) ] ] +-} + +mkPySRTree :: Int -> Int -> Fix SRTree +mkPySRTree nvar np = relabelParams $ sum [mkWith ix | ix <- [0 .. nvar-1]] + where + mkWith ix = fst . (!!(np-1)) $ iterate (\(t, i) -> (cos (t + param i), i+1)) $ (cos (var ix + param 0), 1) + +genBalancedTree :: Int -> Fix SRTree +genBalancedTree = relabelParams . go + where + go 0 = var 0 + go 1 = cos (var 0 + param 0) + go n | even n = go (n `div` 2) * (param 0 + go (n `div` 2 - 1)) + | odd n = cos (param 0 + go (n-1)) + +sizes = (,) <$> [1] <*> [10, 20 .. 1000] +--sizes = (,) <$> [1] <*> [500, 600 .. 2000] +--tests = map (\(ix, iy) -> (ix, iy, mkPySRTree ix iy)) sizes +tests = map (\(ix, iy) -> (ix, iy, genBalancedTree iy)) sizes + +main :: IO () +main = defaultMain [ + bgroup ("PySR " <> show ix <> " " <> show iy) + [ bench "warmup" $ whnf (evalTree xs ps id) t + -- , bench "forwardMode" $ whnf (sum . forwardMode xs ps id) t + , bench "grad" $ whnf (sum . snd . gradParamsFwd xs ps id) t + , bench "grad2" $ whnf (sum . snd . gradParamsRev xs ps id) t + , bench "autodiff" $ whnf (sum . (`autograd` (Data.Vector.toList ps))) t + ] | (ix, iy, t) <- tests + ] + +{- +comp (x, xs) (y, ys) = ((x + sum xs) - (y + sum ys))^2 + +main :: IO () +main = do + let g = mkStdGen 42 + trees <- runRndBalance g lens params + mapM_ (\(t,a,b) -> print (showExpr t, a, b, comp a b)) $ filter (\(t,a,b) -> let c = comp a b in not (isNaN c) && c >= 1e-20) [(t, gradParams xs ps id t, gradParams2 xs ps id t) | t' <- trees, let t = relabelParams (fst $ constsToParam t')] +-} diff --git a/src/Data/SRTree.hs b/src/Data/SRTree.hs --- a/src/Data/SRTree.hs +++ b/src/Data/SRTree.hs @@ -28,7 +28,8 @@ , deriveByParam , derivative , forwardMode - , gradParams + , gradParamsFwd + , gradParamsRev , evalFun , evalOp , inverseFunc @@ -36,6 +37,8 @@ , relabelParams , constsToParam , floatConstsToParam + , paramsToConst + , Fix (..) ) where @@ -57,7 +60,8 @@ , deriveByParam , derivative , forwardMode - , gradParams + , gradParamsFwd + , gradParamsRev , evalFun , evalOp , inverseFunc @@ -65,4 +69,6 @@ , relabelParams , constsToParam , floatConstsToParam + , paramsToConst + , Fix (..) ) diff --git a/src/Data/SRTree/Internal.hs b/src/Data/SRTree/Internal.hs --- a/src/Data/SRTree/Internal.hs +++ b/src/Data/SRTree/Internal.hs @@ -1,6 +1,7 @@ {-# language FlexibleInstances, DeriveFunctor #-} {-# language ScopedTypeVariables #-} {-# language RankNTypes #-} +{-# language ViewPatterns #-} ----------------------------------------------------------------------------- -- | -- Module : Data.SRTree.Internal @@ -32,7 +33,8 @@ , deriveByParam , derivative , forwardMode - , gradParams + , gradParamsFwd + , gradParamsRev , evalFun , evalOp , inverseFunc @@ -40,14 +42,18 @@ , relabelParams , constsToParam , floatConstsToParam + , paramsToConst + , Fix (..) ) where -import Data.SRTree.Recursion ( Fix(Fix), cata, mutu, cataM ) +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) @@ -248,13 +254,13 @@ -- | Count the occurrences of variable indexed as `ix` countOccurrences :: Int -> Fix SRTree -> Int -countOccurrences ix = sum . cata alg +countOccurrences ix = cata alg where - alg (Var iy) = [1 | ix == iy] - alg Param {} = [] - alg Const {} = [] + alg (Var iy) = if ix == iy then 1 else 0 + alg Param {} = 0 + alg Const {} = 0 alg (Uni _ t) = t - alg (Bin _ l r) = l <> r + 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. @@ -406,22 +412,74 @@ alg2 (Bin op l r) = evalOp op (snd l) (snd r) -- | 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`. -gradParams :: (Show a, Num a, Floating a) => V.Vector a -> V.Vector Double -> (Double -> a) -> Fix SRTree -> (a, [a]) -gradParams xss theta f = cata alg +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 where n = V.length theta - alg (Var ix) = (xss ! ix, []) - alg (Param ix) = (f $ theta ! ix, [1]) - alg (Const c) = (f c, []) - alg (Uni f (v, gs)) = let v' = evalFun f v in (v', map (* derivative f v) gs) - alg (Bin Add (v1, l) (v2, r)) = (v1+v2, l ++ r) - alg (Bin Sub (v1, l) (v2, r)) = (v1-v2, l ++ map negate r) - alg (Bin Mul (v1, l) (v2, r)) = (v1*v2, map (*v2) l ++ map (*v1) r) - alg (Bin Div (v1, l) (v2, r)) = (v1/v2, map (/v2) l ++ map ((/v2^2) . (*v1) . negate) r) - alg (Bin Power (v1, l) (v2, r)) = (v1 ** v2, map (* (v1 ** (v2 - 1))) (map (*v2) l ++ map ((*v1).(* log v1)) r)) + 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))) +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) +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) + 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 @@ -494,3 +552,13 @@ 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) + +-- | Convert the parameters into constants in the tree +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 (Bin f l r) = Fix $ Bin f l r diff --git a/src/Data/SRTree/Recursion.hs b/src/Data/SRTree/Recursion.hs --- a/src/Data/SRTree/Recursion.hs +++ b/src/Data/SRTree/Recursion.hs @@ -27,6 +27,9 @@ cata :: Functor f => (f a -> a) -> Fix f -> a cata alg = alg . fmap (cata alg) . unfix +--zigzag :: Functor f => (f a -> a) -> Fix f -> a +--zigzag alg = + cataM :: (Functor f, Monad m) => (forall x . f (m x) -> m (f x)) -> (f a -> m a) -> Fix f -> m a cataM seq alg = cata (seq >=> alg) diff --git a/srtree.cabal b/srtree.cabal --- a/srtree.cabal +++ b/srtree.cabal @@ -5,7 +5,7 @@ -- see: https://github.com/sol/hpack name: srtree -version: 1.0.0.1 +version: 1.0.0.2 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 @@ -39,6 +39,7 @@ build-depends: base >=4.16 && <4.18 , containers ==0.6.* + , dlist ==1.0.* , mtl ==2.2.* , random ==1.2.* , vector >=0.12 && <0.14 @@ -56,6 +57,7 @@ , base >=4.16 && <4.18 , containers ==0.6.* , criterion >=1.5.0 && <1.7 + , dlist ==1.0.* , mtl ==2.2.* , random ==1.2.* , srtree @@ -75,6 +77,7 @@ , ad , base >=4.16 && <4.18 , containers ==0.6.* + , dlist ==1.0.* , mtl ==2.2.* , random ==1.2.* , srtree