module Math.MFSolve
(
MFSolver, SimpleExpr(..), Expr, LinExpr(..), UnaryOp(..), BinaryOp(..),
Dependencies, DepError(..), SimpleVar(..),
evalSimple, evalExpr, fromSimple, toSimple, makeVariable,
makeConstant, hasVar,
getKnown, knownVars, varDefined, nonlinearEqs, dependendVars,
noDeps,
eliminate, addEquation,
dependencies, getValue, getKnownM, varDefinedM, eliminateM,
(=&=), (===), ignore,
runSolver, evalSolver, execSolver, unsafeSolve, showVars)
where
import qualified Data.HashMap.Strict as M
import qualified Data.HashSet as H
import GHC.Generics
import Control.Monad.Except
import Control.Monad.State
import Control.Exception
import Data.Typeable
import Control.Applicative hiding (Const)
import Data.Hashable
import Data.Maybe
import Data.List
import Data.Function(on)
data BinaryOp = Add | Mul
deriving Eq
data UnaryOp =
Sin | Abs | Recip | Signum |
Exp | Log | Cos | Cosh | Atanh |
Tan | Sinh | Asin | Acos | Asinh | Acosh | Atan
deriving (Eq, Generic)
data SimpleExpr v n =
SEBin BinaryOp (SimpleExpr v n) (SimpleExpr v n) |
SEUn UnaryOp (SimpleExpr v n) |
Var v |
Const n
newtype SimpleVar = SimpleVar String
deriving (Eq, Ord, Generic)
data Expr v n = Expr (LinExpr v n) [TrigTerm v n] [NonLinExpr v n]
deriving (Generic)
data LinExpr v n = LinExpr n [(v, n)]
deriving (Generic, Eq, Show)
type Period v n = [(v, n)]
type Phase n = n
type Amplitude v n = LinExpr v n
type TrigTerm v n = (Period v n, [(Phase n, Amplitude v n)])
data NonLinExpr v n =
UnaryApp UnaryOp (Expr v n) |
MulExp (Expr v n) (Expr v n) |
SinExp (Expr v n)
deriving Generic
type LinearMap v n = M.HashMap v (LinExpr v n)
type TrigEq v n = (Period v n, Amplitude v n, Phase n, n)
type TrigEq2 v n = M.HashMap (Period v n)
(M.HashMap v (Expr v n))
pattern LinearE l = Expr l [] []
pattern ConstE c = Expr (LinExpr c []) [] []
pattern LConst c = LinExpr c []
instance (Hashable v, Hashable n) => Hashable (LinExpr v n)
instance (Hashable v, Hashable n) => Hashable (NonLinExpr v n)
instance Hashable UnaryOp
instance (Hashable v, Hashable n) => Hashable (Expr v n)
instance Hashable SimpleVar
instance Show SimpleVar where
show (SimpleVar s) = s
data Dependencies v n = Dependencies
(M.HashMap v (H.HashSet v))
(LinearMap v n)
[TrigEq v n]
(TrigEq2 v n)
[Expr v n]
data DepError v n =
UndefinedVar v |
UnknownVar v n |
InconsistentEq n |
RedundantEq
deriving Typeable
instance (Show v, Show n, Typeable v, Typeable n) => Exception (DepError v n)
instance (Ord n, Num n, Eq n, Show v, Show n) => Show (Expr v n) where
show e = show (toSimple e)
newtype MFSolver v n a = MFSolver {
runSolver :: Dependencies v n -> Either (DepError v n) (Dependencies v n, a) }
instance Monad (MFSolver v n) where
MFSolver s >>= f = MFSolver $ \dep ->
case s dep of
Left err -> Left err
Right (dep2, val) ->
runSolver (f val) dep2
return a = MFSolver $ \dep -> Right (dep, a)
instance (MonadState (Dependencies v n)) (MFSolver v n) where
get = MFSolver $ \dep -> Right (dep, dep)
put dep = MFSolver $ const $ Right (dep, ())
instance (MonadError (DepError v n)) (MFSolver v n) where
throwError e = MFSolver $ const $ Left e
catchError s h = MFSolver $ \dep ->
case runSolver s dep of
Left e -> runSolver (h e) dep
v -> v
instance Functor (MFSolver v n) where
fmap f s = MFSolver $ \dep ->
fmap f <$> runSolver s dep
instance Applicative (MFSolver v n) where
(<*>) = ap
pure = return
withParens :: (Show t1, Show t, Ord t1, Num t1, Eq t1) => SimpleExpr t t1 -> [BinaryOp] -> String
withParens e@(SEBin op _ _) ops
| op `elem` ops = "(" ++ show e ++ ")"
withParens e _ = show e
instance (Show v, Ord n, Show n, Num n, Eq n) => Show (SimpleExpr v n) where
show (Var v) = show v
show (Const n) = show n
show (SEBin Add e1 (SEBin Mul (Const e2) e3))
| e2 < 0 =
show e1 ++ " - " ++ show (SEBin Mul (Const (negate e2)) e3)
show (SEBin Add e1 e2) =
show e1 ++ " + " ++ show e2
show (SEBin Mul (Const 1) e) = show e
show (SEBin Mul e (Const 1)) = show e
show (SEBin Mul e1 (SEUn Recip e2)) =
withParens e1 [Add] ++ "/" ++ withParens e2 [Add, Mul]
show (SEBin Mul e1 e2) =
withParens e1 [Add] ++ "*" ++ withParens e2 [Add]
show (SEUn Exp (SEBin Mul (SEUn Log e1) e2)) =
withParens e1 [Add, Mul] ++ "**" ++ withParens e2 [Add, Mul]
show (SEUn op e) = show op ++ "(" ++ show e ++ ")"
instance Show BinaryOp where
show Add = "+"
show Mul = "*"
instance Show UnaryOp where
show Sin = "sin"
show Abs = "abs"
show Recip = "1/"
show Signum = "sign"
show Exp = "exp"
show Log = "log"
show Cos = "cos"
show Cosh = "cosh"
show Atanh = "atanh"
show Tan = "tan"
show Sinh = "sinh"
show Asin = "asin"
show Acos = "acos"
show Asinh = "asinh"
show Acosh = "acosh"
show Atan = "atan"
instance (Floating n, Ord n, Ord v) => Num (Expr v n) where
(+) = addExpr
(*) = mulExpr
negate = mulExpr (ConstE (1))
abs = unExpr Abs
signum = unExpr Signum
fromInteger = ConstE . fromInteger
instance (Floating n, Ord n, Ord v) => Fractional (Expr v n) where
recip = unExpr Recip
fromRational = ConstE . fromRational
instance (Floating n, Ord n, Ord v) => Floating (Expr v n) where
pi = ConstE pi
exp = unExpr Exp
log = unExpr Log
sin = sinExpr
cos a = sinExpr (a + ConstE (pi/2))
cosh = unExpr Cosh
atanh = unExpr Atanh
tan = unExpr Tan
sinh = unExpr Sinh
asin = unExpr Asin
acos = unExpr Acos
asinh = unExpr Asinh
acosh = unExpr Acosh
atan = unExpr Atan
instance (Show n, Floating n, Ord n, Ord v, Show v) =>Show (Dependencies v n) where
show dep@(Dependencies _ lin _ _ _) =
unlines (map showLin (M.toList lin) ++
map showNl (nonlinearEqs dep))
where showLin (v, e) = show v ++ " = " ++ show (LinearE e)
showNl e = show e ++ " = 0"
instance (Show n, Show v) => Show (DepError v n) where
show (InconsistentEq a) =
"Inconsistent equations, off by " ++ show a
show RedundantEq =
"Redundant Equation."
show (UndefinedVar v) =
error ("Variable is undefined: " ++ show v)
show (UnknownVar v n) =
error ("Value of variable not known: " ++ show v ++ " = " ++ show n)
addSimple :: (Num t1, Eq t1) => SimpleExpr t t1 -> SimpleExpr t t1 -> SimpleExpr t t1
addSimple (Const 0) e = e
addSimple e (Const 0) = e
addSimple e1 e2 = SEBin Add e1 e2
seHasVar :: Eq v => v -> SimpleExpr v t -> Bool
seHasVar v1 (Var v2) = v1 == v2
seHasVar _ (Const _) = False
seHasVar v (SEBin _ e1 e2) =
seHasVar v e1 ||
seHasVar v e2
seHasVar v (SEUn _ e) = seHasVar v e
hasVar :: (Num t, Eq v, Eq t) => v -> Expr v t -> Bool
hasVar v = seHasVar v . toSimple
linToSimple :: (Num n, Eq n) => LinExpr v n -> SimpleExpr v n
linToSimple (LinExpr v t) =
Const v `addSimple`
foldr (addSimple.mul) (Const 0) t
where
mul (v2, 1) = Var v2
mul (v2, c) = SEBin Mul (Const c) (Var v2)
trigToSimple :: (Num n, Eq n) => TrigTerm v n -> SimpleExpr v n
trigToSimple (theta, t) =
foldr (addSimple.makeSin) (Const 0) t
where
makeSin (alpha, n) =
SEBin Mul (linToSimple n)
(SEUn Sin angle) where
angle = linToSimple (LinExpr alpha theta)
nonlinToSimple :: (Num n, Eq n) => NonLinExpr v n -> SimpleExpr v n
nonlinToSimple (UnaryApp o e) =
SEUn o (toSimple e)
nonlinToSimple (MulExp e1 e2) =
SEBin Mul (toSimple e1) (toSimple e2)
nonlinToSimple (SinExp e) =
SEUn Sin (toSimple e)
toSimple :: (Num n, Eq n) => Expr v n -> SimpleExpr v n
toSimple (Expr lin trig nonlin) =
linToSimple lin `addSimple`
foldr (addSimple.trigToSimple)
(Const 0) trig `addSimple`
foldr (addSimple.nonlinToSimple)
(Const 0) nonlin
evalBin :: (Floating n) => BinaryOp -> n -> n -> n
evalBin Add = (+)
evalBin Mul = (*)
evalUn :: (Floating n) => UnaryOp -> n -> n
evalUn Sin = sin
evalUn Abs = abs
evalUn Recip = recip
evalUn Signum = signum
evalUn Exp = exp
evalUn Log = log
evalUn Cos = cos
evalUn Cosh = cosh
evalUn Atanh = atanh
evalUn Tan = tan
evalUn Sinh = sinh
evalUn Asin = asin
evalUn Acos = acos
evalUn Asinh = asinh
evalUn Acosh = acosh
evalUn Atan = atan
evalSimple :: Floating m => (n -> m) -> (v -> m) -> SimpleExpr v n -> m
evalSimple _ s (Var v) = s v
evalSimple g _ (Const c) = g c
evalSimple g s (SEBin f e1 e2) =
evalBin f (evalSimple g s e1) (evalSimple g s e2)
evalSimple g s (SEUn f e) =
evalUn f (evalSimple g s e)
fromSimple :: (Floating n, Ord n, Ord v) => SimpleExpr v n -> Expr v n
fromSimple = evalSimple makeConstant makeVariable
evalExpr :: (Floating n) => (v -> n) -> SimpleExpr v n -> n
evalExpr = evalSimple id
zeroTerm :: (Num n) => LinExpr v n
zeroTerm = LConst 0
makeConstant :: n -> Expr v n
makeConstant = ConstE
makeVariable :: Num n => v -> Expr v n
makeVariable v = LinearE $ LinExpr 0 [(v, 1)]
trigExpr :: (Num n) => [TrigTerm v n] -> Expr v n
trigExpr t = Expr zeroTerm t []
nonlinExpr :: Num n => [NonLinExpr v n] -> Expr v n
nonlinExpr = Expr zeroTerm []
isConst :: LinExpr v n -> Bool
isConst (LConst _) = True
isConst _ = False
linVars :: LinExpr v n -> [v]
linVars (LinExpr _ v) = map fst v
addLin :: (Ord v, Num n, Eq n) => LinExpr v n -> LinExpr v n -> LinExpr v n
addLin (LinExpr c1 terms1) (LinExpr c2 terms2) =
LinExpr (c1+c2) terms3 where
terms3 = filter ((/= 0) . snd) $
merge terms1 terms2 (+)
addExpr :: (Ord n, Ord v, Floating n) => Expr v n -> Expr v n -> Expr v n
addExpr (Expr lt1 trig1 nl1) (Expr lt2 trig2 nl2) =
Expr (addLin lt1 lt2) trig3 (nl1++nl2)
where
trig3 = merge trig1 trig2 addTrigTerms
merge :: Ord k => [(k, v)] -> [(k, v)] -> (v -> v -> v) -> [(k, v)]
merge [] l _ = l
merge l [] _ = l
merge (a@(k,v):as) (b@(k2,v2):bs) f = case compare k k2 of
LT -> a: merge as (b:bs) f
EQ -> (k, f v v2): merge as bs f
GT -> b: merge (a:as) bs f
addTrigTerms :: (Ord a, Ord t, Floating a) => [(a, LinExpr t a)] -> [(a, LinExpr t a)] -> [(a, LinExpr t a)]
addTrigTerms [] p = p
addTrigTerms terms terms2 =
foldr mergeTerms terms terms2
where
mergeTerms (alpha, n) ((beta, m):rest) =
case addTrigTerm alpha n beta m of
Just (_, LConst 0) -> rest
Just (gamma, o) ->
mergeTerms (gamma, o) rest
Nothing ->
(beta, m) : mergeTerms (alpha, n) rest
mergeTerms a [] = [a]
addTrigTerm :: (Ord a, Ord t, Floating a) => a -> LinExpr t a -> a -> LinExpr t a -> Maybe (a, LinExpr t a)
addTrigTerm alpha n beta m
| alpha == beta =
Just (alpha, addLin n m)
| Just r <- termIsMultiple n m =
let gamma = atan (divident/divisor) +
(if divisor < 0 then pi else 0)
divident = r*sin alpha + sin beta
divisor = r*cos alpha + cos beta
o = sqrt(divident*divident + divisor*divisor)
in Just (gamma, mulLinExpr o m)
| otherwise = Nothing
termIsMultiple :: (Ord a, Fractional a, Eq t) => LinExpr t a -> LinExpr t a -> Maybe a
termIsMultiple (LinExpr _ _) (LinExpr 0 []) = Nothing
termIsMultiple (LinExpr 0 []) (LinExpr _ _) = Nothing
termIsMultiple (LinExpr 0 r1@((_, d1):_)) (LinExpr 0 r2@((_, d2):_))
| compareBy r1 r2 (compareTerm (d1/d2)) =
Just (d1/d2)
termIsMultiple (LinExpr c1 r1) (LinExpr c2 r2)
| compareBy r1 r2 (compareTerm (c1/c2)) =
Just (c1/c2)
| otherwise = Nothing
compareTerm :: (Ord a1, Fractional a1, Eq a) => a1 -> (a, a1) -> (a, a1) -> Bool
compareTerm ratio (v3,c3) (v4, c4) =
v3 == v4 && (abs(c3 (c4 * ratio)) <= abs c3*2e-50)
compareBy :: [a] -> [b] -> (a -> b -> Bool) -> Bool
compareBy [] [] _ = True
compareBy (e:l) (e2:l2) f =
f e e2 && compareBy l l2 f
compareBy _ _ _ = False
mulLinExpr :: Num n => n -> LinExpr v n -> LinExpr v n
mulLinExpr x (LinExpr e terms) =
LinExpr (e*x) $ map (fmap (*x)) terms
mulConstTrig :: (Ord n, Num n) => n -> TrigTerm v n -> TrigTerm v n
mulConstTrig c (theta, terms) = (theta, tt) where
tt = map (fmap (mulLinExpr c)) terms
mulLinTrig :: (Ord n, Ord v, Floating n) => LinExpr v n -> TrigTerm v n -> Expr v n
mulLinTrig lt (theta, terms) =
foldr ((+).mul1) 0 terms
where
mul1 (alpha, LinExpr c []) =
trigExpr [(theta, [(alpha, mulLinExpr c lt)])]
mul1 t =
nonlinExpr [MulExp (trigExpr [(theta, [t])])
(Expr lt [] [])]
mulExpr :: (Ord a, Ord t, Floating a) => Expr t a -> Expr t a -> Expr t a
mulExpr (ConstE c) (Expr lt2 trig []) =
Expr (mulLinExpr c lt2)
(map (mulConstTrig c) trig) []
mulExpr (Expr lt2 trig []) (ConstE c) =
Expr (mulLinExpr c lt2)
(map (mulConstTrig c) trig) []
mulExpr (LinearE lt) (Expr (LConst c) trig []) =
LinearE (mulLinExpr c lt) +
foldr ((+).mulLinTrig lt) 0 trig
mulExpr (Expr (LConst c) trig []) (LinearE lt) =
LinearE (mulLinExpr c lt) +
foldr ((+).mulLinTrig lt) 0 trig
mulExpr e1 e2 = nonlinExpr [MulExp e1 e2]
sinExpr :: Floating n => Expr v n -> Expr v n
sinExpr (Expr (LinExpr c t) [] [])
| null t = ConstE (sin c)
| otherwise = trigExpr [(t, [(c, LinExpr 1 [])])]
sinExpr e = nonlinExpr [SinExp e]
unExpr :: Floating n => UnaryOp -> Expr v n -> Expr v n
unExpr f (ConstE c) = ConstE (evalUn f c)
unExpr f e = nonlinExpr [UnaryApp f e]
substVarLin :: (Ord v, Num n, Eq n) => (v -> Maybe (LinExpr v n)) -> LinExpr v n -> LinExpr v n
substVarLin s (LinExpr a terms) =
let substOne (v, c) =
maybe (LinExpr 0 [(v, c)]) (mulLinExpr c) (s v)
in foldr (addLin.substOne) (LinExpr a []) terms
substVarNonLin :: (Ord n, Ord v, Floating n) => (v -> Maybe (LinExpr v n)) -> NonLinExpr v n -> Expr v n
substVarNonLin s (UnaryApp f e1) =
unExpr f (subst s e1)
substVarNonLin s (MulExp e1 e2) =
subst s e1 * subst s e2
substVarNonLin s (SinExp e1) =
sin (subst s e1)
substVarTrig :: (Ord v, Ord n, Floating n) => (v -> Maybe (LinExpr v n)) -> ([(v, n)], [(n, LinExpr v n)]) -> Expr v n
substVarTrig s (period, terms) =
let period2 = LinearE $ substVarLin s (LinExpr 0 period)
terms2 = map (fmap $ LinearE . substVarLin s) terms
in foldr (\(p,a) -> (+ (a * sin (ConstE p + period2))))
0 terms2
subst :: (Ord n, Ord v, Floating n) => (v -> Maybe (LinExpr v n)) -> Expr v n -> Expr v n
subst s (Expr lt trig nl) =
LinearE (substVarLin s lt) +
foldr ((+).substVarTrig s) 0 trig +
foldr ((+).substVarNonLin s) 0 nl
noDeps :: Dependencies v n
noDeps = Dependencies M.empty M.empty [] M.empty []
simpleSubst :: Eq a => a -> b -> a -> Maybe b
simpleSubst x y z
| x == z = Just y
| otherwise = Nothing
trigToExpr :: (Ord n, Ord v, Floating n) => TrigEq v n -> Expr v n
trigToExpr (p, a, ph, c) =
LinearE a * sin(LinearE $ LinExpr ph p) +
ConstE c
trig2ToExpr :: (Ord v, Floating n, Ord n) => M.HashMap v (Expr v n) -> [Expr v n]
trig2ToExpr =
map (\(v,e)-> makeVariable ve)
. M.toList
addEqs :: (Hashable v, Hashable n, RealFrac (Phase n), Ord v, Floating n) => Dependencies v n -> [Expr v n] -> Either (DepError v n) (Dependencies v n)
addEqs = foldM addEquation
addEquation :: (Hashable n, Hashable v, RealFrac (Phase n), Ord v,
Floating n) =>
Dependencies v n
-> Expr v n -> Either (DepError v n) (Dependencies v n)
addEquation deps@(Dependencies _ lin _ _ _) expr =
addEq0 deps $
subst (flip M.lookup lin) expr
select :: [a] -> [(a, [a])]
select [] = []
select (x:xs) =
(x,xs) : [(y,x:ys) | (y,ys) <- select xs]
substDep :: (Hashable v, Ord v, Num n, Eq n) =>
M.HashMap v (H.HashSet v) -> M.HashMap v (LinExpr v n)
-> v -> LinExpr v n -> Bool
-> (M.HashMap v (H.HashSet v), LinearMap v n)
substDep vdep lin v lt insertp =
let depVars = fromMaybe H.empty (M.lookup v vdep)
lin' = (if insertp then M.insert v lt
else id) $
H.foldl' (flip $ M.adjust $
substVarLin $
simpleSubst v lt)
lin depVars
depVars2 | insertp = H.insert v depVars
| otherwise = depVars
tryUnion k m1 m2 =
let xs = H.intersection m1 m2
hasvar v2 = case M.lookup v2 lin' of
Nothing -> False
Just (LinExpr _ vs) ->
any ((==k).fst) vs
in H.filter hasvar xs
`H.union` H.difference m1 xs
`H.union` H.difference m2 xs
vdep' = H.foldl'
(\mp k -> M.insertWith (tryUnion k) k depVars2 mp)
(M.delete v vdep)
(H.fromList $ linVars lt)
in (vdep', lin')
addEq0 :: (Hashable v, Hashable n, RealFrac (Phase n), Ord v, Floating n) => Dependencies v n -> Expr v n -> Either (DepError v n) (Dependencies v n)
addEq0 _ (ConstE c) =
Left $ if c == 0 then RedundantEq
else InconsistentEq c
addEq0 (Dependencies vdep lin trig trig2 nonlin) (Expr lt [] []) =
let (v, _, lt2) = splitMax lt
(vdep', lin') = substDep vdep lin v lt2 True
trig' = map trigToExpr trig
trig2' = concatMap trig2ToExpr $ M.elems trig2
in addEqs (Dependencies vdep' lin' [] M.empty []) (trig'++trig2'++nonlin)
addEq0 deps@(Dependencies vdep lin trig trig2 nl)
(Expr (LinExpr c lt) [(theta, [(alpha, LConst n)])] []) =
if null lt then
addEq0 deps (LinearE $ LinExpr (alpha asin (c/n)) theta)
else
case M.lookup theta trig2 of
Nothing -> addSin (LinExpr c lt) alpha n
Just map2 ->
case foldr ((+).doSubst)
(ConstE c +
ConstE n *
sin (LinearE $ LinExpr alpha theta))
lt of
Expr lt2 [(_, [(alpha2, LConst n2)])] []
| not $ isConst lt2
-> addSin lt2 alpha2 n2
e2 -> addEq0 deps e2
where
doSubst (v,c2) = case M.lookup v map2 of
Nothing -> makeVariable v * ConstE c2
Just e2 -> e2 * ConstE c2
where
addSin l' a' n' =
let (v, c', r) = splitMax l'
trig2' = M.insertWith M.union theta
(M.singleton v $
Expr r [(theta, [(a', LinExpr (n'/negate c') [])])] [])
trig2
in Right $ Dependencies vdep lin trig trig2' nl
addEq0 (Dependencies d lin [] trig2 nl)
(Expr (LConst c) [(theta, [(alpha, n)])] []) =
Right $ Dependencies d lin [(theta, n, alpha, c)] trig2 nl
addEq0 (Dependencies deps lin trig trig2 nl)
(Expr (LConst x) [(theta, [(a, n)])] []) =
case mapMaybe similarTrig $ select trig of
[] -> Right $ Dependencies deps lin ((theta, n, a, x):trig) trig2 nl
l -> addEqs (Dependencies deps lin rest trig2 nl) [lin1, lin2]
where
((b,y), rest) = maximumBy (maxTrig `on` fst) l
maxTrig (t1,_) (t2,_) =
compare ((t1a)`dmod`pi) ((t2a)`dmod`pi)
d = sin(ab)
e = y*cos(ab)x
theta2 = atan (y*d/e)b +
(if (d*e) < 0 then pi else 0)
n2 = sqrt(y*y + e*e/(d*d))
lin1 = LinearE $ LinExpr (theta2) theta
lin2 = LinearE n ConstE n2
where
similarTrig ((t,m,b,y),rest)
| Just r <- termIsMultiple m n,
t == theta &&
(ba) `dmod` pi > pi/8 =
Just ((b,y/r),rest)
| otherwise = Nothing
addEq0 (Dependencies d lin trig trig2 nonlin) e =
Right $ Dependencies d lin trig trig2 (e:nonlin)
deleteDep :: (Hashable k, Hashable b, Eq k, Eq b) =>
M.HashMap b (H.HashSet k)
-> M.HashMap k (LinExpr b n) -> k
-> Maybe (M.HashMap b (H.HashSet k), M.HashMap k (LinExpr b n), LinExpr b n)
deleteDep vdep lin v =
case M.lookup v lin of
Nothing -> Nothing
Just lt -> Just (vdep', lin', lt)
where
lin' = M.delete v lin
vdep' = H.foldl'
(flip $ M.adjust $ H.delete v)
vdep (H.fromList $ linVars lt)
eliminate :: (Hashable n, Show n, Hashable v, RealFrac (Phase n), Ord v, Show v,
Floating n) => Dependencies v n -> v -> (Dependencies v n, [Expr v n])
eliminate (Dependencies vdep lin trig trig2 nonlin) v
| Just (vdep', lin', lt) <- deleteDep vdep lin v =
(Dependencies vdep' lin' trig trig2 nonlin,
[LinearE lt makeVariable v])
| Just vars <- M.lookup v vdep,
(v2:_) <- H.toList vars =
case deleteDep vdep lin v2 of
Nothing ->
error $ "Internal error: found empty dependency on " ++ show v2
Just (vdep', lin', lt) ->
let lt2 = snd $ reArrange v2 lt v
(vdep'', lin'') = substDep vdep' lin' v lt2 False
trig' = map trigToExpr trig
trig2' = concatMap trig2ToExpr $ M.elems trig2
deps = Dependencies vdep'' lin'' [] M.empty []
e = [LinearE lt2 makeVariable v]
in case foldM addEq0
deps $
map (subst $ simpleSubst v lt2)
(trig'++trig2'++nonlin) of
Left _ -> (deps, e)
Right d -> (d, e)
| otherwise =
let (l, trig2') =
M.foldrWithKey trigFold
([], M.empty) trig2
trigFold p t (l2, m2) =
let (l3, m1) = elimTrig p t v
mp | M.null m1 = m2
| otherwise = M.insert p m1 m2
in (l3++l2, mp)
(nlWith, nlWithout) =
partition (hasVar v) $
map trigToExpr trig ++ nonlin
deps = Dependencies vdep lin [] trig2' []
in case foldM addEq0 deps
nlWithout of
Left _ -> (deps, nlWith++l)
Right d -> (d, nlWith++l)
reArrange :: (Show v, Ord v, Fractional n, Eq n) =>
v -> LinExpr v n -> v -> (n, LinExpr v n)
reArrange v2 (LinExpr c vars) v =
case find ((==v).fst.fst) $ select vars of
Nothing ->
error $ "Internal error: variable " ++ show v ++
" not in linear expression "
Just ((_,c2), r) ->
(c2,
LinExpr (c/negate c2) r
`addLin` LinExpr 0 [(v2, 1/c2)])
reArrangeTrig :: (Show v, Ord t1, Ord v, Floating t1) => v -> Expr v t1 -> v -> Expr v t1
reArrangeTrig v2 (Expr lt trig _) v =
let (c2, lt2) = reArrange v2 lt v
in LinearE lt2 trigExpr trig / ConstE c2
elimTrig :: (Show v, Ord v, Hashable v, Floating n, Ord n) =>
Period v n -> M.HashMap v (Expr v n) -> v
-> ([Expr v n], M.HashMap v (Expr v n))
elimTrig p m v
| any ((==v).fst) p =
(trig2ToExpr m, M.empty)
| Just e <- M.lookup v m =
([makeVariable v e],
M.delete v m)
| Just (v2, e) <-
find (hasVar v.snd) $ M.toList m =
let e2 = reArrangeTrig v2 e v
substOne (v3, c)
| v == v3 = e2 * ConstE c
| otherwise = makeVariable v3 * ConstE c
doSubst (Expr (LinExpr c lt) trig _) =
foldr ((+).substOne)
(ConstE c + trigExpr trig) lt
in ([makeVariable v e],
M.map doSubst $ M.delete v2 m)
| otherwise =
([], m)
dmod :: RealFrac a => a -> a -> a
dmod a b = abs((a/b) fromInteger (round (a/b)) * b)
splitMax :: (Ord b, Fractional b, Eq v) => LinExpr v b -> (v, b, LinExpr v b)
splitMax (LinExpr c t) =
let ((v,c2),r) = maximumBy (compare `on` (abs.snd.fst)) $
select t
in (v, c2,
LinExpr (c/c2) $
map (fmap (negate.(/c2))) r)
varDefined :: (Eq v, Hashable v) => v -> Dependencies v n -> Bool
varDefined v (Dependencies _ dep _ _ _) =
case M.lookup v dep of
Nothing -> False
_ -> True
dependendVars :: (Eq n) => Dependencies v n -> [(v, LinExpr v n)]
dependendVars (Dependencies _ lin _ _ _) =
filter (notConst.snd) (M.toList lin)
where
notConst (LinExpr _ []) = False
notConst _ = True
knownVars :: Dependencies v n -> [(v, n)]
knownVars (Dependencies _ lin _ _ _) =
mapMaybe knownVar $ M.toList lin
where
knownVar (v, LinExpr n []) = Just (v, n)
knownVar _ = Nothing
getKnown :: (Eq v, Hashable v) => v -> Dependencies v n -> Either [v] n
getKnown var (Dependencies _ lin _ _ _) =
case M.lookup var lin of
Nothing -> Left []
Just (LinExpr a []) ->
Right a
Just (LinExpr _ v) ->
Left $ map fst v
nonlinearEqs :: (Ord n, Ord v, Floating n) => Dependencies v n -> [Expr v n]
nonlinearEqs (Dependencies _ _ trig trig2 nl) =
map trigToExpr trig ++
map (\(v, e) -> makeVariable v e)
(concatMap M.toList (M.elems trig2)) ++
nl
showVars :: (Show n, Show v, Ord n, Ord v, Floating n) => Either (DepError v n) (Dependencies v n) -> IO ()
showVars (Left e) = print e
showVars (Right dep) = print dep
dependencies :: (MonadState (Dependencies v n) m) => m (Dependencies v n)
dependencies = get
getValue :: (MonadState (Dependencies v n) m,
MonadError (DepError v n) m,
Eq v, Hashable v) =>
v -> m n
getValue v = do
v2 <- getKnownM v
case v2 of
Right e -> return e
Left _ -> throwError $ UndefinedVar v
varDefinedM :: (MonadState (Dependencies v n) m, Hashable v, Eq v) =>
v -> m Bool
varDefinedM v = varDefined v `liftM` dependencies
getKnownM :: (MonadState (Dependencies v n) m, Hashable v, Eq v) =>
v -> m (Either [v] n)
getKnownM v = getKnown v `liftM` dependencies
eliminateM :: (MonadState (Dependencies v n) m, Hashable n, Hashable v,
Show n, Show v, RealFrac n, Ord v, Floating n) =>
v -> m [Expr v n]
eliminateM v = do
dep <- dependencies
let (dep2, e) = eliminate dep v
put dep2
return e
infixr 1 === , =&=
(===) :: (MonadState (Dependencies v n) m,
MonadError (DepError v n) m,
Eq v, Hashable v, Hashable n,
RealFrac n, Floating n, Ord v) =>
Expr v n -> Expr v n -> m ()
(===) lhs rhs = do
deps <- dependencies
case addEquation deps (lhs rhs) of
Left e -> throwError e
Right dep -> put dep
(=&=) :: (MonadState (Dependencies v n) m,
MonadError (DepError v n) m,
Eq v, Hashable v, Hashable n,
RealFrac n, Floating n, Ord v) =>
(Expr v n, Expr v n) -> (Expr v n, Expr v n) -> m ()
(=&=) (a, b) (c, d) =
do ignore $ a === c
b === d
ignore :: MonadError (DepError v n) m => m () -> m ()
ignore m = m `catchError` (
\e -> case e of
RedundantEq -> return ()
_ -> throwError e)
unsafeSolve :: (Typeable n, Typeable v, Show n, Show v) => MFSolver v n a -> Dependencies v n -> a
unsafeSolve s dep = case runSolver s dep of
Right (_, v) -> v
Left e -> throw e
evalSolver :: MFSolver v n a -> Dependencies v n -> Either (DepError v n) a
evalSolver s dep = snd <$> runSolver s dep
execSolver :: MFSolver v n a -> Dependencies v n ->
Either (DepError v n) (Dependencies v n)
execSolver s dep = fst <$> runSolver s dep