Add plot method
[linpy.git] / pypol / polyhedra.py
1
2 import functools
3 import numbers
4
5 from . import islhelper
6
7 from .islhelper import mainctx, libisl
8 from .linexprs import Expression, Symbol, Rational
9 from .domains import Domain
10
11
12 __all__ = [
13 'Polyhedron',
14 'Lt', 'Le', 'Eq', 'Ne', 'Ge', 'Gt',
15 'Empty', 'Universe',
16 ]
17
18
19 class Polyhedron(Domain):
20
21 __slots__ = (
22 '_equalities',
23 '_inequalities',
24 '_constraints',
25 '_symbols',
26 '_dimension',
27 )
28
29 def __new__(cls, equalities=None, inequalities=None):
30 if isinstance(equalities, str):
31 if inequalities is not None:
32 raise TypeError('too many arguments')
33 return cls.fromstring(equalities)
34 elif isinstance(equalities, Polyhedron):
35 if inequalities is not None:
36 raise TypeError('too many arguments')
37 return equalities
38 elif isinstance(equalities, Domain):
39 if inequalities is not None:
40 raise TypeError('too many arguments')
41 return equalities.aspolyhedron()
42 if equalities is None:
43 equalities = []
44 else:
45 for i, equality in enumerate(equalities):
46 if not isinstance(equality, Expression):
47 raise TypeError('equalities must be linear expressions')
48 equalities[i] = equality.scaleint()
49 if inequalities is None:
50 inequalities = []
51 else:
52 for i, inequality in enumerate(inequalities):
53 if not isinstance(inequality, Expression):
54 raise TypeError('inequalities must be linear expressions')
55 inequalities[i] = inequality.scaleint()
56 symbols = cls._xsymbols(equalities + inequalities)
57 islbset = cls._toislbasicset(equalities, inequalities, symbols)
58 return cls._fromislbasicset(islbset, symbols)
59
60 @property
61 def equalities(self):
62 return self._equalities
63
64 @property
65 def inequalities(self):
66 return self._inequalities
67
68 @property
69 def constraints(self):
70 return self._constraints
71
72 @property
73 def polyhedra(self):
74 return self,
75
76 def disjoint(self):
77 return self
78
79 def isuniverse(self):
80 islbset = self._toislbasicset(self.equalities, self.inequalities,
81 self.symbols)
82 universe = bool(libisl.isl_basic_set_is_universe(islbset))
83 libisl.isl_basic_set_free(islbset)
84 return universe
85
86 def aspolyhedron(self):
87 return self
88
89 def subs(self, symbol, expression=None):
90 equalities = [equality.subs(symbol, expression)
91 for equality in self.equalities]
92 inequalities = [inequality.subs(symbol, expression)
93 for inequality in self.inequalities]
94 return Polyhedron(equalities, inequalities)
95
96 @classmethod
97 def _fromislbasicset(cls, islbset, symbols):
98 islconstraints = islhelper.isl_basic_set_constraints(islbset)
99 equalities = []
100 inequalities = []
101 for islconstraint in islconstraints:
102 constant = libisl.isl_constraint_get_constant_val(islconstraint)
103 constant = islhelper.isl_val_to_int(constant)
104 coefficients = {}
105 for index, symbol in enumerate(symbols):
106 coefficient = libisl.isl_constraint_get_coefficient_val(islconstraint,
107 libisl.isl_dim_set, index)
108 coefficient = islhelper.isl_val_to_int(coefficient)
109 if coefficient != 0:
110 coefficients[symbol] = coefficient
111 expression = Expression(coefficients, constant)
112 if libisl.isl_constraint_is_equality(islconstraint):
113 equalities.append(expression)
114 else:
115 inequalities.append(expression)
116 libisl.isl_basic_set_free(islbset)
117 self = object().__new__(Polyhedron)
118 self._equalities = tuple(equalities)
119 self._inequalities = tuple(inequalities)
120 self._constraints = tuple(equalities + inequalities)
121 self._symbols = cls._xsymbols(self._constraints)
122 self._dimension = len(self._symbols)
123 return self
124
125 @classmethod
126 def _toislbasicset(cls, equalities, inequalities, symbols):
127 dimension = len(symbols)
128 indices = {symbol: index for index, symbol in enumerate(symbols)}
129 islsp = libisl.isl_space_set_alloc(mainctx, 0, dimension)
130 islbset = libisl.isl_basic_set_universe(libisl.isl_space_copy(islsp))
131 islls = libisl.isl_local_space_from_space(islsp)
132 for equality in equalities:
133 isleq = libisl.isl_equality_alloc(libisl.isl_local_space_copy(islls))
134 for symbol, coefficient in equality.coefficients():
135 islval = str(coefficient).encode()
136 islval = libisl.isl_val_read_from_str(mainctx, islval)
137 index = indices[symbol]
138 isleq = libisl.isl_constraint_set_coefficient_val(isleq,
139 libisl.isl_dim_set, index, islval)
140 if equality.constant != 0:
141 islval = str(equality.constant).encode()
142 islval = libisl.isl_val_read_from_str(mainctx, islval)
143 isleq = libisl.isl_constraint_set_constant_val(isleq, islval)
144 islbset = libisl.isl_basic_set_add_constraint(islbset, isleq)
145 for inequality in inequalities:
146 islin = libisl.isl_inequality_alloc(libisl.isl_local_space_copy(islls))
147 for symbol, coefficient in inequality.coefficients():
148 islval = str(coefficient).encode()
149 islval = libisl.isl_val_read_from_str(mainctx, islval)
150 index = indices[symbol]
151 islin = libisl.isl_constraint_set_coefficient_val(islin,
152 libisl.isl_dim_set, index, islval)
153 if inequality.constant != 0:
154 islval = str(inequality.constant).encode()
155 islval = libisl.isl_val_read_from_str(mainctx, islval)
156 islin = libisl.isl_constraint_set_constant_val(islin, islval)
157 islbset = libisl.isl_basic_set_add_constraint(islbset, islin)
158 return islbset
159
160 @classmethod
161 def fromstring(cls, string):
162 domain = Domain.fromstring(string)
163 if not isinstance(domain, Polyhedron):
164 raise ValueError('non-polyhedral expression: {!r}'.format(string))
165 return domain
166
167 def __repr__(self):
168 if self.isempty():
169 return 'Empty'
170 elif self.isuniverse():
171 return 'Universe'
172 else:
173 strings = []
174 for equality in self.equalities:
175 strings.append('0 == {}'.format(equality))
176 for inequality in self.inequalities:
177 strings.append('0 <= {}'.format(inequality))
178 if len(strings) == 1:
179 return strings[0]
180 else:
181 return 'And({})'.format(', '.join(strings))
182
183 @classmethod
184 def fromsympy(cls, expr):
185 domain = Domain.fromsympy(expr)
186 if not isinstance(domain, Polyhedron):
187 raise ValueError('non-polyhedral expression: {!r}'.format(expr))
188 return domain
189
190 def tosympy(self):
191 import sympy
192 constraints = []
193 for equality in self.equalities:
194 constraints.append(sympy.Eq(equality.tosympy(), 0))
195 for inequality in self.inequalities:
196 constraints.append(sympy.Ge(inequality.tosympy(), 0))
197 return sympy.And(*constraints)
198
199 def plot(self):
200 import matplotlib.pyplot as plt
201 from matplotlib.path import Path
202 import matplotlib.patches as patches
203
204 if len(self.symbols)> 3:
205 raise TypeError
206
207 elif len(self.symbols) == 2:
208 verts = self.vertices()
209 points = []
210 codes = [Path.MOVETO]
211 for vert in verts:
212 pairs = ()
213 for sym in sorted(vert, key=Symbol.sortkey):
214 num = vert.get(sym)
215 pairs = pairs + (num,)
216 points.append(pairs)
217 points.append((0.0, 0.0))
218 num = len(points)
219 while num > 2:
220 codes.append(Path.LINETO)
221 num = num - 1
222 else:
223 codes.append(Path.CLOSEPOLY)
224 path = Path(points, codes)
225 fig = plt.figure()
226 ax = fig.add_subplot(111)
227 patch = patches.PathPatch(path, facecolor='blue', lw=2)
228 ax.add_patch(patch)
229 ax.set_xlim(-5,5)
230 ax.set_ylim(-5,5)
231 plt.show()
232
233 elif len(self.symbols)==3:
234 return 0
235
236 return points
237
238
239 def _polymorphic(func):
240 @functools.wraps(func)
241 def wrapper(left, right):
242 if isinstance(left, numbers.Rational):
243 left = Rational(left)
244 elif not isinstance(left, Expression):
245 raise TypeError('left must be a a rational number '
246 'or a linear expression')
247 if isinstance(right, numbers.Rational):
248 right = Rational(right)
249 elif not isinstance(right, Expression):
250 raise TypeError('right must be a a rational number '
251 'or a linear expression')
252 return func(left, right)
253 return wrapper
254
255 @_polymorphic
256 def Lt(left, right):
257 return Polyhedron([], [right - left - 1])
258
259 @_polymorphic
260 def Le(left, right):
261 return Polyhedron([], [right - left])
262
263 @_polymorphic
264 def Eq(left, right):
265 return Polyhedron([left - right], [])
266
267 @_polymorphic
268 def Ne(left, right):
269 return ~Eq(left, right)
270
271 @_polymorphic
272 def Gt(left, right):
273 return Polyhedron([], [left - right - 1])
274
275 @_polymorphic
276 def Ge(left, right):
277 return Polyhedron([], [left - right])
278
279
280 Empty = Eq(1, 0)
281
282 Universe = Polyhedron([])