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