module ToySolver.Arith.LPUtil
( toStandardForm
, toStandardForm'
) where
import Control.Exception
import Control.Monad
import Control.Monad.State
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Maybe
import Data.VectorSpace
import qualified Data.Interval as Interval
import ToySolver.Data.OrdRel
import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar
import qualified ToySolver.Arith.BoundsInference as BI
toStandardForm
:: (LA.Expr Rational, [OrdRel (LA.Expr Rational)])
-> ( (LA.Expr Rational, [(LA.Expr Rational, Rational)])
, Model Rational -> Model Rational
)
toStandardForm :: (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
Model Rational -> Model Rational)
toStandardForm prob1 :: (Expr Rational, [OrdRel (Expr Rational)])
prob1@(Expr Rational
obj, [OrdRel (Expr Rational)]
cs) = ((Expr Rational, [(Expr Rational, Rational)])
prob2, Model Rational -> Model Rational
mt)
where
vs :: IntSet
vs = forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` forall a. Variables a => a -> IntSet
vars [OrdRel (Expr Rational)]
cs
((Expr Rational, [(Expr Rational, Rational)])
prob2,VarMap (Expr Rational)
s) = (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
VarMap (Expr Rational))
toStandardForm' (Expr Rational, [OrdRel (Expr Rational)])
prob1
mt :: Model Rational -> Model Rational
mt :: Model Rational -> Model Rational
mt Model Rational
m = forall a. [(Var, a)] -> IntMap a
IM.fromAscList forall a b. (a -> b) -> a -> b
$ do
Var
v <- IntSet -> [Var]
IS.toAscList IntSet
vs
case forall a. Var -> IntMap a -> Maybe a
IM.lookup Var
v VarMap (Expr Rational)
s of
Just Expr Rational
def -> forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, forall m e v. Eval m e v => m -> e -> v
LA.eval Model Rational
m Expr Rational
def)
Maybe (Expr Rational)
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, Model Rational
m forall a. IntMap a -> Var -> a
IM.! Var
v)
type M = State Var
toStandardForm'
:: (LA.Expr Rational, [OrdRel (LA.Expr Rational)])
-> ( (LA.Expr Rational, [(LA.Expr Rational, Rational)])
, VarMap (LA.Expr Rational)
)
toStandardForm' :: (Expr Rational, [OrdRel (Expr Rational)])
-> ((Expr Rational, [(Expr Rational, Rational)]),
VarMap (Expr Rational))
toStandardForm' (Expr Rational
obj, [OrdRel (Expr Rational)]
cs) = ((Expr Rational, [(Expr Rational, Rational)]),
VarMap (Expr Rational))
m
where
vs :: IntSet
vs = forall a. Variables a => a -> IntSet
vars Expr Rational
obj IntSet -> IntSet -> IntSet
`IS.union` forall a. Variables a => a -> IntSet
vars [OrdRel (Expr Rational)]
cs
v1 :: Var
v1 = if IntSet -> Bool
IS.null IntSet
vs then Var
0 else IntSet -> Var
IS.findMax IntSet
vs forall a. Num a => a -> a -> a
+ Var
1
initialBounds :: IntMap (Interval Rational)
initialBounds = forall a. [(Var, a)] -> IntMap a
IM.fromList [(Var
v, forall r. Ord r => Interval r
Interval.whole) | Var
v <- IntSet -> [Var]
IS.toList IntSet
vs]
bounds :: IntMap (Interval Rational)
bounds = forall r.
RealFrac r =>
BoundsEnv r -> [Atom r] -> IntSet -> Var -> BoundsEnv r
BI.inferBounds IntMap (Interval Rational)
initialBounds [OrdRel (Expr Rational)]
cs IntSet
IS.empty Var
10
gensym :: M Var
gensym :: M Var
gensym = do
Var
v <- forall s (m :: * -> *). MonadState s m => m s
get
forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ Var
vforall a. Num a => a -> a -> a
+Var
1
forall (m :: * -> *) a. Monad m => a -> m a
return Var
v
m :: ((Expr Rational, [(Expr Rational, Rational)]),
VarMap (Expr Rational))
m = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
evalState Var
v1 forall a b. (a -> b) -> a -> b
$ do
VarMap (Expr Rational)
s <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall (f :: * -> *) a. Foldable f => f (IntMap a) -> IntMap a
IM.unions forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (forall a. IntMap a -> [(Var, a)]
IM.toList IntMap (Interval Rational)
bounds) forall a b. (a -> b) -> a -> b
$ \(Var
v,Interval Rational
i) -> do
case forall r. Interval r -> Extended r
Interval.lowerBound Interval Rational
i of
Extended Rational
Interval.NegInf -> do
Var
v1 <- M Var
gensym
Var
v2 <- M Var
gensym
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Var -> a -> IntMap a
IM.singleton Var
v (forall r. Num r => Var -> Expr r
LA.var Var
v1 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. Num r => Var -> Expr r
LA.var Var
v2)
Interval.Finite Rational
lb
| Rational
lb forall a. Ord a => a -> a -> Bool
>= Rational
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. IntMap a
IM.empty
| Bool
otherwise -> do
Var
v1 <- M Var
gensym
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Var -> a -> IntMap a
IM.singleton Var
v (forall r. Num r => Var -> Expr r
LA.var Var
v1 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. (Num r, Eq r) => r -> Expr r
LA.constant Rational
lb)
let obj2 :: Expr Rational
obj2 = forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst VarMap (Expr Rational)
s Expr Rational
obj
[(Expr Rational, Rational)]
cs2 <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [OrdRel (Expr Rational)]
cs forall a b. (a -> b) -> a -> b
$ \(OrdRel Expr Rational
lhs RelOp
op Expr Rational
rhs) -> do
case forall r. Num r => Var -> Expr r -> (r, Expr r)
LA.extract Var
LA.unitVar (forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst VarMap (Expr Rational)
s (Expr Rational
lhs forall v. AdditiveGroup v => v -> v -> v
^-^ Expr Rational
rhs)) of
(Rational
c,Expr Rational
e) -> do
let (Expr Rational
lhs2,RelOp
op2,Rational
rhs2) =
if -Rational
c forall a. Ord a => a -> a -> Bool
>= Rational
0
then (Expr Rational
e,RelOp
op,-Rational
c)
else (forall v. AdditiveGroup v => v -> v
negateV Expr Rational
e, RelOp -> RelOp
flipOp RelOp
op, Rational
c)
case RelOp
op2 of
RelOp
Eql -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2,Rational
rhs2)
RelOp
Le -> do
Var
v <- M Var
gensym
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2 forall v. AdditiveGroup v => v -> v -> v
^+^ forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
RelOp
Ge -> do
case forall r. Expr r -> [(r, Var)]
LA.terms Expr Rational
lhs2 of
[(Rational
1,Var
_)] | Rational
rhs2forall a. Ord a => a -> a -> Bool
<=Rational
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
[(Rational, Var)]
_ -> do
Var
v <- M Var
gensym
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Expr Rational
lhs2 forall v. AdditiveGroup v => v -> v -> v
^-^ forall r. Num r => Var -> Expr r
LA.var Var
v, Rational
rhs2)
RelOp
_ -> forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"ToySolver.LPUtil.toStandardForm: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show RelOp
op2 forall a. [a] -> [a] -> [a]
++ [Char]
" is not supported"
forall a. HasCallStack => Bool -> a -> a
assert (forall (t :: * -> *). Foldable t => t Bool -> Bool
and [forall a. Maybe a -> Bool
isNothing forall a b. (a -> b) -> a -> b
$ forall r. Num r => Var -> Expr r -> Maybe r
LA.lookupCoeff Var
LA.unitVar Expr Rational
c | (Expr Rational
c,Rational
_) <- [(Expr Rational, Rational)]
cs2]) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return ()
forall (m :: * -> *) a. Monad m => a -> m a
return ((Expr Rational
obj2,[(Expr Rational, Rational)]
cs2),VarMap (Expr Rational)
s)