{-# language FlexibleInstances, DeriveFunctor #-}
{-# language ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.SRTree.Internal 
-- Copyright   :  (c) Fabricio Olivetti 2021 - 2021
-- License     :  BSD3
-- Maintainer  :  fabricio.olivetti@gmail.com
-- Stability   :  experimental
-- Portability :  FlexibleInstances, DeriveFunctor, ScopedTypeVariables
--
-- Expression tree for Symbolic Regression
--
-----------------------------------------------------------------------------

module Data.SRTree.Internal 
         ( SRTree(..)
         , Function(..)
         , OptIntPow(..)
         , traverseIx
         , arity
         , getChildren
         , countNodes
         , countVarNodes
         , countOccurrences
         , deriveBy
         , deriveParamBy
         , simplify
         , derivative
         , evalFun
         , inverseFunc
         , evalTree
         , evalTreeMap
         , evalTreeWithMap
         , evalTreeWithVector
         , relabelOccurrences
         , relabelParams
         )
         where

import Data.Bifunctor

import Data.Map.Strict (Map(..), (!), (!?), insert, fromList)
import qualified Data.Map.Strict as M
import qualified Data.Vector as V
import Control.Monad.State
import Control.Monad.Reader
import Control.Applicative hiding (Const)

-- | Tree structure to be used with Symbolic Regression algorithms.
-- This structure is parametrized by the indexing type to retrieve the values
-- of a variable and the type of the output value.
data SRTree ix val = 
   Empty 
 | Var ix
 | Const val
 | Param ix
 | Fun Function (SRTree ix val)
 | Pow (SRTree ix val) Int
 | SRTree ix val `Add`     SRTree ix val
 | SRTree ix val `Sub`     SRTree ix val
 | SRTree ix val `Mul`     SRTree ix val
 | SRTree ix val `Div`     SRTree ix val
 | SRTree ix val `Power`   SRTree ix val
 | SRTree ix val `LogBase` SRTree ix val
     deriving (Int -> SRTree ix val -> ShowS
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall ix val. (Show ix, Show val) => Int -> SRTree ix val -> ShowS
forall ix val. (Show ix, Show val) => [SRTree ix val] -> ShowS
forall ix val. (Show ix, Show val) => SRTree ix val -> String
showList :: [SRTree ix val] -> ShowS
$cshowList :: forall ix val. (Show ix, Show val) => [SRTree ix val] -> ShowS
show :: SRTree ix val -> String
$cshow :: forall ix val. (Show ix, Show val) => SRTree ix val -> String
showsPrec :: Int -> SRTree ix val -> ShowS
$cshowsPrec :: forall ix val. (Show ix, Show val) => Int -> SRTree ix val -> ShowS
Show, SRTree ix val -> SRTree ix val -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall ix val.
(Eq ix, Eq val) =>
SRTree ix val -> SRTree ix val -> Bool
/= :: SRTree ix val -> SRTree ix val -> Bool
$c/= :: forall ix val.
(Eq ix, Eq val) =>
SRTree ix val -> SRTree ix val -> Bool
== :: SRTree ix val -> SRTree ix val -> Bool
$c== :: forall ix val.
(Eq ix, Eq val) =>
SRTree ix val -> SRTree ix val -> Bool
Eq, SRTree ix val -> SRTree ix val -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {ix} {val}. (Ord ix, Ord val) => Eq (SRTree ix val)
forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Bool
forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Ordering
forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> SRTree ix val
min :: SRTree ix val -> SRTree ix val -> SRTree ix val
$cmin :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> SRTree ix val
max :: SRTree ix val -> SRTree ix val -> SRTree ix val
$cmax :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> SRTree ix val
>= :: SRTree ix val -> SRTree ix val -> Bool
$c>= :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Bool
> :: SRTree ix val -> SRTree ix val -> Bool
$c> :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Bool
<= :: SRTree ix val -> SRTree ix val -> Bool
$c<= :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Bool
< :: SRTree ix val -> SRTree ix val -> Bool
$c< :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Bool
compare :: SRTree ix val -> SRTree ix val -> Ordering
$ccompare :: forall ix val.
(Ord ix, Ord val) =>
SRTree ix val -> SRTree ix val -> Ordering
Ord, forall a b. a -> SRTree ix b -> SRTree ix a
forall a b. (a -> b) -> SRTree ix a -> SRTree ix b
forall ix a b. a -> SRTree ix b -> SRTree ix a
forall ix a b. (a -> b) -> SRTree ix a -> SRTree ix b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> SRTree ix b -> SRTree ix a
$c<$ :: forall ix a b. a -> SRTree ix b -> SRTree ix a
fmap :: forall a b. (a -> b) -> SRTree ix a -> SRTree ix b
$cfmap :: forall ix a b. (a -> b) -> SRTree ix a -> SRTree ix b
Functor)

-- | Functions that can be applied to a subtree.
data Function = 
    Id 
  | Abs
  | Sin 
  | Cos 
  | Tan 
  | Sinh
  | Cosh
  | Tanh 
  | ASin 
  | ACos 
  | ATan
  | ASinh
  | ACosh 
  | ATanh 
  | Sqrt 
  | Cbrt 
  | Square 
  | Log 
  | Exp 
     deriving (Int -> Function -> ShowS
[Function] -> ShowS
Function -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Function] -> ShowS
$cshowList :: [Function] -> ShowS
show :: Function -> String
$cshow :: Function -> String
showsPrec :: Int -> Function -> ShowS
$cshowsPrec :: Int -> Function -> ShowS
Show, ReadPrec [Function]
ReadPrec Function
Int -> ReadS Function
ReadS [Function]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Function]
$creadListPrec :: ReadPrec [Function]
readPrec :: ReadPrec Function
$creadPrec :: ReadPrec Function
readList :: ReadS [Function]
$creadList :: ReadS [Function]
readsPrec :: Int -> ReadS Function
$creadsPrec :: Int -> ReadS Function
Read, Function -> Function -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Function -> Function -> Bool
$c/= :: Function -> Function -> Bool
== :: Function -> Function -> Bool
$c== :: Function -> Function -> Bool
Eq, Eq Function
Function -> Function -> Bool
Function -> Function -> Ordering
Function -> Function -> Function
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Function -> Function -> Function
$cmin :: Function -> Function -> Function
max :: Function -> Function -> Function
$cmax :: Function -> Function -> Function
>= :: Function -> Function -> Bool
$c>= :: Function -> Function -> Bool
> :: Function -> Function -> Bool
$c> :: Function -> Function -> Bool
<= :: Function -> Function -> Bool
$c<= :: Function -> Function -> Bool
< :: Function -> Function -> Bool
$c< :: Function -> Function -> Bool
compare :: Function -> Function -> Ordering
$ccompare :: Function -> Function -> Ordering
Ord, Int -> Function
Function -> Int
Function -> [Function]
Function -> Function
Function -> Function -> [Function]
Function -> Function -> Function -> [Function]
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
enumFromThenTo :: Function -> Function -> Function -> [Function]
$cenumFromThenTo :: Function -> Function -> Function -> [Function]
enumFromTo :: Function -> Function -> [Function]
$cenumFromTo :: Function -> Function -> [Function]
enumFromThen :: Function -> Function -> [Function]
$cenumFromThen :: Function -> Function -> [Function]
enumFrom :: Function -> [Function]
$cenumFrom :: Function -> [Function]
fromEnum :: Function -> Int
$cfromEnum :: Function -> Int
toEnum :: Int -> Function
$ctoEnum :: Int -> Function
pred :: Function -> Function
$cpred :: Function -> Function
succ :: Function -> Function
$csucc :: Function -> Function
Enum)

-- | A class for optimized `(^^)` operators for specific types.
-- This was created because the integer power operator for
-- interval arithmetic must be aware of the dependency problem,
-- thus the default `(^)` doesn't work.
class OptIntPow a where
  (^.) :: a -> Int -> a
  infix 8 ^.
  
instance OptIntPow Double where
  ^. :: Double -> Int -> Double
(^.) = forall a b. (Fractional a, Integral b) => a -> b -> a
(^^)
  {-# INLINE (^.) #-}
instance OptIntPow Float where
  ^. :: Float -> Int -> Float
(^.) = forall a b. (Fractional a, Integral b) => a -> b -> a
(^^)
  {-# INLINE (^.) #-}
  
 
instance (Eq ix, Eq val, Num val, OptIntPow val) => OptIntPow (SRTree ix val) where
  SRTree ix val
t ^. :: SRTree ix val -> Int -> SRTree ix val
^. Int
0         = SRTree ix val
1
  SRTree ix val
t ^. Int
1         = SRTree ix val
t
  (Const val
c) ^. Int
k = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c forall a. OptIntPow a => a -> Int -> a
^. Int
k 
  SRTree ix val
t ^. Int
k         = forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow SRTree ix val
t Int
k
  {-# INLINE (^.) #-}
       
instance (Eq ix, Eq val, Num val) => Num (SRTree ix val) where
  SRTree ix val
0 + :: SRTree ix val -> SRTree ix val -> SRTree ix val
+ SRTree ix val
r                   = SRTree ix val
r
  SRTree ix val
l + SRTree ix val
0                   = SRTree ix val
l
  (Const val
c1) + (Const val
c2) = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c1 forall a. Num a => a -> a -> a
+ val
c2
  SRTree ix val
l + SRTree ix val
r                   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add SRTree ix val
l SRTree ix val
r
  {-# INLINE (+) #-}

  SRTree ix val
0 - :: SRTree ix val -> SRTree ix val -> SRTree ix val
- SRTree ix val
r                   = (-SRTree ix val
1) forall a. Num a => a -> a -> a
* SRTree ix val
r
  SRTree ix val
l - SRTree ix val
0                   = SRTree ix val
l
  (Const val
c1) - (Const val
c2) = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c1 forall a. Num a => a -> a -> a
- val
c2
  SRTree ix val
l - SRTree ix val
r                   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub SRTree ix val
l SRTree ix val
r
  {-# INLINE (-) #-}

  SRTree ix val
0 * :: SRTree ix val -> SRTree ix val -> SRTree ix val
* SRTree ix val
r                   = SRTree ix val
0
  SRTree ix val
l * SRTree ix val
0                   = SRTree ix val
0
  SRTree ix val
1 * SRTree ix val
r                   = SRTree ix val
r
  SRTree ix val
l * SRTree ix val
1                   = SRTree ix val
l 
  (Const val
c1) * (Const val
c2) = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c1 forall a. Num a => a -> a -> a
* val
c2
  SRTree ix val
l * SRTree ix val
r                   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul SRTree ix val
l SRTree ix val
r
  {-# INLINE (*) #-}
    
  abs :: SRTree ix val -> SRTree ix val
abs         = forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Abs
  {-# INLINE abs #-}
  
  negate :: SRTree ix val -> SRTree ix val
negate (Const val
x) = forall ix val. val -> SRTree ix val
Const (forall a. Num a => a -> a
negate val
x)
  negate SRTree ix val
t         = forall ix val. val -> SRTree ix val
Const (-val
1) forall a. Num a => a -> a -> a
* SRTree ix val
t
  {-# INLINE negate #-}
  
  signum :: SRTree ix val -> SRTree ix val
signum SRTree ix val
t    = case SRTree ix val
t of
                  Const val
x -> forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ forall a. Num a => a -> a
signum val
x
                  SRTree ix val
_       -> forall ix val. val -> SRTree ix val
Const val
0
  fromInteger :: Integer -> SRTree ix val
fromInteger Integer
x = forall ix val. val -> SRTree ix val
Const (forall a. Num a => Integer -> a
fromInteger Integer
x)
  {-# INLINE fromInteger #-}

instance (Eq ix, Eq val, Fractional val) => Fractional (SRTree ix val) where
  SRTree ix val
0 / :: SRTree ix val -> SRTree ix val -> SRTree ix val
/ SRTree ix val
r                   = SRTree ix val
0
  SRTree ix val
l / SRTree ix val
1                   = SRTree ix val
l
  (Const val
c1) / (Const val
c2) = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c1forall a. Fractional a => a -> a -> a
/val
c2
  SRTree ix val
l / SRTree ix val
r                   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div SRTree ix val
l SRTree ix val
r
  {-# INLINE (/) #-}
  
  fromRational :: Rational -> SRTree ix val
fromRational = forall ix val. val -> SRTree ix val
Const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Fractional a => Rational -> a
fromRational
  {-# INLINE fromRational #-}
  
instance (Eq ix, Eq val, Floating val) => Floating (SRTree ix val) where  
  pi :: SRTree ix val
pi      = forall ix val. val -> SRTree ix val
Const  forall a. Floating a => a
pi
  {-# INLINE pi #-}
  exp :: SRTree ix val -> SRTree ix val
exp     = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Exp
  {-# INLINE exp #-}
  log :: SRTree ix val -> SRTree ix val
log     = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Log
  {-# INLINE log #-}
  sqrt :: SRTree ix val -> SRTree ix val
sqrt    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Sqrt
  {-# INLINE sqrt #-}
  sin :: SRTree ix val -> SRTree ix val
sin     = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Sin
  {-# INLINE sin #-}
  cos :: SRTree ix val -> SRTree ix val
cos     = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Cos
  {-# INLINE cos #-}
  tan :: SRTree ix val -> SRTree ix val
tan     = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Tan
  {-# INLINE tan #-}
  asin :: SRTree ix val -> SRTree ix val
asin    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ASin
  {-# INLINE asin #-}
  acos :: SRTree ix val -> SRTree ix val
acos    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ACos
  {-# INLINE acos #-}
  atan :: SRTree ix val -> SRTree ix val
atan    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ATan
  {-# INLINE atan #-}
  sinh :: SRTree ix val -> SRTree ix val
sinh    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Sinh
  {-# INLINE sinh #-}
  cosh :: SRTree ix val -> SRTree ix val
cosh    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Cosh
  {-# INLINE cosh #-}
  tanh :: SRTree ix val -> SRTree ix val
tanh    = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
Tanh
  {-# INLINE tanh #-}
  asinh :: SRTree ix val -> SRTree ix val
asinh   = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ASinh
  {-# INLINE asinh #-}
  acosh :: SRTree ix val -> SRTree ix val
acosh   = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ACosh
  {-# INLINE acosh #-}
  atanh :: SRTree ix val -> SRTree ix val
atanh   = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
ATanh
  {-# INLINE atanh #-}

  SRTree ix val
0 ** :: SRTree ix val -> SRTree ix val -> SRTree ix val
** SRTree ix val
r  = SRTree ix val
0
  SRTree ix val
1 ** SRTree ix val
r  = SRTree ix val
1
  SRTree ix val
l ** SRTree ix val
0  = SRTree ix val
1
  SRTree ix val
l ** SRTree ix val
1  = SRTree ix val
l
  SRTree ix val
l ** SRTree ix val
r  = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power SRTree ix val
l SRTree ix val
r
  {-# INLINE (**) #-}

  logBase :: SRTree ix val -> SRTree ix val -> SRTree ix val
logBase SRTree ix val
1 SRTree ix val
r = SRTree ix val
0
  logBase SRTree ix val
l SRTree ix val
r = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase SRTree ix val
l SRTree ix val
r
  {-# INLINE logBase #-}

instance Bifunctor SRTree where
  first :: forall a b c. (a -> b) -> SRTree a c -> SRTree b c
first a -> b
f SRTree a c
Empty         = forall ix val. SRTree ix val
Empty
  first a -> b
f (Var a
ix)      = forall ix val. ix -> SRTree ix val
Var forall a b. (a -> b) -> a -> b
$ a -> b
f a
ix
  first a -> b
f (Param a
ix)    = forall ix val. ix -> SRTree ix val
Param forall a b. (a -> b) -> a -> b
$ a -> b
f a
ix
  first a -> b
f (Fun Function
g SRTree a c
t)     = forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g forall a b. (a -> b) -> a -> b
$ forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
t
  first a -> b
f (Pow SRTree a c
t Int
k)     = forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
t) Int
k
  first a -> b
f (Add SRTree a c
l SRTree a c
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  first a -> b
f (Sub SRTree a c
l SRTree a c
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  first a -> b
f (Mul SRTree a c
l SRTree a c
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  first a -> b
f (Div SRTree a c
l SRTree a c
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  first a -> b
f (Power SRTree a c
l SRTree a c
r)   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  first a -> b
f (LogBase SRTree a c
l SRTree a c
r) = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
l) (forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first a -> b
f SRTree a c
r)
  {-# INLINE first #-}
  
  second :: forall b c a. (b -> c) -> SRTree a b -> SRTree a c
second                = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
  {-# INLINE second #-}

instance Applicative (SRTree ix) where
  pure :: forall a. a -> SRTree ix a
pure = forall ix val. val -> SRTree ix val
Const

  SRTree ix (a -> b)
Empty         <*> :: forall a b. SRTree ix (a -> b) -> SRTree ix a -> SRTree ix b
<*> SRTree ix a
t = forall ix val. SRTree ix val
Empty
  Var ix
ix        <*> SRTree ix a
t = forall ix val. ix -> SRTree ix val
Var ix
ix
  Param ix
ix      <*> SRTree ix a
t = forall ix val. ix -> SRTree ix val
Param ix
ix
  Const a -> b
f       <*> SRTree ix a
t = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f SRTree ix a
t
  Fun Function
g SRTree ix (a -> b)
tf      <*> SRTree ix a
t = forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g forall a b. (a -> b) -> a -> b
$ SRTree ix (a -> b)
tf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t
  Pow SRTree ix (a -> b)
tf Int
k      <*> SRTree ix a
t = forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow (SRTree ix (a -> b)
tf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) Int
k
  Add SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf     <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
  Sub SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf     <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
  Mul SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf     <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
  Div SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf     <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
  Power SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf   <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
  LogBase SRTree ix (a -> b)
lf SRTree ix (a -> b)
rf <*> SRTree ix a
t = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase (SRTree ix (a -> b)
lf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t) (SRTree ix (a -> b)
rf forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> SRTree ix a
t)
 
instance Foldable (SRTree ix) where
  foldMap :: forall m a. Monoid m => (a -> m) -> SRTree ix a -> m
foldMap a -> m
f SRTree ix a
Empty      = forall a. Monoid a => a
mempty
  foldMap a -> m
f (Var ix
ix)   = forall a. Monoid a => a
mempty
  foldMap a -> m
f (Param ix
ix) = forall a. Monoid a => a
mempty
  foldMap a -> m
f (Const a
c)  = a -> m
f a
c
  foldMap a -> m
f SRTree ix a
t          = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f) forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> [SRTree ix val]
getChildren SRTree ix a
t

instance Traversable (SRTree ix) where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SRTree ix a -> f (SRTree ix b)
traverse a -> f b
mf SRTree ix a
Empty         = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall ix val. SRTree ix val
Empty
  traverse a -> f b
mf (Var ix
ix)      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. ix -> SRTree ix val
Var ix
ix
  traverse a -> f b
mf (Param ix
ix)    = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. ix -> SRTree ix val
Param ix
ix
  traverse a -> f b
mf (Const a
c)     = forall ix val. val -> SRTree ix val
Const forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
mf a
c
  traverse a -> f b
mf (Fun Function
g SRTree ix a
t)     = forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
t
  traverse a -> f b
mf (Pow SRTree ix a
t Int
k)     = (forall ix val. SRTree ix val -> Int -> SRTree ix val
`Pow` Int
k) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
t
  traverse a -> f b
mf (Add SRTree ix a
l SRTree ix a
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r
  traverse a -> f b
mf (Sub SRTree ix a
l SRTree ix a
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r
  traverse a -> f b
mf (Mul SRTree ix a
l SRTree ix a
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r
  traverse a -> f b
mf (Div SRTree ix a
l SRTree ix a
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r
  traverse a -> f b
mf (Power SRTree ix a
l SRTree ix a
r)   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r
  traverse a -> f b
mf (LogBase SRTree ix a
l SRTree ix a
r) = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
mf SRTree ix a
r

-- | Same as `traverse` but for the first type parameter.
traverseIx :: Applicative f => (ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx :: forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
Empty         = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall ix val. SRTree ix val
Empty
traverseIx ixa -> f ixb
mf (Var ixa
ix)      = forall ix val. ix -> SRTree ix val
Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ixa -> f ixb
mf ixa
ix
traverseIx ixa -> f ixb
mf (Param ixa
ix)    = forall ix val. ix -> SRTree ix val
Param forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ixa -> f ixb
mf ixa
ix
traverseIx ixa -> f ixb
mf (Const val
c)     = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. val -> SRTree ix val
Const val
c
traverseIx ixa -> f ixb
mf (Fun Function
g SRTree ixa val
t)     = forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
t
traverseIx ixa -> f ixb
mf (Pow SRTree ixa val
t Int
k)     = (forall ix val. SRTree ix val -> Int -> SRTree ix val
`Pow` Int
k) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
t
traverseIx ixa -> f ixb
mf (Add SRTree ixa val
l SRTree ixa val
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
traverseIx ixa -> f ixb
mf (Sub SRTree ixa val
l SRTree ixa val
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
traverseIx ixa -> f ixb
mf (Mul SRTree ixa val
l SRTree ixa val
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
traverseIx ixa -> f ixb
mf (Div SRTree ixa val
l SRTree ixa val
r)     = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
traverseIx ixa -> f ixb
mf (Power SRTree ixa val
l SRTree ixa val
r)   = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
traverseIx ixa -> f ixb
mf (LogBase SRTree ixa val
l SRTree ixa val
r) = forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ixa -> f ixb
mf SRTree ixa val
r
{-# INLINE traverseIx #-}

-- | Arity of the current node
arity :: SRTree ix val -> Int
arity :: forall ix a. SRTree ix a -> Int
arity SRTree ix val
Empty     = Int
0
arity (Var ix
_)   = Int
0
arity (Param ix
_) = Int
0
arity (Const val
_) = Int
0
arity (Fun Function
_ SRTree ix val
_) = Int
1
arity (Pow SRTree ix val
_ Int
_) = Int
1
arity SRTree ix val
_         = Int
2
{-# INLINE arity #-}

-- | Get the children of a node. Returns an empty list in case of a leaf node.
getChildren :: SRTree ix val -> [SRTree ix val]
getChildren :: forall ix val. SRTree ix val -> [SRTree ix val]
getChildren SRTree ix val
Empty         = []
getChildren (Var ix
_)       = []
getChildren (Param ix
_)     = []
getChildren (Const val
_)     = []
getChildren (Fun Function
_ SRTree ix val
t)     = [SRTree ix val
t]
getChildren (Pow SRTree ix val
t Int
_)     = [SRTree ix val
t]
getChildren (Add SRTree ix val
l SRTree ix val
r)     = [SRTree ix val
l, SRTree ix val
r]
getChildren (Sub SRTree ix val
l SRTree ix val
r)     = [SRTree ix val
l, SRTree ix val
r]
getChildren (Mul SRTree ix val
l SRTree ix val
r)     = [SRTree ix val
l, SRTree ix val
r]
getChildren (Div SRTree ix val
l SRTree ix val
r)     = [SRTree ix val
l, SRTree ix val
r]
getChildren (Power SRTree ix val
l SRTree ix val
r)   = [SRTree ix val
l, SRTree ix val
r]
getChildren (LogBase SRTree ix val
l SRTree ix val
r) = [SRTree ix val
l, SRTree ix val
r]
{-# INLINE getChildren #-}

-- Support function to simplify operations applied to const subtrees.
evalToConst :: Floating val => SRTree ix val -> SRTree ix val  
evalToConst :: forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst (Fun Function
g (Const val
c))               = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ forall val. Floating val => Function -> val -> val
evalFun Function
g val
c
evalToConst (Power (Const val
c1) (Const val
c2))   = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c1forall a. Floating a => a -> a -> a
**val
c2
evalToConst (LogBase (Const val
c1) (Const val
c2)) = forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ forall a. Floating a => a -> a -> a
logBase val
c1 val
c2
evalToConst SRTree ix val
t                               = SRTree ix val
t
{-# INLINE evalToConst #-}

-- Support function to sum the types of nodes specified by `f`.
sumCounts :: (SRTree ix val -> Int) -> Int -> SRTree ix val -> Int
sumCounts :: forall ix val.
(SRTree ix val -> Int) -> Int -> SRTree ix val -> Int
sumCounts SRTree ix val -> Int
f Int
val = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\SRTree ix val
c Int
v -> SRTree ix val -> Int
f SRTree ix val
c forall a. Num a => a -> a -> a
+ Int
v) Int
val forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. SRTree ix val -> [SRTree ix val]
getChildren
{-# INLINE sumCounts #-}

-- | Count the number of nodes in a tree.
countNodes :: SRTree ix val -> Int
countNodes :: forall ix a. SRTree ix a -> Int
countNodes SRTree ix val
Empty = Int
0
countNodes SRTree ix val
t     = forall ix val.
(SRTree ix val -> Int) -> Int -> SRTree ix val -> Int
sumCounts forall ix a. SRTree ix a -> Int
countNodes Int
1 SRTree ix val
t
{-# INLINE countNodes #-}

-- | Count the number of `Var` nodes
countVarNodes :: SRTree ix val -> Int
countVarNodes :: forall ix a. SRTree ix a -> Int
countVarNodes (Var ix
_) = Int
1
countVarNodes SRTree ix val
t       = forall ix val.
(SRTree ix val -> Int) -> Int -> SRTree ix val -> Int
sumCounts forall ix a. SRTree ix a -> Int
countVarNodes Int
0 SRTree ix val
t
{-# INLINE countVarNodes #-}

-- | Count the occurrences of variable indexed as `ix`
countOccurrences :: Eq ix => SRTree ix val -> ix -> Int
countOccurrences :: forall ix val. Eq ix => SRTree ix val -> ix -> Int
countOccurrences (Var ix
ix) ix
iy = if ix
ixforall a. Eq a => a -> a -> Bool
==ix
iy then Int
1 else Int
0
countOccurrences SRTree ix val
t        ix
iy = forall ix val.
(SRTree ix val -> Int) -> Int -> SRTree ix val -> Int
sumCounts (forall ix val. Eq ix => SRTree ix val -> ix -> Int
`countOccurrences` ix
iy) Int
0 SRTree ix val
t
{-# INLINE countOccurrences #-}

-- | Creates an `SRTree` representing the partial derivative of the input by the variable indexed by `ix`.
deriveBy :: (Eq ix, Eq val, Floating val, OptIntPow val) => ix -> SRTree ix val -> SRTree ix val
deriveBy :: forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
_  SRTree ix val
Empty    = forall ix val. SRTree ix val
Empty
deriveBy ix
dx (Var ix
ix)
  | ix
dx forall a. Eq a => a -> a -> Bool
== ix
ix  = SRTree ix val
1
  | Bool
otherwise = SRTree ix val
0
deriveBy ix
dx (Param ix
ix) = SRTree ix val
0
deriveBy ix
dx (Const val
val) = SRTree ix val
0
deriveBy ix
dx (Fun Function
g SRTree ix val
t)   =
  case forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
t of
    SRTree ix val
0  -> SRTree ix val
0
    SRTree ix val
1  -> forall ix val.
(Eq ix, Eq val, Floating val) =>
Function -> SRTree ix val -> SRTree ix val
derivative Function
g SRTree ix val
t
    SRTree ix val
t' -> forall ix val.
(Eq ix, Eq val, Floating val) =>
Function -> SRTree ix val -> SRTree ix val
derivative Function
g SRTree ix val
t forall a. Num a => a -> a -> a
* SRTree ix val
t'
deriveBy ix
dx (Pow SRTree ix val
t Int
0)   = SRTree ix val
0    
deriveBy ix
dx (Pow SRTree ix val
t Int
1)   = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
t
deriveBy ix
dx (Pow SRTree ix val
t Int
k)   = 
  case forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
t of
    SRTree ix val
0 -> SRTree ix val
0
    Const val
val -> forall ix val. val -> SRTree ix val
Const (val
val forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) forall a. Num a => a -> a -> a
* (SRTree ix val
t forall a. OptIntPow a => a -> Int -> a
^. (Int
kforall a. Num a => a -> a -> a
-Int
1))
    SRTree ix val
t'        -> forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k forall a. Num a => a -> a -> a
* (SRTree ix val
t forall a. OptIntPow a => a -> Int -> a
^. (Int
kforall a. Num a => a -> a -> a
-Int
1)) forall a. Num a => a -> a -> a
* SRTree ix val
t'
deriveBy ix
dx (Add SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
+ forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
r
deriveBy ix
dx (Sub SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
- forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
r
deriveBy ix
dx (Mul SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
* SRTree ix val
r forall a. Num a => a -> a -> a
+ SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
r
deriveBy ix
dx (Div SRTree ix val
l SRTree ix val
r)     = (forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
* SRTree ix val
r forall a. Num a => a -> a -> a
- SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
r) forall a. Fractional a => a -> a -> a
/ SRTree ix val
r forall a. OptIntPow a => a -> Int -> a
^. Int
2
deriveBy ix
dx (Power SRTree ix val
l SRTree ix val
r)   = SRTree ix val
l forall a. Floating a => a -> a -> a
** (SRTree ix val
rforall a. Num a => a -> a -> a
-SRTree ix val
1) forall a. Num a => a -> a -> a
* (SRTree ix val
r forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
+ SRTree ix val
l forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx SRTree ix val
r)
deriveBy ix
dx (LogBase SRTree ix val
l SRTree ix val
r) = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveBy ix
dx (forall a. Floating a => a -> a
log SRTree ix val
l forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a -> a
log SRTree ix val
r)
{-# INLINE deriveBy #-}

-- | Creates an `SRTree` representing the partial derivative of the input by the parameter indexed by `ix`.
deriveParamBy :: (Eq ix, Eq val, Floating val, OptIntPow val) => ix -> SRTree ix val -> SRTree ix val
deriveParamBy :: forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
_  SRTree ix val
Empty    = forall ix val. SRTree ix val
Empty
deriveParamBy ix
dx (Var ix
ix) = SRTree ix val
0
deriveParamBy ix
dx (Param ix
ix)
  | ix
dx forall a. Eq a => a -> a -> Bool
== ix
ix  = SRTree ix val
1
  | Bool
otherwise = SRTree ix val
0
deriveParamBy ix
dx (Const val
val) = SRTree ix val
0
deriveParamBy ix
dx (Fun Function
g SRTree ix val
t)   =
  case forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
t of
    SRTree ix val
0  -> SRTree ix val
0
    SRTree ix val
1  -> forall ix val.
(Eq ix, Eq val, Floating val) =>
Function -> SRTree ix val -> SRTree ix val
derivative Function
g SRTree ix val
t
    SRTree ix val
t' -> forall ix val.
(Eq ix, Eq val, Floating val) =>
Function -> SRTree ix val -> SRTree ix val
derivative Function
g SRTree ix val
t forall a. Num a => a -> a -> a
* SRTree ix val
t'
deriveParamBy ix
dx (Pow SRTree ix val
t Int
0)   = SRTree ix val
0    
deriveParamBy ix
dx (Pow SRTree ix val
t Int
1)   = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
t
deriveParamBy ix
dx (Pow SRTree ix val
t Int
k)   = 
  case forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
t of
    SRTree ix val
0 -> SRTree ix val
0
    Const val
val -> forall ix val. val -> SRTree ix val
Const (val
val forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) forall a. Num a => a -> a -> a
* (SRTree ix val
t forall a. OptIntPow a => a -> Int -> a
^. (Int
kforall a. Num a => a -> a -> a
-Int
1))
    SRTree ix val
t'        -> forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k forall a. Num a => a -> a -> a
* (SRTree ix val
t forall a. OptIntPow a => a -> Int -> a
^. (Int
kforall a. Num a => a -> a -> a
-Int
1)) forall a. Num a => a -> a -> a
* SRTree ix val
t'
deriveParamBy ix
dx (Add SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
+ forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
r
deriveParamBy ix
dx (Sub SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
- forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
r
deriveParamBy ix
dx (Mul SRTree ix val
l SRTree ix val
r)     = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
* SRTree ix val
r forall a. Num a => a -> a -> a
+ SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
r
deriveParamBy ix
dx (Div SRTree ix val
l SRTree ix val
r)     = (forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
* SRTree ix val
r forall a. Num a => a -> a -> a
- SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
r) forall a. Fractional a => a -> a -> a
/ SRTree ix val
r forall a. OptIntPow a => a -> Int -> a
^. Int
2
deriveParamBy ix
dx (Power SRTree ix val
l SRTree ix val
r)   = SRTree ix val
l forall a. Floating a => a -> a -> a
** (SRTree ix val
rforall a. Num a => a -> a -> a
-SRTree ix val
1) forall a. Num a => a -> a -> a
* (SRTree ix val
r forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
l forall a. Num a => a -> a -> a
+ SRTree ix val
l forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log SRTree ix val
l forall a. Num a => a -> a -> a
* forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx SRTree ix val
r)
deriveParamBy ix
dx (LogBase SRTree ix val
l SRTree ix val
r) = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
ix -> SRTree ix val -> SRTree ix val
deriveParamBy ix
dx (forall a. Floating a => a -> a
log SRTree ix val
l forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a -> a
log SRTree ix val
r)
{-# INLINE deriveParamBy #-}

-- | Simplifies the `SRTree`.
simplify :: (Eq ix, Eq val, Floating val, OptIntPow val) => SRTree ix val -> SRTree ix val
simplify :: forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify (Fun Function
g SRTree ix val
t) = forall val ix. Floating val => SRTree ix val -> SRTree ix val
evalToConst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g forall a b. (a -> b) -> a -> b
$ forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
t
simplify (Pow SRTree ix val
t Int
0) = SRTree ix val
1    
simplify (Pow SRTree ix val
t Int
1) = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
t
simplify (Pow SRTree ix val
t Int
k) =
  case forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
t of
    Const val
c -> forall ix val. val -> SRTree ix val
Const forall a b. (a -> b) -> a -> b
$ val
c forall a. OptIntPow a => a -> Int -> a
^. Int
k
    SRTree ix val
t'      -> forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow SRTree ix val
t' Int
k
    
simplify (Add SRTree ix val
l SRTree ix val
r)
  | SRTree ix val
l' forall a. Eq a => a -> a -> Bool
== SRTree ix val
r' = SRTree ix val
2 forall a. Num a => a -> a -> a
* SRTree ix val
l' 
  | Bool
otherwise = SRTree ix val
l' forall a. Num a => a -> a -> a
+ SRTree ix val
r' 
  where 
      l' :: SRTree ix val
l' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l 
      r' :: SRTree ix val
r' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r
simplify (Sub SRTree ix val
l SRTree ix val
r)
  | SRTree ix val
l' forall a. Eq a => a -> a -> Bool
== SRTree ix val
r' = SRTree ix val
0
  | Bool
otherwise = SRTree ix val
l' forall a. Num a => a -> a -> a
- SRTree ix val
r' 
  where 
      l' :: SRTree ix val
l' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l 
      r' :: SRTree ix val
r' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r
simplify (Mul SRTree ix val
l SRTree ix val
r)
  | SRTree ix val
l' forall a. Eq a => a -> a -> Bool
== SRTree ix val
r'  = forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow SRTree ix val
l' Int
2
  | Bool
otherwise = SRTree ix val
l' forall a. Num a => a -> a -> a
* SRTree ix val
r' 
  where 
      l' :: SRTree ix val
l' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l 
      r' :: SRTree ix val
r' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r
simplify (Div SRTree ix val
l SRTree ix val
r)
  | SRTree ix val
l' forall a. Eq a => a -> a -> Bool
== SRTree ix val
r'  = SRTree ix val
1
  | Bool
otherwise = SRTree ix val
l' forall a. Fractional a => a -> a -> a
/ SRTree ix val
r' 
  where 
      l' :: SRTree ix val
l' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l 
      r' :: SRTree ix val
r' = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r

simplify (Power SRTree ix val
l SRTree ix val
r)   = forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l forall a. Floating a => a -> a -> a
** forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r
simplify (LogBase SRTree ix val
l SRTree ix val
r) = forall a. Floating a => a -> a -> a
logBase (forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
l) (forall ix val.
(Eq ix, Eq val, Floating val, OptIntPow val) =>
SRTree ix val -> SRTree ix val
simplify SRTree ix val
r)
simplify SRTree ix val
t             = SRTree ix val
t
{-# INLINE simplify #-}

-- | Derivative of a Function
derivative :: (Eq ix, Eq val, Floating val) => Function -> SRTree ix val -> SRTree ix val
derivative :: forall ix val.
(Eq ix, Eq val, Floating val) =>
Function -> SRTree ix val -> SRTree ix val
derivative Function
Id      = forall a b. a -> b -> a
const SRTree ix val
1
derivative Function
Abs     = \SRTree ix val
x -> SRTree ix val
x forall a. Fractional a => a -> a -> a
/ forall a. Num a => a -> a
abs SRTree ix val
x
derivative Function
Sin     = forall a. Floating a => a -> a
cos
derivative Function
Cos     = forall a. Num a => a -> a
negateforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall a. Floating a => a -> a
sin
derivative Function
Tan     = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Floating a => a -> a -> a
**SRTree ix val
2.0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
cos
derivative Function
Sinh    = forall a. Floating a => a -> a
cosh
derivative Function
Cosh    = forall a. Floating a => a -> a
sinh
derivative Function
Tanh    = (SRTree ix val
1forall a. Num a => a -> a -> a
-) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Floating a => a -> a -> a
**SRTree ix val
2.0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
tanh
derivative Function
ASin    = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
1forall a. Num a => a -> a -> a
-) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
ACos    = forall a. Num a => a -> a
negate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
1forall a. Num a => a -> a -> a
-) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
ATan    = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
1forall a. Num a => a -> a -> a
+) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
ASinh   = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
1forall a. Num a => a -> a -> a
+) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
ACosh   = \SRTree ix val
x -> SRTree ix val
1 forall a. Fractional a => a -> a -> a
/ (forall a. Floating a => a -> a
sqrt (SRTree ix val
xforall a. Num a => a -> a -> a
-SRTree ix val
1) forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
sqrt (SRTree ix val
xforall a. Num a => a -> a -> a
+SRTree ix val
1))
derivative Function
ATanh   = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
1forall a. Num a => a -> a -> a
-) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
Sqrt    = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
2forall a. Num a => a -> a -> a
*) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt
derivative Function
Cbrt    = forall a. Fractional a => a -> a
recip forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SRTree ix val
3forall a. Num a => a -> a -> a
*) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
cbrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
derivative Function
Square  = (SRTree ix val
2forall a. Num a => a -> a -> a
*)
derivative Function
Exp     = forall a. Floating a => a -> a
exp
derivative Function
Log     = forall a. Fractional a => a -> a
recip
{-# INLINE derivative #-}

-- | Evaluates a function.
evalFun :: Floating val => Function -> val -> val
evalFun :: forall val. Floating val => Function -> val -> val
evalFun Function
Id      = forall a. a -> a
id
evalFun Function
Abs     = forall a. Num a => a -> a
abs
evalFun Function
Sin     = forall a. Floating a => a -> a
sin
evalFun Function
Cos     = forall a. Floating a => a -> a
cos
evalFun Function
Tan     = forall a. Floating a => a -> a
tan
evalFun Function
Sinh    = forall a. Floating a => a -> a
sinh
evalFun Function
Cosh    = forall a. Floating a => a -> a
cosh
evalFun Function
Tanh    = forall a. Floating a => a -> a
tanh
evalFun Function
ASin    = forall a. Floating a => a -> a
asin
evalFun Function
ACos    = forall a. Floating a => a -> a
acos
evalFun Function
ATan    = forall a. Floating a => a -> a
atan
evalFun Function
ASinh   = forall a. Floating a => a -> a
asinh
evalFun Function
ACosh   = forall a. Floating a => a -> a
acosh
evalFun Function
ATanh   = forall a. Floating a => a -> a
atanh
evalFun Function
Sqrt    = forall a. Floating a => a -> a
sqrt
evalFun Function
Cbrt    = forall a. Floating a => a -> a
cbrt
evalFun Function
Square  = (forall a b. (Num a, Integral b) => a -> b -> a
^Integer
2)
evalFun Function
Exp     = forall a. Floating a => a -> a
exp
evalFun Function
Log     = forall a. Floating a => a -> a
log
{-# INLINE evalFun #-}

cbrt :: Floating val => val -> val
cbrt :: forall a. Floating a => a -> a
cbrt val
x = forall a. Num a => a -> a
signum val
x forall a. Num a => a -> a -> a
* forall a. Num a => a -> a
abs val
x forall a. Floating a => a -> a -> a
** (val
1forall a. Fractional a => a -> a -> a
/val
3)
{-# INLINE cbrt #-}

-- | Returns the inverse of a function. This is a partial function.
inverseFunc :: Function -> Function
inverseFunc :: Function -> Function
inverseFunc Function
Id     = Function
Id
inverseFunc Function
Sin    = Function
ASin
inverseFunc Function
Cos    = Function
ACos
inverseFunc Function
Tan    = Function
ATan
inverseFunc Function
Tanh   = Function
ATanh
inverseFunc Function
ASin   = Function
Sin
inverseFunc Function
ACos   = Function
Cos
inverseFunc Function
ATan   = Function
Tan
inverseFunc Function
ATanh  = Function
Tanh
inverseFunc Function
Sqrt   = Function
Square
inverseFunc Function
Square = Function
Sqrt
inverseFunc Function
Log    = Function
Exp
inverseFunc Function
Exp    = Function
Log
inverseFunc Function
x      = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show Function
x forall a. [a] -> [a] -> [a]
++ String
" has no support for inverse function"
{-# INLINE inverseFunc #-}

-- | Evaluates a tree with the variables stored in a `Reader` monad.
evalTree :: (Floating val, OptIntPow val) => SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree :: forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
Empty         = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
evalTree (Const val
c)     = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just val
c
evalTree (Var ix
ix)      = forall x a. x -> Reader (x -> a) a
askAbout ix
ix
evalTree (Param ix
ix)    = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just val
1.0 -- TODO: askAbout paramIx
evalTree (Fun Function
f SRTree ix val
t)     = forall val. Floating val => Function -> val -> val
evalFun Function
f forall (f :: * -> *) (g :: * -> *) a b.
(Applicative f, Applicative g) =>
(a -> b) -> f (g a) -> f (g b)
<$$> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
t
evalTree (Pow SRTree ix val
t Int
k)     = (forall a. OptIntPow a => a -> Int -> a
^. Int
k) forall (f :: * -> *) (g :: * -> *) a b.
(Applicative f, Applicative g) =>
(a -> b) -> f (g a) -> f (g b)
<$$> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
t
evalTree (Add SRTree ix val
l SRTree ix val
r)     = forall a. Num a => a -> a -> a
(+)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r
evalTree (Sub SRTree ix val
l SRTree ix val
r)     = (-)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r
evalTree (Mul SRTree ix val
l SRTree ix val
r)     = forall a. Num a => a -> a -> a
(*)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r
evalTree (Div SRTree ix val
l SRTree ix val
r)     = forall a. Fractional a => a -> a -> a
(/)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r
evalTree (Power SRTree ix val
l SRTree ix val
r)   = forall a. Floating a => a -> a -> a
(**) forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r
evalTree (LogBase SRTree ix val
l SRTree ix val
r) = forall a. Floating a => a -> a -> a
logBase forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
r

-- | Evaluates a tree with the variables stored in a `Reader` monad while mapping the constant 
-- values to a different type.
evalTreeMap :: (Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) => (v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap :: forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
Empty         = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
evalTreeMap v1 -> v2
f (Const v1
c)     = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ v1 -> v2
f v1
c
evalTreeMap v1 -> v2
f (Var ix
ix)      = forall x a. x -> Reader (x -> a) a
askAbout ix
ix
evalTreeMap v1 -> v2
f (Param ix
ix)    = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ v1 -> v2
f v1
1.0 -- TODO: askAbout paramIx
evalTreeMap v1 -> v2
f (Fun Function
g SRTree ix v1
t)     = forall val. Floating val => Function -> val -> val
evalFun Function
g forall (f :: * -> *) (g :: * -> *) a b.
(Applicative f, Applicative g) =>
(a -> b) -> f (g a) -> f (g b)
<$$> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
t
evalTreeMap v1 -> v2
f (Pow SRTree ix v1
t Int
k)     = (forall a. OptIntPow a => a -> Int -> a
^. Int
k) forall (f :: * -> *) (g :: * -> *) a b.
(Applicative f, Applicative g) =>
(a -> b) -> f (g a) -> f (g b)
<$$> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
t
evalTreeMap v1 -> v2
f (Add SRTree ix v1
l SRTree ix v1
r)     = forall a. Num a => a -> a -> a
(+)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r
evalTreeMap v1 -> v2
f (Sub SRTree ix v1
l SRTree ix v1
r)     = (-)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r
evalTreeMap v1 -> v2
f (Mul SRTree ix v1
l SRTree ix v1
r)     = forall a. Num a => a -> a -> a
(*)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r
evalTreeMap v1 -> v2
f (Div SRTree ix v1
l SRTree ix v1
r)     = forall a. Fractional a => a -> a -> a
(/)  forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r
evalTreeMap v1 -> v2
f (Power SRTree ix v1
l SRTree ix v1
r)   = forall a. Floating a => a -> a -> a
(**) forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r
evalTreeMap v1 -> v2
f (LogBase SRTree ix v1
l SRTree ix v1
r) = forall a. Floating a => a -> a -> a
logBase forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
l forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall v1 v2 ix.
(Floating v1, OptIntPow v1, Floating v2, OptIntPow v2) =>
(v1 -> v2) -> SRTree ix v1 -> Reader (ix -> Maybe v2) (Maybe v2)
evalTreeMap v1 -> v2
f SRTree ix v1
r

-- lift functions inside nested applicatives.
(<$$>) :: (Applicative f, Applicative g) => (a -> b) -> f (g a) -> f (g b)
<$$> :: forall (f :: * -> *) (g :: * -> *) a b.
(Applicative f, Applicative g) =>
(a -> b) -> f (g a) -> f (g b)
(<$$>) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
{-# INLINE (<$$>) #-}
(<$*>) :: (Applicative f, Applicative g) => (a -> b -> c) -> f (g a) -> f (g b -> g c)
a -> b -> c
op <$*> :: forall (f :: * -> *) (g :: * -> *) a b c.
(Applicative f, Applicative g) =>
(a -> b -> c) -> f (g a) -> f (g b -> g c)
<$*> f (g a)
m = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> c
op forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (g a)
m
{-# INLINE (<$*>) #-}

-- applies the argument `x` in the function carried by the `Reader` monad.
askAbout :: x -> Reader (x -> a) a
askAbout :: forall x a. x -> Reader (x -> a) a
askAbout x
x = forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (forall a b. (a -> b) -> a -> b
$ x
x)
{-# INLINE askAbout #-}

-- | Example of using `evalTree` with a Map.
evalTreeWithMap :: (Ord ix, Floating val, OptIntPow val) => SRTree ix val -> Map ix val -> Maybe val
evalTreeWithMap :: forall ix val.
(Ord ix, Floating val, OptIntPow val) =>
SRTree ix val -> Map ix val -> Maybe val
evalTreeWithMap SRTree ix val
t Map ix val
m = forall r a. Reader r a -> r -> a
runReader (forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree ix val
t) (Map ix val
m forall k a. Ord k => Map k a -> k -> Maybe a
!?)
{-# INLINE evalTreeWithMap #-}

-- | Example of using `evalTree` with a Vector.
evalTreeWithVector :: (Floating val, OptIntPow val) => SRTree Int val -> V.Vector val -> Maybe val
evalTreeWithVector :: forall val.
(Floating val, OptIntPow val) =>
SRTree Int val -> Vector val -> Maybe val
evalTreeWithVector SRTree Int val
t Vector val
v = forall r a. Reader r a -> r -> a
runReader (forall val ix.
(Floating val, OptIntPow val) =>
SRTree ix val -> Reader (ix -> Maybe val) (Maybe val)
evalTree SRTree Int val
t) (Vector val
v forall a. Vector a -> Int -> Maybe a
V.!?)
{-# INLINE evalTreeWithVector #-}

-- | Relabel occurences of a var into a tuple (ix, Int).
relabelOccurrences :: forall ix val . Ord ix => SRTree ix val -> SRTree (ix, Int) val
relabelOccurrences :: forall ix val. Ord ix => SRTree ix val -> SRTree (ix, Int) val
relabelOccurrences SRTree ix val
t = forall (f :: * -> *) ixa ixb val.
Applicative f =>
(ixa -> f ixb) -> SRTree ixa val -> f (SRTree ixb val)
traverseIx ix -> State (Map ix Int) (ix, Int)
updVar SRTree ix val
t forall s a. State s a -> s -> a
`evalState` forall k a. Map k a
M.empty 
  where
    updVar :: ix -> State (Map ix Int) (ix, Int)
    updVar :: ix -> State (Map ix Int) (ix, Int)
updVar ix
ix = do
      Map ix Int
s <- forall s (m :: * -> *). MonadState s m => m s
get
      case Map ix Int
s forall k a. Ord k => Map k a -> k -> Maybe a
!? ix
ix of
        Maybe Int
Nothing -> do forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
insert ix
ix Int
0 Map ix Int
s
                      forall (f :: * -> *) a. Applicative f => a -> f a
pure (ix
ix, Int
0)
        Just Int
c  -> do forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
insert ix
ix (Int
cforall a. Num a => a -> a -> a
+Int
1) Map ix Int
s
                      forall (f :: * -> *) a. Applicative f => a -> f a
pure (ix
ix, Int
cforall a. Num a => a -> a -> a
+Int
1)

-- | Relabel the parameters sequentially starting from 0
relabelParams :: Num ix => SRTree ix val -> SRTree ix val
relabelParams :: forall ix val. Num ix => SRTree ix val -> SRTree ix val
relabelParams SRTree ix val
t = (forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
t) forall s a. State s a -> s -> a
`evalState` ix
0
  where
    toState :: Num ix => SRTree ix val -> State ix (SRTree ix val)
    toState :: forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState (Param ix
x) = do ix
n <- forall s (m :: * -> *). MonadState s m => m s
get; forall s (m :: * -> *). MonadState s m => s -> m ()
put (ix
nforall a. Num a => a -> a -> a
+ix
1); forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. ix -> SRTree ix val
Param ix
n)
    toState (Add SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add SRTree ix val
l' SRTree ix val
r')
    toState (Sub SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub SRTree ix val
l' SRTree ix val
r')
    toState (Mul SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul SRTree ix val
l' SRTree ix val
r')
    toState (Div SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div SRTree ix val
l' SRTree ix val
r')
    toState (Power SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power SRTree ix val
l' SRTree ix val
r')
    toState (LogBase SRTree ix val
l SRTree ix val
r) = do SRTree ix val
l' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
l; SRTree ix val
r' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
r; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase SRTree ix val
l' SRTree ix val
r')
    toState (Fun Function
f SRTree ix val
n) = do SRTree ix val
n' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
n; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
f SRTree ix val
n')
    toState (Pow SRTree ix val
n Int
i) = do SRTree ix val
n' <- forall ix val. Num ix => SRTree ix val -> State ix (SRTree ix val)
toState SRTree ix val
n; forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow SRTree ix val
n' Int
i)
    toState SRTree ix val
n = forall (f :: * -> *) a. Applicative f => a -> f a
pure SRTree ix val
n