module Math.LinearRecursive.Monad
(
VectorLike(..)
, LinearCombination
, Variable
, (<+>)
, (<->)
, (<*)
, (*>)
, zeroVector
, Polynomial
, P.x
, LinearRecursive
, newVariable
, newVariables
, (<:-)
, (<+-)
, runLinearRecursive
, simulateLinearRecursive
, getConstant
, getPartialSum
, getStep
, getPowerOf
, getPolynomial
, getPartialSumWith
) where
import Control.Monad (zipWithM_)
import Control.Applicative ((<$>))
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Math.LinearRecursive.Internal.Vector
import Math.LinearRecursive.Internal.Matrix
import Math.LinearRecursive.Internal.Polynomial hiding (fromList, toList, x)
import qualified Math.LinearRecursive.Internal.Polynomial as P
type LinearCombination = Vector
type Variable = Vector1
data LRVariable a = LRV { initialValue :: a, dependency :: LinearCombination a }
dmap :: Num a => (LinearCombination a -> LinearCombination a) -> LRVariable a -> LRVariable a
dmap f (LRV val dep) = LRV val (f dep)
type LRVariables a = IntMap (LRVariable a)
data LinearRecursive a b = LR { unLR :: Int -> (b, Int, LRVariables a -> LRVariables a) }
instance Num a => Functor (LinearRecursive a) where
fmap f m = m >>= return . f
instance Num a => Monad (LinearRecursive a) where
return a = LR (const (a, 0, id))
a >>= b = LR $ \v -> let (ra, nva, ma) = unLR a v
(rb, nvb, mb) = unLR (b ra) (v + nva)
in
(rb, nva + nvb, mb . ma)
newVariable :: Num a => a -> LinearRecursive a (Variable a)
newVariable val0 = LR $ \v -> (vector1 v, 1, IntMap.insert v variable)
where
variable = LRV { initialValue = val0, dependency = zeroVector }
newVariables :: Num a => [a] -> LinearRecursive a [Variable a]
newVariables vals = do
ret <- mapM newVariable vals
zipWithM_ (<:-) (tail ret) ret
return ret
getConstant :: Num a => a -> LinearRecursive a (LinearCombination a)
getConstant val = do
one <- newVariable 1
one <:- one
return (toVector one *> val)
getPartialSum :: Num a => LinearCombination a -> LinearRecursive a (LinearCombination a)
getPartialSum val = do
s <- newVariable 0
s <:- s <+> val
return (toVector s)
getStep :: Num a => LinearRecursive a (LinearCombination a)
getStep = getConstant 1 >>= getPartialSum
getPowerOf :: Num a => a -> LinearRecursive a (LinearCombination a)
getPowerOf a = do
prod <- newVariable 1
prod <:- prod *> a
return (toVector prod)
inverseTrans :: Num a => [Polynomial a] -> Matrix a
inverseTrans polys = inverseMatrixDiag1 ma
where
n = length polys
ma = matrix [[vcomponent (unPoly polyi) j | j <- [0..n1]] | polyi <- polys]
getPartialSumWith :: (Num a, VectorLike v) => Polynomial a -> v a -> LinearRecursive a (LinearCombination a)
getPartialSumWith poly v
| n < 0 = return zeroVector
| otherwise = do
basisValue <- go (toVector v) 0
let vars = map (foldl (<+>) zeroVector . zipWith (*>) basisValue) trans
return $ foldl (<+>) zeroVector [ powi *> coeffi
| (i, powi) <- zip [0..] vars
, let coeffi = vcomponent vec i
]
where
n = degree poly
basisPoly = scanl (*) 1 [P.x + fromIntegral i | i <- [1..n]]
go prev pos | pos > n = return []
| otherwise = do
next <- (*> fromIntegral (pos `max` 1)) . (<+> prev) <$> getPartialSum prev
(next:) <$> go next (pos + 1)
trans = unMatrix (inverseTrans basisPoly)
vec = unPoly poly
getPolynomial :: Num a => Polynomial a -> LinearRecursive a (LinearCombination a)
getPolynomial poly = newVariable 1 >>= getPartialSumWith poly
(<+-) :: (Num a, VectorLike v) => Variable a -> v a -> LinearRecursive a ()
(<+-) var dep = LR (const ((), 0, IntMap.adjust (dmap (<+>toVector dep)) (unVector1 var)))
(<:-) :: (Num a, VectorLike v) => Variable a -> v a -> LinearRecursive a ()
(<:-) var dep = LR (const ((), 0, IntMap.adjust (dmap (const (toVector dep))) (unVector1 var)))
infix 1 <:-,<+-
buildMatrix :: Num a => LRVariables a -> (Matrix a, Matrix a)
buildMatrix mapping = (matrix trans, matrix $ map (: []) initValues)
where
initValues = map initialValue (IntMap.elems mapping)
rawDep = map (unVector'.dependency) (IntMap.elems mapping)
varCount = length initValues
trans = map (\m -> [IntMap.findWithDefault 0 i m | i <- [0..varCount1]]) rawDep
runLinearRecursive :: (Num a, Integral b, VectorLike v) => LinearRecursive a (v a) -> b -> a
runLinearRecursive _ steps | steps < 0 = error "runLinearRecursive: steps must be non-negative"
runLinearRecursive m steps = sum [head (res !! i) * ai | (i, ai) <- IntMap.assocs (unVector' target)]
where
(target, _, g) = unLR m 0
dep = g IntMap.empty
(trans, initCol) = buildMatrix dep
res = unMatrix' (trans^steps * initCol)
simulateLinearRecursive :: (Num a, VectorLike v) => LinearRecursive a (v a) -> [a]
simulateLinearRecursive m = map (\res -> sum [head (res !! i) * ai | (i, ai) <- IntMap.assocs (unVector' target)]) cols
where
(target, _, g) = unLR m 0
dep = g IntMap.empty
(trans, initCol) = buildMatrix dep
cols = map unMatrix' $ scanl (flip (*)) initCol (repeat trans)