msakai / toysolver / 513

Committed 24 Nov 2024 08:49AM UTC coverage: 68.675% (-1.0%) from 69.681%

Build # 513

Build Type

push

github

Committed by

web-flow

Commit Message

Merge 496ed4263 into 46e1509c4

Run Details

5 of 121 new or added lines in 7 files covered. (4.13%)

105 existing lines in 14 files now uncovered.

9745 of 14190 relevant lines covered (68.68%)

0.69 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

70.37

/src/ToySolver/Data/LA.hs

{-# OPTIONS_HADDOCK show-extensions #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Data.LA
-- Copyright   :  (c) Masahiro Sakai 2011
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- Some definition for Theory of Linear Arithmetics.
--
-----------------------------------------------------------------------------
module ToySolver.Data.LA
  (
  -- * Expression of linear arithmetics
    Expr
  , Var

  -- ** Conversion
  , var
  , constant
  , terms
  , fromTerms
  , coeffMap
  , fromCoeffMap
  , unitVar
  , asConst

  -- ** Query
  , coeff
  , lookupCoeff
  , extract
  , extractMaybe

  -- ** Operations
  , mapCoeff
  , mapCoeffWithVar
  , evalLinear
  , lift1
  , applySubst
  , applySubst1
  , showExpr

  -- * Atomic formula of linear arithmetics
  , Atom (..)
  , showAtom
  , applySubstAtom
  , applySubst1Atom
  , solveFor
  , module ToySolver.Data.OrdRel

  -- * Evaluation of expression and atomic formula
  , Eval (..)

  -- * misc
  , BoundsEnv
  , computeInterval
  ) where

import Control.Monad
import Control.DeepSeq
import Data.List
import Data.Maybe
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.IntSet as IntSet
import Data.Interval
import Data.VectorSpace

import qualified ToySolver.Data.OrdRel as OrdRel
import ToySolver.Data.OrdRel
import ToySolver.Data.IntVar

-----------------------------------------------------------------------------

-- | Linear combination of variables and constants.
-- Non-negative keys are used for variables's coefficients.
-- key 'unitVar' is used for constants.
newtype Expr r
  = Expr
  { -- | a mapping from variables to coefficients
    coeffMap :: IntMap r
  } deriving (Eq, Ord)

-- | Create a @Expr@ from a mapping from variables to coefficients.
fromCoeffMap :: (Num r, Eq r) => IntMap r -> Expr r
fromCoeffMap m = normalizeExpr (Expr m)

-- | terms contained in the expression.
terms :: Expr r -> [(r,Var)]
terms (Expr m) = [(c,v) | (v,c) <- IntMap.toList m]

-- | Create a @Expr@ from a list of terms.
fromTerms :: (Num r, Eq r) => [(r,Var)] -> Expr r
fromTerms ts = fromCoeffMap $ IntMap.fromListWith (+) [(x,c) | (c,x) <- ts]

instance Variables (Expr r) where
  vars (Expr m) = IntSet.delete unitVar (IntMap.keysSet m)

instance Show r => Show (Expr r) where
  showsPrec d m  = showParen (d > 10) $
    showString "fromTerms " . shows (terms m)

instance (Num r, Eq r, Read r) => Read (Expr r) where
  readsPrec p = readParen (p > 10) $ \ r -> do
    ("fromTerms",s) <- lex r
    (xs,t) <- reads s
    return (fromTerms xs, t)

instance NFData r => NFData (Expr r) where
  rnf (Expr m) = rnf m

-- | Special variable that should always be evaluated to 1.
unitVar :: Var
unitVar = -1

asConst :: Num r => Expr r -> Maybe r
asConst (Expr m) =
  case IntMap.toList m of
    [] -> Just 0
    [(v,x)] | v==unitVar -> Just x
    _ -> Nothing

normalizeExpr :: (Num r, Eq r) => Expr r -> Expr r
normalizeExpr (Expr t) = Expr $ IntMap.filter (0/=) t

-- | variable
var :: Num r => Var -> Expr r
var v = Expr $ IntMap.singleton v 1

-- | constant
constant :: (Num r, Eq r) => r -> Expr r
constant c = normalizeExpr $ Expr $ IntMap.singleton unitVar c

-- | map coefficients.
mapCoeff :: (Num b, Eq b) => (a -> b) -> Expr a -> Expr b
mapCoeff f (Expr t) = Expr $ IntMap.mapMaybe g t
  where
    g c = if c' == 0 then Nothing else Just c'
      where c' = f c

-- | map coefficients.
mapCoeffWithVar :: (Num b, Eq b) => (a -> Var -> b) -> Expr a -> Expr b
mapCoeffWithVar f (Expr t) = Expr $ IntMap.mapMaybeWithKey g t
  where
    g v c = if c' == 0 then Nothing else Just c'
      where c' = f c v

instance (Num r, Eq r) => AdditiveGroup (Expr r) where
  Expr t ^+^ e2 | IntMap.null t = e2
  e1 ^+^ Expr t | IntMap.null t = e1
  e1 ^+^ e2 = normalizeExpr $ plus e1 e2
  zeroV = Expr $ IntMap.empty
  negateV = ((-1) *^)

instance (Num r, Eq r) => VectorSpace (Expr r) where
  type Scalar (Expr r) = r
  1 *^ e = e
  0 *^ e = zeroV
  c *^ e = mapCoeff (c*) e

plus :: Num r => Expr r -> Expr r -> Expr r
plus (Expr t1) (Expr t2) = Expr $ IntMap.unionWith (+) t1 t2

instance Num r => Eval (Model r) (Expr r) r where
  eval m (Expr t) = sum [(m' IntMap.! v) * c | (v,c) <- IntMap.toList t]
    where m' = IntMap.insert unitVar 1 m

-- | evaluate the expression under the model.
evalLinear :: VectorSpace a => Model a -> a -> Expr (Scalar a) -> a
evalLinear m u (Expr t) = sumV [c *^ (m' IntMap.! v) | (v,c) <- IntMap.toList t]
  where m' = IntMap.insert unitVar u m

lift1 :: VectorSpace x => x -> (Var -> x) -> Expr (Scalar x) -> x
lift1 unit f (Expr t) = sumV [c *^ (g v) | (v,c) <- IntMap.toList t]
  where
    g v
      | v==unitVar = unit
      | otherwise   = f v

applySubst :: (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
applySubst s (Expr m) = sumV (map f (IntMap.toList m))
  where
    f (v,c) = c *^ (
      case IntMap.lookup v s of
        Just tm -> tm
        Nothing -> var v)

-- | applySubst1 x e e1 == e1[e/x]
applySubst1 :: (Num r, Eq r) => Var -> Expr r -> Expr r -> Expr r
applySubst1 x e e1 =
  case extractMaybe x e1 of
    Nothing -> e1
    Just (c,e2) -> c *^ e ^+^ e2

-- | lookup a coefficient of the variable.
-- @
--   coeff v e == fst (extract v e)
-- @
coeff :: Num r => Var -> Expr r -> r
coeff v (Expr m) = IntMap.findWithDefault 0 v m

-- | lookup a coefficient of the variable.
-- It returns @Nothing@ if the expression does not contain @v@.
-- @
--   lookupCoeff v e == fmap fst (extractMaybe v e)
-- @
lookupCoeff :: Num r => Var -> Expr r -> Maybe r
lookupCoeff v (Expr m) = IntMap.lookup v m

-- | @extract v e@ returns @(c, e')@ such that @e == c *^ v ^+^ e'@
extract :: Num r => Var -> Expr r -> (r, Expr r)
extract v (Expr m) = (IntMap.findWithDefault 0 v m, Expr (IntMap.delete v m))
{-
-- Alternative implementation which may be faster but allocte more memory
extract v (Expr m) =
  case IntMap.updateLookupWithKey (\_ _ -> Nothing) v m of
    (Nothing, _) -> (0, Expr m)
    (Just c, m2) -> (c, Expr m2)
-}

-- | @extractMaybe v e@ returns @Just (c, e')@ such that @e == c *^ v ^+^ e'@
-- if @e@ contains v, and returns @Nothing@ otherwise.
extractMaybe :: Num r => Var -> Expr r -> Maybe (r, Expr r)
extractMaybe v (Expr m) =
  case IntMap.lookup v m of
    Nothing -> Nothing
    Just c -> Just (c, Expr (IntMap.delete v m))
{-
-- Alternative implementation which may be faster but allocte more memory
extractMaybe v (Expr m) =
  case IntMap.updateLookupWithKey (\_ _ -> Nothing) v m of
    (Nothing, _) -> Nothing
    (Just c, m2) -> Just (c, Expr m2)
-}

showExpr :: (Num r, Eq r, Show r) => Expr r -> String
showExpr = showExprWith f
  where
    f x = "x" ++ show x

showExprWith :: (Num r, Show r, Eq r) => (Var -> String) -> Expr r -> String
showExprWith env (Expr m) = foldr (.) id xs ""
  where
    xs = intersperse (showString " + ") ts
    ts = [if c==1
            then showString (env x)
            else showsPrec 8 c . showString "*" . showString (env x)
          | (x,c) <- IntMap.toList m, x /= unitVar] ++
         [showsPrec 7 c | c <- maybeToList (IntMap.lookup unitVar m)]

-----------------------------------------------------------------------------

-- | Atomic Formula of Linear Arithmetics
type Atom r = OrdRel (Expr r)

showAtom :: (Num r, Eq r, Show r) => Atom r -> String
showAtom (OrdRel lhs op rhs) = showExpr lhs ++ showOp op ++ showExpr rhs

applySubstAtom :: (Num r, Eq r) => VarMap (Expr r) -> Atom r -> Atom r
applySubstAtom s (OrdRel lhs op rhs) = OrdRel (applySubst s lhs) op (applySubst s rhs)

-- | applySubst1 x e phi == phi[e/x]
applySubst1Atom :: (Num r, Eq r) => Var -> Expr r -> Atom r -> Atom r
applySubst1Atom x e (OrdRel lhs op rhs) = OrdRel (applySubst1 x e lhs) op (applySubst1 x e rhs)

-- | Solve linear (in)equation for the given variable.
--
-- @solveFor a v@ returns @Just (op, e)@ such that @Atom v op e@
-- is equivalent to @a@.
solveFor :: (Real r, Fractional r) => Atom r -> Var -> Maybe (RelOp, Expr r)
solveFor (OrdRel lhs op rhs) v = do
  (c,e) <- extractMaybe v (lhs ^-^ rhs)
  return ( if c < 0 then flipOp op else op
         , (1/c) *^ negateV e
         )

-----------------------------------------------------------------------------

type BoundsEnv r = VarMap (Interval r)

-- | compute bounds for a @Expr@ with respect to @BoundsEnv@.
computeInterval :: (Real r, Fractional r) => BoundsEnv r -> Expr r -> Interval r
computeInterval b = eval b . mapCoeff Data.Interval.singleton

-----------------------------------------------------------------------------

1	{-# OPTIONS_HADDOCK show-extensions #-}
2	{-# LANGUAGE FlexibleInstances #-}
3	{-# LANGUAGE MultiParamTypeClasses #-}
4	{-# LANGUAGE TypeFamilies #-}
5	{-# LANGUAGE TypeSynonymInstances #-}
6	-----------------------------------------------------------------------------
7	-- \|
8	-- Module : ToySolver.Data.LA
9	-- Copyright : (c) Masahiro Sakai 2011
10	-- License : BSD-style
11	--
12	-- Maintainer : masahiro.sakai@gmail.com
13	-- Stability : provisional
14	-- Portability : non-portable
15	--
16	-- Some definition for Theory of Linear Arithmetics.
17	--
18	-----------------------------------------------------------------------------
19	module ToySolver.Data.LA
20	(
21	-- * Expression of linear arithmetics
22	Expr
23	, Var
24
25	-- ** Conversion
26	, var
27	, constant
28	, terms
29	, fromTerms
30	, coeffMap
31	, fromCoeffMap
32	, unitVar
33	, asConst
34
35	-- ** Query
36	, coeff
37	, lookupCoeff
38	, extract
39	, extractMaybe
40
41	-- ** Operations
42	, mapCoeff
43	, mapCoeffWithVar
44	, evalLinear
45	, lift1
46	, applySubst
47	, applySubst1
48	, showExpr
49
50	-- * Atomic formula of linear arithmetics
51	, Atom (..)
52	, showAtom
53	, applySubstAtom
54	, applySubst1Atom
55	, solveFor
56	, module ToySolver.Data.OrdRel
57
58	-- * Evaluation of expression and atomic formula
59	, Eval (..)
60
61	-- * misc
62	, BoundsEnv
63	, computeInterval
64	) where
65
66	import Control.Monad
67	import Control.DeepSeq
68	import Data.List
69	import Data.Maybe
70	import Data.IntMap.Strict (IntMap)
71	import qualified Data.IntMap.Strict as IntMap
72	import qualified Data.IntSet as IntSet
73	import Data.Interval
74	import Data.VectorSpace
75
76	import qualified ToySolver.Data.OrdRel as OrdRel
77	import ToySolver.Data.OrdRel
78	import ToySolver.Data.IntVar
79
80	-----------------------------------------------------------------------------
81
82	-- \| Linear combination of variables and constants.
83	-- Non-negative keys are used for variables's coefficients.
84	-- key 'unitVar' is used for constants.
85	newtype Expr r
86	= Expr
87	{ -- \| a mapping from variables to coefficients
88	coeffMap :: IntMap r	1✔
89	} deriving (Eq, Ord)	×
90
91	-- \| Create a @Expr@ from a mapping from variables to coefficients.
92	fromCoeffMap :: (Num r, Eq r) => IntMap r -> Expr r
93	fromCoeffMap m = normalizeExpr (Expr m)	1✔
94
95	-- \| terms contained in the expression.
96	terms :: Expr r -> [(r,Var)]
97	terms (Expr m) = [(c,v) \| (v,c) <- IntMap.toList m]	1✔
98
99	-- \| Create a @Expr@ from a list of terms.
100	fromTerms :: (Num r, Eq r) => [(r,Var)] -> Expr r
101	fromTerms ts = fromCoeffMap $ IntMap.fromListWith (+) [(x,c) \| (c,x) <- ts]	1✔
102
103	instance Variables (Expr r) where
104	vars (Expr m) = IntSet.delete unitVar (IntMap.keysSet m)	1✔
105
106	instance Show r => Show (Expr r) where	×
107	showsPrec d m = showParen (d > 10) $	×
108	showString "fromTerms " . shows (terms m)	×
109
110	instance (Num r, Eq r, Read r) => Read (Expr r) where	×
111	readsPrec p = readParen (p > 10) $ \ r -> do	×
112	("fromTerms",s) <- lex r	×
113	(xs,t) <- reads s	×
114	return (fromTerms xs, t)	×
115
116	instance NFData r => NFData (Expr r) where
117	rnf (Expr m) = rnf m	×
118
119	-- \| Special variable that should always be evaluated to 1.
120	unitVar :: Var
121	unitVar = -1	1✔
122
123	asConst :: Num r => Expr r -> Maybe r
124	asConst (Expr m) =	1✔
125	case IntMap.toList m of	1✔
126	[] -> Just 0	1✔
127	[(v,x)] \| v==unitVar -> Just x	1✔
128	_ -> Nothing	1✔
129
130	normalizeExpr :: (Num r, Eq r) => Expr r -> Expr r
131	normalizeExpr (Expr t) = Expr $ IntMap.filter (0/=) t	1✔
132
133	-- \| variable
134	var :: Num r => Var -> Expr r
135	var v = Expr $ IntMap.singleton v 1	1✔
136
137	-- \| constant
138	constant :: (Num r, Eq r) => r -> Expr r
139	constant c = normalizeExpr $ Expr $ IntMap.singleton unitVar c	1✔
140
141	-- \| map coefficients.
142	mapCoeff :: (Num b, Eq b) => (a -> b) -> Expr a -> Expr b
143	mapCoeff f (Expr t) = Expr $ IntMap.mapMaybe g t	1✔
144	where
145	g c = if c' == 0 then Nothing else Just c'	1✔
146	where c' = f c	1✔
147
148	-- \| map coefficients.
149	mapCoeffWithVar :: (Num b, Eq b) => (a -> Var -> b) -> Expr a -> Expr b
150	mapCoeffWithVar f (Expr t) = Expr $ IntMap.mapMaybeWithKey g t	1✔
151	where
152	g v c = if c' == 0 then Nothing else Just c'	×
153	where c' = f c v	1✔
154
155	instance (Num r, Eq r) => AdditiveGroup (Expr r) where	1✔
156	Expr t ^+^ e2 \| IntMap.null t = e2	1✔
157	e1 ^+^ Expr t \| IntMap.null t = e1	1✔
158	e1 ^+^ e2 = normalizeExpr $ plus e1 e2	1✔
159	zeroV = Expr $ IntMap.empty	1✔
160	negateV = ((-1) *^)	1✔
161
162	instance (Num r, Eq r) => VectorSpace (Expr r) where
163	type Scalar (Expr r) = r
164	1 *^ e = e	1✔
UNCOV 165	0 *^ e = zeroV	×
166	c ^ e = mapCoeff (c) e	1✔
167
168	plus :: Num r => Expr r -> Expr r -> Expr r
169	plus (Expr t1) (Expr t2) = Expr $ IntMap.unionWith (+) t1 t2	1✔
170
171	instance Num r => Eval (Model r) (Expr r) r where
172	eval m (Expr t) = sum [(m' IntMap.! v) * c \| (v,c) <- IntMap.toList t]	1✔
173	where m' = IntMap.insert unitVar 1 m	1✔
174
175	-- \| evaluate the expression under the model.
176	evalLinear :: VectorSpace a => Model a -> a -> Expr (Scalar a) -> a
177	evalLinear m u (Expr t) = sumV [c *^ (m' IntMap.! v) \| (v,c) <- IntMap.toList t]	1✔
178	where m' = IntMap.insert unitVar u m	1✔
179
180	lift1 :: VectorSpace x => x -> (Var -> x) -> Expr (Scalar x) -> x
181	lift1 unit f (Expr t) = sumV [c *^ (g v) \| (v,c) <- IntMap.toList t]	1✔
182	where
183	g v	1✔
184	\| v==unitVar = unit	1✔
185	\| otherwise = f v	×
186
187	applySubst :: (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
188	applySubst s (Expr m) = sumV (map f (IntMap.toList m))	1✔
189	where
190	f (v,c) = c *^ (	1✔
191	case IntMap.lookup v s of	1✔
192	Just tm -> tm	1✔
193	Nothing -> var v)	1✔
194
195	-- \| applySubst1 x e e1 == e1[e/x]
196	applySubst1 :: (Num r, Eq r) => Var -> Expr r -> Expr r -> Expr r
197	applySubst1 x e e1 =	1✔
198	case extractMaybe x e1 of	1✔
199	Nothing -> e1	1✔
200	Just (c,e2) -> c *^ e ^+^ e2	1✔
201
202	-- \| lookup a coefficient of the variable.
203	-- @
204	-- coeff v e == fst (extract v e)
205	-- @
206	coeff :: Num r => Var -> Expr r -> r
207	coeff v (Expr m) = IntMap.findWithDefault 0 v m	1✔
208
209	-- \| lookup a coefficient of the variable.
210	-- It returns @Nothing@ if the expression does not contain @v@.
211	-- @
212	-- lookupCoeff v e == fmap fst (extractMaybe v e)
213	-- @
214	lookupCoeff :: Num r => Var -> Expr r -> Maybe r
215	lookupCoeff v (Expr m) = IntMap.lookup v m	1✔
216
217	-- \| @extract v e@ returns @(c, e')@ such that @e == c *^ v ^+^ e'@
218	extract :: Num r => Var -> Expr r -> (r, Expr r)
219	extract v (Expr m) = (IntMap.findWithDefault 0 v m, Expr (IntMap.delete v m))	1✔
220	{-
221	-- Alternative implementation which may be faster but allocte more memory
222	extract v (Expr m) =
223	case IntMap.updateLookupWithKey (\_ _ -> Nothing) v m of
224	(Nothing, _) -> (0, Expr m)
225	(Just c, m2) -> (c, Expr m2)
226	-}
227
228	-- \| @extractMaybe v e@ returns @Just (c, e')@ such that @e == c *^ v ^+^ e'@
229	-- if @e@ contains v, and returns @Nothing@ otherwise.
230	extractMaybe :: Num r => Var -> Expr r -> Maybe (r, Expr r)
231	extractMaybe v (Expr m) =	1✔
232	case IntMap.lookup v m of	1✔
233	Nothing -> Nothing	1✔
234	Just c -> Just (c, Expr (IntMap.delete v m))	1✔
235	{-
236	-- Alternative implementation which may be faster but allocte more memory
237	extractMaybe v (Expr m) =
238	case IntMap.updateLookupWithKey (\_ _ -> Nothing) v m of
239	(Nothing, _) -> Nothing
240	(Just c, m2) -> Just (c, Expr m2)
241	-}
242
243	showExpr :: (Num r, Eq r, Show r) => Expr r -> String
244	showExpr = showExprWith f	×
245	where
246	f x = "x" ++ show x	×
247
248	showExprWith :: (Num r, Show r, Eq r) => (Var -> String) -> Expr r -> String
249	showExprWith env (Expr m) = foldr (.) id xs ""	×
250	where
251	xs = intersperse (showString " + ") ts	×
252	ts = [if c==1	×
253	then showString (env x)	×
254	else showsPrec 8 c . showString "*" . showString (env x)	×
255	\| (x,c) <- IntMap.toList m, x /= unitVar] ++	×
256	[showsPrec 7 c \| c <- maybeToList (IntMap.lookup unitVar m)]	×
257
258	-----------------------------------------------------------------------------
259
260	-- \| Atomic Formula of Linear Arithmetics
261	type Atom r = OrdRel (Expr r)
262
263	showAtom :: (Num r, Eq r, Show r) => Atom r -> String
264	showAtom (OrdRel lhs op rhs) = showExpr lhs ++ showOp op ++ showExpr rhs	×
265
266	applySubstAtom :: (Num r, Eq r) => VarMap (Expr r) -> Atom r -> Atom r
267	applySubstAtom s (OrdRel lhs op rhs) = OrdRel (applySubst s lhs) op (applySubst s rhs)	1✔
268
269	-- \| applySubst1 x e phi == phi[e/x]
270	applySubst1Atom :: (Num r, Eq r) => Var -> Expr r -> Atom r -> Atom r
271	applySubst1Atom x e (OrdRel lhs op rhs) = OrdRel (applySubst1 x e lhs) op (applySubst1 x e rhs)	×
272
273	-- \| Solve linear (in)equation for the given variable.
274	--
275	-- @solveFor a v@ returns @Just (op, e)@ such that @Atom v op e@
276	-- is equivalent to @a@.
277	solveFor :: (Real r, Fractional r) => Atom r -> Var -> Maybe (RelOp, Expr r)
278	solveFor (OrdRel lhs op rhs) v = do	1✔
279	(c,e) <- extractMaybe v (lhs ^-^ rhs)	1✔
280	return ( if c < 0 then flipOp op else op	1✔
281	, (1/c) *^ negateV e	1✔
282	)
283
284	-----------------------------------------------------------------------------
285
286	type BoundsEnv r = VarMap (Interval r)
287
288	-- \| compute bounds for a @Expr@ with respect to @BoundsEnv@.
289	computeInterval :: (Real r, Fractional r) => BoundsEnv r -> Expr r -> Interval r
290	computeInterval b = eval b . mapCoeff Data.Interval.singleton	1✔
291
292	-----------------------------------------------------------------------------

msakai / toysolver / 513

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous