from itertools import count id_counter = count(start=1) def _get_var(): return f"x_{next(id_counter)}" def _is_number(x) -> bool: """Check if x is a number.""" try: float(x) return True except ValueError: return False except TypeError: return False class Variable: # def __init__(self, scip_var): def __init__(self): # self.scip_var = scip_var self.scip_var = _get_var() # property name: # def __get__(self): # return bytes(SCIPvarGetName(self.scip_var)).decode("utf-8") def ptr(self): # return <size_t>(self.scip_var) return self.scip_var def __repr__(self): return str(self.scip_var) # return self.name def __add__(self, other): return self.to_expr().__add__(other) def __iadd__(self, other): self = self.__add__(other) return self def __radd__(self, other): return self.to_expr().__radd__(other) def __mul__(self, other): return self.to_expr().__mul__(other) def __rmul__(self, other): return self.to_expr().__rmul__(other) def __truediv__(self, other): return self.to_expr().__truediv__(other) def __rtruediv__(self, other): return self.to_expr().__rtruediv__(other) def __pow__(self, other): return self.to_expr().__pow__(other) def __neg__(self): return self.to_expr().__neg__() def __sub__(self, other): return self.to_expr().__sub__(other) def __rsub__(self, other): return self.to_expr().__rsub__(other) def __richcmp__(self, other, op: int): return self.to_expr().__richcmp__(other, op) def to_expr(self): return MonomialExpr.from_var(self) class Term: """A monomial term consisting of one or more variables.""" __slots__ = ("vars", "ptrs") def __init__(self, *vars): self.vars = tuple(sorted(vars, key=lambda v: v.ptr())) self.ptrs = tuple(v.ptr() for v in self.vars) def __getitem__(self, idx): return self.vars[idx] def __hash__(self): return self.ptrs.__hash__() def __eq__(self, other): return self.ptrs == other.ptrs def __len__(self): return len(self.vars) def __mul__(self, other): if not isinstance(other, Term): raise TypeError( f"unsupported operand type(s) for *: 'Term' and '{type(other)}'" ) return Term(*self.vars, *other.vars) def __repr__(self): return f"Term({', '.join(map(str, self.vars))})" CONST = Term() class Expr: """Expression representing for `expression1 + expression2 + constant`.""" # cdef public children def __init__(self, children: dict = None): self.children = children or {} def __hash__(self): return frozenset(self.children.items()).__hash__() def __getitem__(self, key): return self.children.get(key, 0.0) def __iter__(self): return iter(self.children) def __next__(self): try: return next(self.children) except: raise StopIteration def __add__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, Expr): return SumExpr({self: 1.0, other: 1.0}) # elif isinstance(other, MatrixExpr): # return other.__add__(self) raise TypeError( f"unsupported operand type(s) for +: 'Expr' and '{type(other)}'" ) def __mul__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, Expr): return ProdExpr(self, other) # elif isinstance(other, MatrixExpr): # return other.__mul__(self) raise TypeError( f"unsupported operand type(s) for *: 'Expr' and '{type(other)}'" ) def __truediv__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, ConstExpr) and other[CONST] == 0: raise ZeroDivisionError("division by zero") if hash(self) == hash(other): return ConstExpr(1.0) return self.__mul__(other.__pow__(-1.0)) def __rtruediv__(self, other): return Expr.to_const_or_var(other).__truediv__(self) def __pow__(self, other): other = Expr.to_const_or_var(other) if not isinstance(other, ConstExpr): raise TypeError("exponent must be a number") if other[CONST] == 0: return ConstExpr(1.0) return PowExpr(self, other[CONST]) def __sub__(self, other): return self.__add__(-other) def __neg__(self): return self.__mul__(-1.0) def __iadd__(self, other): self = self.__add__(other) return self def __radd__(self, other): return self.__add__(other) def __rmul__(self, other): return self.__mul__(other) def __rsub__(self, other): return self.__neg__().__add__(other) def __richcmp__(self, other, op: int): other = Expr.to_const_or_var(other) if op == 1: # <= if isinstance(other, Expr): if isinstance(other, ConstExpr): return ExprCons(self, rhs=other[CONST]) return (self - other) <= 0 # elif isinstance(other, MatrixExpr): # return self.__richcmp__(other, 5) raise TypeError(f"Unsupported type {type(other)}") elif op == 5: # >= if isinstance(other, Expr): if isinstance(other, ConstExpr): return ExprCons(self, lhs=other[CONST]) return (self - other) >= 0 # elif isinstance(other, MatrixExpr): # return self.__richcmp__(other, 1) raise TypeError(f"Unsupported type {type(other)}") elif op == 2: # == if isinstance(other, Expr): if isinstance(other, ConstExpr): return ExprCons(self, lhs=other[CONST], rhs=other[CONST]) return (self - other) == 0 # elif isinstance(other, MatrixExpr): # return self.__richcmp__(other, 2) raise TypeError(f"Unsupported type {type(other)}") raise NotImplementedError( "Can only support constraints with '<=', '>=', or '=='." ) def __repr__(self): return f"Expr({self.children})" def degree(self): """Computes the highest degree of children""" return max(map(len, self.children)) if self.children else 0 @staticmethod def to_const_or_var(x): """Convert a number or variable to an expression.""" if _is_number(x): return PolynomialExpr.to_subclass({CONST: x}) elif isinstance(x, Variable): return PolynomialExpr.to_subclass({Term(x): 1.0}) return x def to_dict(self, other: dict = None) -> dict: """Merge two dictionaries by summing values of common keys""" other = other or {} if not isinstance(other, dict): raise TypeError("other must be a dict") res = self.children.copy() for child, coef in other.items(): res[child] = res.get(child, 0.0) + coef return res class SumExpr(Expr): def __add__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, SumExpr): return SumExpr(self.to_dict(other.children)) return super().__add__(other) def __mul__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, ConstExpr): if other[CONST] == 0: return ConstExpr(0.0) return SumExpr({i: self[i] * other[CONST] for i in self if self[i] != 0}) return super().__mul__(other) class PolynomialExpr(SumExpr): """Expression representing for `2*x**3 + 4*x*y + constant`.""" def __init__(self, children: dict = None): if children and not all(isinstance(t, Term) for t in children): raise TypeError("All keys must be Term instances") super().__init__(children) def __add__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, PolynomialExpr): return PolynomialExpr.to_subclass(self.to_dict(other.children)) return super().__add__(other) def __mul__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, PolynomialExpr): children = {} for i in self: for j in other: child = i * j children[child] = children.get(child, 0.0) + self[i] * other[j] return PolynomialExpr.to_subclass(children) return super().__mul__(other) def __truediv__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, ConstExpr): return self.__mul__(1.0 / other[CONST]) return super().__truediv__(other) def __pow__(self, other): other = Expr.to_const_or_var(other) if ( isinstance(other, Expr) and isinstance(other, ConstExpr) and other[CONST].is_integer() and other[CONST] > 0 ): res = 1 for _ in range(int(other[CONST])): res *= self return res return super().__pow__(other) @classmethod def to_subclass(cls, children: dict = None): if len(children) == 0: return ConstExpr(0.0) elif len(children) == 1: if CONST in children: return ConstExpr(children[CONST]) return MonomialExpr(children) return cls(children) class ConstExpr(PolynomialExpr): """Expression representing for `constant`.""" def __init__(self, constant: float = 0): super().__init__({CONST: constant}) def __pow__(self, other, modulo): other = Expr.to_const_or_var(other) if isinstance(other, ConstExpr): return ConstExpr(self[CONST] ** other[CONST]) return super().__pow__(other, modulo) class MonomialExpr(PolynomialExpr): """Expression representing for `x**3`.""" def __init__(self, children: dict = None): if children and len(children) != 1: raise ValueError("MonomialExpr must have exactly one child") super().__init__(children) @staticmethod def from_var(var: Variable, coef: float = 1.0): return MonomialExpr({Term(var): coef}) class FuncExpr(Expr): def degree(self): return float("inf") class ProdExpr(FuncExpr): """Expression representing for `coefficient * expression`.""" def __init__(self, *children, coef: float = 1.0): super().__init__({i: 1.0 for i in children}) self.coef = coef def __hash__(self): return (frozenset(self), self.coef).__hash__() def __add__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, ProdExpr) and hash(frozenset(self)) == hash( frozenset(other) ): return ProdExpr(*self, coef=self.coef + other.coef) return super().__add__(other) def __mul__(self, other): other = Expr.to_const_or_var(other) if isinstance(other, ConstExpr): if other[CONST] == 0: return ConstExpr(0.0) return ProdExpr(*self, coef=self.coef * other[CONST]) return super().__mul__(other) def __repr__(self): return f"ProdExpr({{{tuple(self)}: {self.coef}}})" def _normalize(self): if self.coef == 0: return ConstExpr(0.0) return self class PowExpr(FuncExpr): """Expression representing for `pow(expression, exponent)`.""" def __init__(self, base, expo: float = 1.0): super().__init__({base: 1.0}) self.expo = expo def __hash__(self): return (frozenset(self), self.expo).__hash__() def __repr__(self): return f"PowExpr({tuple(self)}, {self.expo})" def _normalize(self): if self.expo == 0: self = ConstExpr(1.0) elif self.expo == 1: self = tuple(self)[0] class UnaryExpr(FuncExpr): def __init__(self, expr): super().__init__((expr,), (1.0,)) def __hash__(self): return frozenset(self).__hash__() def __repr__(self): return f"{type(self).__name__}({tuple(self)[0]})" class ExpExpr(UnaryExpr): """Expression representing for `exp(expression)`.""" class LogExpr(UnaryExpr): """Expression representing for `log(expression)`.""" class SqrtExpr(UnaryExpr): """Expression representing for `sqrt(expression)`.""" class SinExpr(UnaryExpr): """Expression representing for `sin(expression)`.""" class CosExpr(UnaryExpr): """Expression representing for `cos(expression)`.""" # cdef class ExprCons: class ExprCons: """Constraints with a polynomial expressions and lower/upper bounds.""" # cdef public expr # cdef public _lhs # cdef public _rhs def __init__(self, expr, lhs=None, rhs=None): self.expr = expr self._lhs = lhs self._rhs = rhs assert not (lhs is None and rhs is None) self.normalize() def normalize(self): """move constant children in expression to bounds""" if isinstance(self.expr, Expr): c = self.expr[CONST] self.expr -= c assert self.expr[CONST] == 0.0 self.expr.normalize() else: return if self._lhs is not None: self._lhs -= c if self._rhs is not None: self._rhs -= c def __richcmp__(self, other, op): if op == 1: # <= if self._rhs is not None: raise TypeError("ExprCons already has upper bound") assert self._lhs is not None if not _is_number(other): raise TypeError("Ranged ExprCons is not well defined!") return ExprCons(self.expr, lhs=self._lhs, rhs=float(other)) elif op == 5: # >= if self._lhs is not None: raise TypeError("ExprCons already has lower bound") assert self._lhs is None assert self._rhs is not None if not _is_number(other): raise TypeError("Ranged ExprCons is not well defined!") return ExprCons(self.expr, lhs=float(other), rhs=self._rhs) else: raise NotImplementedError( "Ranged ExprCons can only support with '<=' or '>='." ) def __repr__(self): return "ExprCons(%s, %s, %s)" % (self.expr, self._lhs, self._rhs) def __bool__(self): """Make sure that equality of expressions is not asserted with ==""" msg = """Can't evaluate constraints as booleans.  If you want to add a ranged constraint of the form:  lhs <= expression <= rhs you have to use parenthesis to break the Python syntax for chained comparisons:  lhs <= (expression <= rhs) """ raise TypeError(msg)