Add methods Domain.subs(), Polyhedron.subs()
[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, Rational
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.aspolyhedron()
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.scaleint()
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.scaleint()
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 aspolyhedron(self):
86 return self
87
88 def subs(self, symbol, expression=None):
89 equalities = [equality.subs(symbol, expression)
90 for equality in self.equalities]
91 inequalities = [inequality.subs(symbol, expression)
92 for inequality in self.inequalities]
93 return Polyhedron(equalities, inequalities)
94
95 @classmethod
96 def _fromislbasicset(cls, islbset, symbols):
97 islconstraints = islhelper.isl_basic_set_constraints(islbset)
98 equalities = []
99 inequalities = []
100 for islconstraint in islconstraints:
101 constant = libisl.isl_constraint_get_constant_val(islconstraint)
102 constant = islhelper.isl_val_to_int(constant)
103 coefficients = {}
104 for index, symbol in enumerate(symbols):
105 coefficient = libisl.isl_constraint_get_coefficient_val(islconstraint,
106 libisl.isl_dim_set, index)
107 coefficient = islhelper.isl_val_to_int(coefficient)
108 if coefficient != 0:
109 coefficients[symbol] = coefficient
110 expression = Expression(coefficients, constant)
111 if libisl.isl_constraint_is_equality(islconstraint):
112 equalities.append(expression)
113 else:
114 inequalities.append(expression)
115 libisl.isl_basic_set_free(islbset)
116 self = object().__new__(Polyhedron)
117 self._equalities = tuple(equalities)
118 self._inequalities = tuple(inequalities)
119 self._constraints = tuple(equalities + inequalities)
120 self._symbols = cls._xsymbols(self._constraints)
121 self._dimension = len(self._symbols)
122 return self
123
124 @classmethod
125 def _toislbasicset(cls, equalities, inequalities, symbols):
126 dimension = len(symbols)
127 indices = {symbol: index for index, symbol in enumerate(symbols)}
128 islsp = libisl.isl_space_set_alloc(mainctx, 0, dimension)
129 islbset = libisl.isl_basic_set_universe(libisl.isl_space_copy(islsp))
130 islls = libisl.isl_local_space_from_space(islsp)
131 for equality in equalities:
132 isleq = libisl.isl_equality_alloc(libisl.isl_local_space_copy(islls))
133 for symbol, coefficient in equality.coefficients():
134 islval = str(coefficient).encode()
135 islval = libisl.isl_val_read_from_str(mainctx, islval)
136 index = indices[symbol]
137 isleq = libisl.isl_constraint_set_coefficient_val(isleq,
138 libisl.isl_dim_set, index, islval)
139 if equality.constant != 0:
140 islval = str(equality.constant).encode()
141 islval = libisl.isl_val_read_from_str(mainctx, islval)
142 isleq = libisl.isl_constraint_set_constant_val(isleq, islval)
143 islbset = libisl.isl_basic_set_add_constraint(islbset, isleq)
144 for inequality in inequalities:
145 islin = libisl.isl_inequality_alloc(libisl.isl_local_space_copy(islls))
146 for symbol, coefficient in inequality.coefficients():
147 islval = str(coefficient).encode()
148 islval = libisl.isl_val_read_from_str(mainctx, islval)
149 index = indices[symbol]
150 islin = libisl.isl_constraint_set_coefficient_val(islin,
151 libisl.isl_dim_set, index, islval)
152 if inequality.constant != 0:
153 islval = str(inequality.constant).encode()
154 islval = libisl.isl_val_read_from_str(mainctx, islval)
155 islin = libisl.isl_constraint_set_constant_val(islin, islval)
156 islbset = libisl.isl_basic_set_add_constraint(islbset, islin)
157 return islbset
158
159 @classmethod
160 def fromstring(cls, string):
161 domain = Domain.fromstring(string)
162 if not isinstance(domain, Polyhedron):
163 raise ValueError('non-polyhedral expression: {!r}'.format(string))
164 return domain
165
166 def __repr__(self):
167 if self.isempty():
168 return 'Empty'
169 elif self.isuniverse():
170 return 'Universe'
171 else:
172 strings = []
173 for equality in self.equalities:
174 strings.append('0 == {}'.format(equality))
175 for inequality in self.inequalities:
176 strings.append('0 <= {}'.format(inequality))
177 if len(strings) == 1:
178 return strings[0]
179 else:
180 return 'And({})'.format(', '.join(strings))
181
182 @classmethod
183 def fromsympy(cls, expr):
184 domain = Domain.fromsympy(expr)
185 if not isinstance(domain, Polyhedron):
186 raise ValueError('non-polyhedral expression: {!r}'.format(expr))
187 return domain
188
189 def tosympy(self):
190 import sympy
191 constraints = []
192 for equality in self.equalities:
193 constraints.append(sympy.Eq(equality.tosympy(), 0))
194 for inequality in self.inequalities:
195 constraints.append(sympy.Ge(inequality.tosympy(), 0))
196 return sympy.And(*constraints)
197
198
199 def _polymorphic(func):
200 @functools.wraps(func)
201 def wrapper(left, right):
202 if isinstance(left, numbers.Rational):
203 left = Rational(left)
204 elif not isinstance(left, Expression):
205 raise TypeError('left must be a a rational number '
206 'or a linear expression')
207 if isinstance(right, numbers.Rational):
208 right = Rational(right)
209 elif not isinstance(right, Expression):
210 raise TypeError('right must be a a rational number '
211 'or a linear expression')
212 return func(left, right)
213 return wrapper
214
215 @_polymorphic
216 def Lt(left, right):
217 return Polyhedron([], [right - left - 1])
218
219 @_polymorphic
220 def Le(left, right):
221 return Polyhedron([], [right - left])
222
223 @_polymorphic
224 def Eq(left, right):
225 return Polyhedron([left - right], [])
226
227 @_polymorphic
228 def Ne(left, right):
229 return ~Eq(left, right)
230
231 @_polymorphic
232 def Gt(left, right):
233 return Polyhedron([], [left - right - 1])
234
235 @_polymorphic
236 def Ge(left, right):
237 return Polyhedron([], [left - right])
238
239
240 Empty = Eq(1, 0)
241
242 Universe = Polyhedron([])