Minor changes in plot functions
[linpy.git] / pypol / domains.py
1 import ast
2 import functools
3 import re
4 import math
5
6 from fractions import Fraction
7
8 from . import islhelper
9 from .islhelper import mainctx, libisl
10 from .linexprs import Expression, Symbol, Rational
11 from .geometry import GeometricObject, Point, Vector
12
13
14 __all__ = [
15 'Domain',
16 'And', 'Or', 'Not',
17 ]
18
19
20 @functools.total_ordering
21 class Domain(GeometricObject):
22
23 __slots__ = (
24 '_polyhedra',
25 '_symbols',
26 '_dimension',
27 )
28
29 def __new__(cls, *polyhedra):
30 from .polyhedra import Polyhedron
31 if len(polyhedra) == 1:
32 argument = polyhedra[0]
33 if isinstance(argument, str):
34 return cls.fromstring(argument)
35 elif isinstance(argument, GeometricObject):
36 return argument.aspolyhedron()
37 else:
38 raise TypeError('argument must be a string '
39 'or a GeometricObject instance')
40 else:
41 for polyhedron in polyhedra:
42 if not isinstance(polyhedron, Polyhedron):
43 raise TypeError('arguments must be Polyhedron instances')
44 symbols = cls._xsymbols(polyhedra)
45 islset = cls._toislset(polyhedra, symbols)
46 return cls._fromislset(islset, symbols)
47
48 @classmethod
49 def _xsymbols(cls, iterator):
50 """
51 Return the ordered tuple of symbols present in iterator.
52 """
53 symbols = set()
54 for item in iterator:
55 symbols.update(item.symbols)
56 return tuple(sorted(symbols, key=Symbol.sortkey))
57
58 @property
59 def polyhedra(self):
60 return self._polyhedra
61
62 @property
63 def symbols(self):
64 return self._symbols
65
66 @property
67 def dimension(self):
68 return self._dimension
69
70 def disjoint(self):
71 """
72 Returns this set as disjoint.
73 """
74 islset = self._toislset(self.polyhedra, self.symbols)
75 islset = libisl.isl_set_make_disjoint(mainctx, islset)
76 return self._fromislset(islset, self.symbols)
77
78 def isempty(self):
79 """
80 Returns true if this set is an Empty set.
81 """
82 islset = self._toislset(self.polyhedra, self.symbols)
83 empty = bool(libisl.isl_set_is_empty(islset))
84 libisl.isl_set_free(islset)
85 return empty
86
87 def __bool__(self):
88 return not self.isempty()
89
90 def isuniverse(self):
91 """
92 Returns true if this set is the Universe set.
93 """
94 islset = self._toislset(self.polyhedra, self.symbols)
95 universe = bool(libisl.isl_set_plain_is_universe(islset))
96 libisl.isl_set_free(islset)
97 return universe
98
99 def isbounded(self):
100 """
101 Returns true if this set is bounded.
102 """
103 islset = self._toislset(self.polyhedra, self.symbols)
104 bounded = bool(libisl.isl_set_is_bounded(islset))
105 libisl.isl_set_free(islset)
106 return bounded
107
108 def __eq__(self, other):
109 """
110 Returns true if two sets are equal.
111 """
112 symbols = self._xsymbols([self, other])
113 islset1 = self._toislset(self.polyhedra, symbols)
114 islset2 = other._toislset(other.polyhedra, symbols)
115 equal = bool(libisl.isl_set_is_equal(islset1, islset2))
116 libisl.isl_set_free(islset1)
117 libisl.isl_set_free(islset2)
118 return equal
119
120 def isdisjoint(self, other):
121 """
122 Return True if two sets have a null intersection.
123 """
124 symbols = self._xsymbols([self, other])
125 islset1 = self._toislset(self.polyhedra, symbols)
126 islset2 = self._toislset(other.polyhedra, symbols)
127 equal = bool(libisl.isl_set_is_disjoint(islset1, islset2))
128 libisl.isl_set_free(islset1)
129 libisl.isl_set_free(islset2)
130 return equal
131
132 def issubset(self, other):
133 """
134 Report whether another set contains this set.
135 """
136 symbols = self._xsymbols([self, other])
137 islset1 = self._toislset(self.polyhedra, symbols)
138 islset2 = self._toislset(other.polyhedra, symbols)
139 equal = bool(libisl.isl_set_is_subset(islset1, islset2))
140 libisl.isl_set_free(islset1)
141 libisl.isl_set_free(islset2)
142 return equal
143
144 def __le__(self, other):
145 """
146 Returns true if this set is less than or equal to another set.
147 """
148 return self.issubset(other)
149
150 def __lt__(self, other):
151 """
152 Returns true if this set is less than another set.
153 """
154 symbols = self._xsymbols([self, other])
155 islset1 = self._toislset(self.polyhedra, symbols)
156 islset2 = self._toislset(other.polyhedra, symbols)
157 equal = bool(libisl.isl_set_is_strict_subset(islset1, islset2))
158 libisl.isl_set_free(islset1)
159 libisl.isl_set_free(islset2)
160 return equal
161
162 def complement(self):
163 """
164 Returns the complement of this set.
165 """
166 islset = self._toislset(self.polyhedra, self.symbols)
167 islset = libisl.isl_set_complement(islset)
168 return self._fromislset(islset, self.symbols)
169
170 def __invert__(self):
171 """
172 Returns the complement of this set.
173 """
174 return self.complement()
175
176 def simplify(self):
177 """
178 Returns a set without redundant constraints.
179 """
180 islset = self._toislset(self.polyhedra, self.symbols)
181 islset = libisl.isl_set_remove_redundancies(islset)
182 return self._fromislset(islset, self.symbols)
183
184 def aspolyhedron(self):
185 """
186 Returns polyhedral hull of set.
187 """
188 from .polyhedra import Polyhedron
189 islset = self._toislset(self.polyhedra, self.symbols)
190 islbset = libisl.isl_set_polyhedral_hull(islset)
191 return Polyhedron._fromislbasicset(islbset, self.symbols)
192
193 def asdomain(self):
194 return self
195
196 def project(self, dims):
197 """
198 Return new set with given dimensions removed.
199 """
200 islset = self._toislset(self.polyhedra, self.symbols)
201 n = 0
202 for index, symbol in reversed(list(enumerate(self.symbols))):
203 if symbol in dims:
204 n += 1
205 elif n > 0:
206 islset = libisl.isl_set_project_out(islset, libisl.isl_dim_set, index + 1, n)
207 n = 0
208 if n > 0:
209 islset = libisl.isl_set_project_out(islset, libisl.isl_dim_set, 0, n)
210 dims = [symbol for symbol in self.symbols if symbol not in dims]
211 return Domain._fromislset(islset, dims)
212
213 def sample(self):
214 """
215 Returns a single subset of the input.
216 """
217 islset = self._toislset(self.polyhedra, self.symbols)
218 islpoint = libisl.isl_set_sample_point(islset)
219 if bool(libisl.isl_point_is_void(islpoint)):
220 libisl.isl_point_free(islpoint)
221 raise ValueError('domain must be non-empty')
222 point = {}
223 for index, symbol in enumerate(self.symbols):
224 coordinate = libisl.isl_point_get_coordinate_val(islpoint,
225 libisl.isl_dim_set, index)
226 coordinate = islhelper.isl_val_to_int(coordinate)
227 point[symbol] = coordinate
228 libisl.isl_point_free(islpoint)
229 return point
230
231 def intersection(self, *others):
232 """
233 Return the intersection of two sets as a new set.
234 """
235 if len(others) == 0:
236 return self
237 symbols = self._xsymbols((self,) + others)
238 islset1 = self._toislset(self.polyhedra, symbols)
239 for other in others:
240 islset2 = other._toislset(other.polyhedra, symbols)
241 islset1 = libisl.isl_set_intersect(islset1, islset2)
242 return self._fromislset(islset1, symbols)
243
244 def __and__(self, other):
245 """
246 Return the intersection of two sets as a new set.
247 """
248 return self.intersection(other)
249
250 def union(self, *others):
251 """
252 Return the union of sets as a new set.
253 """
254 if len(others) == 0:
255 return self
256 symbols = self._xsymbols((self,) + others)
257 islset1 = self._toislset(self.polyhedra, symbols)
258 for other in others:
259 islset2 = other._toislset(other.polyhedra, symbols)
260 islset1 = libisl.isl_set_union(islset1, islset2)
261 return self._fromislset(islset1, symbols)
262
263 def __or__(self, other):
264 """
265 Return a new set with elements from both sets.
266 """
267 return self.union(other)
268
269 def __add__(self, other):
270 """
271 Return new set containing all elements in both sets.
272 """
273 return self.union(other)
274
275 def difference(self, other):
276 """
277 Return the difference of two sets as a new set.
278 """
279 symbols = self._xsymbols([self, other])
280 islset1 = self._toislset(self.polyhedra, symbols)
281 islset2 = other._toislset(other.polyhedra, symbols)
282 islset = libisl.isl_set_subtract(islset1, islset2)
283 return self._fromislset(islset, symbols)
284
285 def __sub__(self, other):
286 """
287 Return the difference of two sets as a new set.
288 """
289 return self.difference(other)
290
291 def lexmin(self):
292 """
293 Return a new set containing the lexicographic minimum of the elements in the set.
294 """
295 islset = self._toislset(self.polyhedra, self.symbols)
296 islset = libisl.isl_set_lexmin(islset)
297 return self._fromislset(islset, self.symbols)
298
299 def lexmax(self):
300 """
301 Return a new set containing the lexicographic maximum of the elements in the set.
302 """
303 islset = self._toislset(self.polyhedra, self.symbols)
304 islset = libisl.isl_set_lexmax(islset)
305 return self._fromislset(islset, self.symbols)
306
307 def num_parameters(self):
308 """
309 Return the total number of parameters, input, output or set dimensions.
310 """
311 islbset = self._toislbasicset(self.equalities, self.inequalities, self.symbols)
312 num = libisl.isl_basic_set_dim(islbset, libisl.isl_dim_set)
313 return num
314
315 def involves_dims(self, dims):
316 """
317 Returns true if set depends on given dimensions.
318 """
319 islset = self._toislset(self.polyhedra, self.symbols)
320 dims = sorted(dims)
321 symbols = sorted(list(self.symbols))
322 n = 0
323 if len(dims)>0:
324 for dim in dims:
325 if dim in symbols:
326 first = symbols.index(dims[0])
327 n +=1
328 else:
329 first = 0
330 else:
331 return False
332 value = bool(libisl.isl_set_involves_dims(islset, libisl.isl_dim_set, first, n))
333 libisl.isl_set_free(islset)
334 return value
335
336 _RE_COORDINATE = re.compile(r'\((?P<num>\-?\d+)\)(/(?P<den>\d+))?')
337
338 def vertices(self):
339 """
340 Return a list of vertices for this Polygon.
341 """
342 from .polyhedra import Polyhedron
343 islbset = self._toislbasicset(self.equalities, self.inequalities, self.symbols)
344 vertices = libisl.isl_basic_set_compute_vertices(islbset);
345 vertices = islhelper.isl_vertices_vertices(vertices)
346 points = []
347 for vertex in vertices:
348 expr = libisl.isl_vertex_get_expr(vertex)
349 coordinates = []
350 if islhelper.isl_version < '0.13':
351 constraints = islhelper.isl_basic_set_constraints(expr)
352 for constraint in constraints:
353 constant = libisl.isl_constraint_get_constant_val(constraint)
354 constant = islhelper.isl_val_to_int(constant)
355 for index, symbol in enumerate(self.symbols):
356 coefficient = libisl.isl_constraint_get_coefficient_val(constraint,
357 libisl.isl_dim_set, index)
358 coefficient = islhelper.isl_val_to_int(coefficient)
359 if coefficient != 0:
360 coordinate = -Fraction(constant, coefficient)
361 coordinates.append((symbol, coordinate))
362 else:
363 string = islhelper.isl_multi_aff_to_str(expr)
364 matches = self._RE_COORDINATE.finditer(string)
365 for symbol, match in zip(self.symbols, matches):
366 numerator = int(match.group('num'))
367 denominator = match.group('den')
368 denominator = 1 if denominator is None else int(denominator)
369 coordinate = Fraction(numerator, denominator)
370 coordinates.append((symbol, coordinate))
371 points.append(Point(coordinates))
372 return points
373
374 def points(self):
375 """
376 Returns the points contained in the set.
377 """
378 if not self.isbounded():
379 raise ValueError('domain must be bounded')
380 from .polyhedra import Universe, Eq
381 islset = self._toislset(self.polyhedra, self.symbols)
382 islpoints = islhelper.isl_set_points(islset)
383 points = []
384 for islpoint in islpoints:
385 coordinates = {}
386 for index, symbol in enumerate(self.symbols):
387 coordinate = libisl.isl_point_get_coordinate_val(islpoint,
388 libisl.isl_dim_set, index)
389 coordinate = islhelper.isl_val_to_int(coordinate)
390 coordinates[symbol] = coordinate
391 points.append(Point(coordinates))
392 return points
393
394 @classmethod
395 def _polygon_inner_point(cls, points):
396 symbols = points[0].symbols
397 coordinates = {symbol: 0 for symbol in symbols}
398 for point in points:
399 for symbol, coordinate in point.coordinates():
400 coordinates[symbol] += coordinate
401 for symbol in symbols:
402 coordinates[symbol] /= len(points)
403 return Point(coordinates)
404
405 @classmethod
406 def _sort_polygon_2d(cls, points):
407 if len(points) <= 3:
408 return points
409 o = cls._polygon_inner_point(points)
410 angles = {}
411 for m in points:
412 om = Vector(o, m)
413 dx, dy = (coordinate for symbol, coordinate in om.coordinates())
414 angle = math.atan2(dy, dx)
415 angles[m] = angle
416 return sorted(points, key=angles.get)
417
418 @classmethod
419 def _sort_polygon_3d(cls, points):
420 if len(points) <= 3:
421 return points
422 o = cls._polygon_inner_point(points)
423 a = points[0]
424 oa = Vector(o, a)
425 norm_oa = oa.norm()
426 for b in points[1:]:
427 ob = Vector(o, b)
428 u = oa.cross(ob)
429 if not u.isnull():
430 u = u.asunit()
431 break
432 else:
433 raise ValueError('degenerate polygon')
434 angles = {a: 0.}
435 for m in points[1:]:
436 om = Vector(o, m)
437 normprod = norm_oa * om.norm()
438 cosinus = max(oa.dot(om) / normprod, -1.)
439 sinus = u.dot(oa.cross(om)) / normprod
440 angle = math.acos(cosinus)
441 angle = math.copysign(angle, sinus)
442 angles[m] = angle
443 return sorted(points, key=angles.get)
444
445 def faces(self):
446 faces = []
447 for polyhedron in self.polyhedra:
448 vertices = polyhedron.vertices()
449 for constraint in polyhedron.constraints:
450 face = []
451 for vertex in vertices:
452 if constraint.subs(vertex.coordinates()) == 0:
453 face.append(vertex)
454 if len(face) >= 3:
455 faces.append(face)
456 return faces
457
458 def _plot_2d(self, plot=None, **kwargs):
459 import matplotlib.pyplot as plt
460 from matplotlib.patches import Polygon
461 if plot is None:
462 fig = plt.figure()
463 plot = fig.add_subplot(1, 1, 1)
464 xmin, xmax = plot.get_xlim()
465 ymin, ymax = plot.get_ylim()
466 for polyhedron in self.polyhedra:
467 vertices = polyhedron._sort_polygon_2d(polyhedron.vertices())
468 xys = [tuple(vertex.values()) for vertex in vertices]
469 xs, ys = zip(*xys)
470 xmin, xmax = min(xmin, float(min(xs))), max(xmax, float(max(xs)))
471 ymin, ymax = min(ymin, float(min(ys))), max(ymax, float(max(ys)))
472 plot.add_patch(Polygon(xys, closed=True, **kwargs))
473 plot.set_xlim(xmin, xmax)
474 plot.set_ylim(ymin, ymax)
475 return plot
476
477 def _plot_3d(self, plot=None, **kwargs):
478 import matplotlib.pyplot as plt
479 from mpl_toolkits.mplot3d import Axes3D
480 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
481 if plot is None:
482 fig = plt.figure()
483 axes = Axes3D(fig)
484 else:
485 axes = plot
486 xmin, xmax = axes.get_xlim()
487 ymin, ymax = axes.get_ylim()
488 zmin, zmax = axes.get_zlim()
489 poly_xyzs = []
490 for vertices in self.faces():
491 vertices = self._sort_polygon_3d(vertices)
492 vertices.append(vertices[0])
493 face_xyzs = [tuple(vertex.values()) for vertex in vertices]
494 xs, ys, zs = zip(*face_xyzs)
495 xmin, xmax = min(xmin, float(min(xs))), max(xmax, float(max(xs)))
496 ymin, ymax = min(ymin, float(min(ys))), max(ymax, float(max(ys)))
497 zmin, zmax = min(zmin, float(min(zs))), max(zmax, float(max(zs)))
498 poly_xyzs.append(face_xyzs)
499 collection = Poly3DCollection(poly_xyzs, **kwargs)
500 axes.add_collection3d(collection)
501 axes.set_xlim(xmin, xmax)
502 axes.set_ylim(ymin, ymax)
503 axes.set_zlim(zmin, zmax)
504 return axes
505
506 def plot(self, plot=None, **kwargs):
507 """
508 Display plot of this set.
509 """
510 if not self.isbounded():
511 raise ValueError('domain must be bounded')
512 elif self.dimension == 2:
513 return self._plot_2d(plot=plot, **kwargs)
514 elif self.dimension == 3:
515 return self._plot_3d(plot=plot, **kwargs)
516 else:
517 raise ValueError('polyhedron must be 2 or 3-dimensional')
518
519 def __contains__(self, point):
520 for polyhedron in self.polyhedra:
521 if point in polyhedron:
522 return True
523 return False
524
525 def subs(self, symbol, expression=None):
526 polyhedra = [polyhedron.subs(symbol, expression)
527 for polyhedron in self.polyhedra]
528 return Domain(*polyhedra)
529
530 @classmethod
531 def _fromislset(cls, islset, symbols):
532 from .polyhedra import Polyhedron
533 islset = libisl.isl_set_remove_divs(islset)
534 islbsets = islhelper.isl_set_basic_sets(islset)
535 libisl.isl_set_free(islset)
536 polyhedra = []
537 for islbset in islbsets:
538 polyhedron = Polyhedron._fromislbasicset(islbset, symbols)
539 polyhedra.append(polyhedron)
540 if len(polyhedra) == 0:
541 from .polyhedra import Empty
542 return Empty
543 elif len(polyhedra) == 1:
544 return polyhedra[0]
545 else:
546 self = object().__new__(Domain)
547 self._polyhedra = tuple(polyhedra)
548 self._symbols = cls._xsymbols(polyhedra)
549 self._dimension = len(self._symbols)
550 return self
551
552 @classmethod
553 def _toislset(cls, polyhedra, symbols):
554 polyhedron = polyhedra[0]
555 islbset = polyhedron._toislbasicset(polyhedron.equalities,
556 polyhedron.inequalities, symbols)
557 islset1 = libisl.isl_set_from_basic_set(islbset)
558 for polyhedron in polyhedra[1:]:
559 islbset = polyhedron._toislbasicset(polyhedron.equalities,
560 polyhedron.inequalities, symbols)
561 islset2 = libisl.isl_set_from_basic_set(islbset)
562 islset1 = libisl.isl_set_union(islset1, islset2)
563 return islset1
564
565 @classmethod
566 def _fromast(cls, node):
567 from .polyhedra import Polyhedron
568 if isinstance(node, ast.Module) and len(node.body) == 1:
569 return cls._fromast(node.body[0])
570 elif isinstance(node, ast.Expr):
571 return cls._fromast(node.value)
572 elif isinstance(node, ast.UnaryOp):
573 domain = cls._fromast(node.operand)
574 if isinstance(node.operand, ast.invert):
575 return Not(domain)
576 elif isinstance(node, ast.BinOp):
577 domain1 = cls._fromast(node.left)
578 domain2 = cls._fromast(node.right)
579 if isinstance(node.op, ast.BitAnd):
580 return And(domain1, domain2)
581 elif isinstance(node.op, ast.BitOr):
582 return Or(domain1, domain2)
583 elif isinstance(node, ast.Compare):
584 equalities = []
585 inequalities = []
586 left = Expression._fromast(node.left)
587 for i in range(len(node.ops)):
588 op = node.ops[i]
589 right = Expression._fromast(node.comparators[i])
590 if isinstance(op, ast.Lt):
591 inequalities.append(right - left - 1)
592 elif isinstance(op, ast.LtE):
593 inequalities.append(right - left)
594 elif isinstance(op, ast.Eq):
595 equalities.append(left - right)
596 elif isinstance(op, ast.GtE):
597 inequalities.append(left - right)
598 elif isinstance(op, ast.Gt):
599 inequalities.append(left - right - 1)
600 else:
601 break
602 left = right
603 else:
604 return Polyhedron(equalities, inequalities)
605 raise SyntaxError('invalid syntax')
606
607 _RE_BRACES = re.compile(r'^\{\s*|\s*\}$')
608 _RE_EQ = re.compile(r'([^<=>])=([^<=>])')
609 _RE_AND = re.compile(r'\band\b|,|&&|/\\|∧|∩')
610 _RE_OR = re.compile(r'\bor\b|;|\|\||\\/|∨|∪')
611 _RE_NOT = re.compile(r'\bnot\b|!|¬')
612 _RE_NUM_VAR = Expression._RE_NUM_VAR
613 _RE_OPERATORS = re.compile(r'(&|\||~)')
614
615 @classmethod
616 def fromstring(cls, string):
617 # remove curly brackets
618 string = cls._RE_BRACES.sub(r'', string)
619 # replace '=' by '=='
620 string = cls._RE_EQ.sub(r'\1==\2', string)
621 # replace 'and', 'or', 'not'
622 string = cls._RE_AND.sub(r' & ', string)
623 string = cls._RE_OR.sub(r' | ', string)
624 string = cls._RE_NOT.sub(r' ~', string)
625 # add implicit multiplication operators, e.g. '5x' -> '5*x'
626 string = cls._RE_NUM_VAR.sub(r'\1*\2', string)
627 # add parentheses to force precedence
628 tokens = cls._RE_OPERATORS.split(string)
629 for i, token in enumerate(tokens):
630 if i % 2 == 0:
631 token = '({})'.format(token)
632 tokens[i] = token
633 string = ''.join(tokens)
634 tree = ast.parse(string, 'eval')
635 return cls._fromast(tree)
636
637 def __repr__(self):
638 assert len(self.polyhedra) >= 2
639 strings = [repr(polyhedron) for polyhedron in self.polyhedra]
640 return 'Or({})'.format(', '.join(strings))
641
642 def _repr_latex_(self):
643 strings = []
644 for polyhedron in self.polyhedra:
645 strings.append('({})'.format(polyhedron._repr_latex_().strip('$')))
646 return '${}$'.format(' \\vee '.join(strings))
647
648 @classmethod
649 def fromsympy(cls, expr):
650 import sympy
651 from .polyhedra import Lt, Le, Eq, Ne, Ge, Gt
652 funcmap = {
653 sympy.And: And, sympy.Or: Or, sympy.Not: Not,
654 sympy.Lt: Lt, sympy.Le: Le,
655 sympy.Eq: Eq, sympy.Ne: Ne,
656 sympy.Ge: Ge, sympy.Gt: Gt,
657 }
658 if expr.func in funcmap:
659 args = [Domain.fromsympy(arg) for arg in expr.args]
660 return funcmap[expr.func](*args)
661 elif isinstance(expr, sympy.Expr):
662 return Expression.fromsympy(expr)
663 raise ValueError('non-domain expression: {!r}'.format(expr))
664
665 def tosympy(self):
666 import sympy
667 polyhedra = [polyhedron.tosympy() for polyhedron in polyhedra]
668 return sympy.Or(*polyhedra)
669
670
671 def And(*domains):
672 """
673 Return the intersection of two sets as a new set.
674 """
675 if len(domains) == 0:
676 from .polyhedra import Universe
677 return Universe
678 else:
679 return domains[0].intersection(*domains[1:])
680
681 def Or(*domains):
682 """
683 Return the union of sets as a new set.
684 """
685 if len(domains) == 0:
686 from .polyhedra import Empty
687 return Empty
688 else:
689 return domains[0].union(*domains[1:])
690
691 def Not(domain):
692 """
693 Returns the complement of this set.
694 """
695 return ~domain