Support Gurobi expression trees

cpmpy/solvers/gurobi.py

@@ -42,19 +42,23 @@
     ==============
 """
 
+import cpmpy as cp
 from typing import Optional, List
 import warnings
 
+import math
+
 from .solver_interface import SolverInterface, SolverStatus, ExitStatus, Callback
 from ..exceptions import NotSupportedError
 from ..expressions.core import Expression, Comparison, Operator, BoolVal
 from ..expressions.utils import argvals, argval, is_any_list, is_num
 from ..expressions.variables import _BoolVarImpl, NegBoolView, _IntVarImpl, _NumVarImpl, intvar
 from ..expressions.globalconstraints import DirectConstraint
 from ..transformations.comparison import only_numexpr_equality
-from ..transformations.flatten_model import flatten_constraint, flatten_objective
+from ..transformations.flatten_model import flatten_constraint, flatten_objective, get_or_make_var_or_list
 from ..transformations.get_variables import get_variables
 from ..transformations.linearize import linearize_constraint, linearize_reified_variables, only_positive_bv, only_positive_bv_wsum, decompose_linear, decompose_linear_objective
+from ..transformations.decompose_global import decompose_in_tree
 from ..transformations.normalize import toplevel_list
 from ..transformations.reification import only_implies, reify_rewrite, only_bv_reifies
 from ..transformations.safening import no_partial_functions, safen_objective
@@ -192,6 +196,9 @@ def solve(self, time_limit:Optional[float]=None, solution_callback=None, **kwarg
         for param, val in kwargs.items():
             self.grb_model.setParam(param, val)
 
+        # write LP file for debugging
+        self.grb_model.write("/tmp/model.lp")
+
         _ = self.grb_model.optimize(callback=solution_callback)
         grb_objective = self.grb_model.getObjective()
 
@@ -259,7 +266,7 @@ def solver_var(self, cpm_var):
 
         # special case, negative-bool-view. Should be eliminated in linearize
         if isinstance(cpm_var, NegBoolView):
-            raise NotSupportedError("Negative literals should not be left as part of any equation. Please report.")
+            return 1 - self.solver_var(cpm_var._bv)
 
         # create if it does not exit
         if cpm_var not in self._varmap:
@@ -297,7 +304,7 @@ def objective(self, expr, minimize=True):
                                                       supported=self.supported_global_constraints,
                                                       supported_reified=self.supported_reified_global_constraints,
                                                       csemap=self._csemap)
-        obj, flat_cons = flatten_objective(obj, csemap=self._csemap)
+        obj, flat_cons = flatten_objective(obj, csemap=self._csemap, supported={"pow", "mul"})
         obj = only_positive_bv_wsum(obj)  # remove negboolviews
 
         self.add(safe_cons + decomp_cons + flat_cons)
@@ -358,20 +365,37 @@ def transform(self, cpm_expr):
         # apply transformations, then post internally
         # expressions have to be linearized to fit in MIP model. See /transformations/linearize
         cpm_cons = toplevel_list(cpm_expr)
-        cpm_cons = no_partial_functions(cpm_cons, safen_toplevel={"mod", "div", "element"})  # linearize and decompose expect safe exprs
+        cpm_cons = no_partial_functions(cpm_cons, safen_toplevel=frozenset(["mod", "div", "element"]))  # linearize and decompose expect safe exprs
         cpm_cons = decompose_linear(cpm_cons,
                                     supported=self.supported_global_constraints,
                                     supported_reified=self.supported_reified_global_constraints,
                                     csemap=self._csemap)
-        cpm_cons = flatten_constraint(cpm_cons, csemap=self._csemap)  # flat normal form
+        cpm_cons = flatten_constraint(cpm_cons, csemap=self._csemap, supported=frozenset(["mul", "pow", "-", "sum"]))  # flat normal form
         cpm_cons = reify_rewrite(cpm_cons, supported=frozenset(['sum', 'wsum']), csemap=self._csemap)  # constraints that support reification
         cpm_cons = only_numexpr_equality(cpm_cons, supported=frozenset(["sum", "wsum", "sub"]), csemap=self._csemap)  # supports >, <, !=
-        cpm_cons = linearize_reified_variables(cpm_cons, min_values=2, csemap=self._csemap)
+
+        # cpm_cons = linearize_reified_variables(cpm_cons, min_values=2, csemap=self._csemap)
         cpm_cons = only_bv_reifies(cpm_cons, csemap=self._csemap)
-        cpm_cons = only_implies(cpm_cons, csemap=self._csemap)  # anything that can create full reif should go above...
+        cpm_cons = only_implies(
+            cpm_cons,
+            csemap=self._csemap,
+            is_supported=lambda cpm_expr:
+                (
+                    cpm_expr.name == "==" and
+                    (
+                        (isinstance(cpm_expr.args[1], Operator) and cpm_expr.args[1].name in {"or", "and"}) or
+                        isinstance(cpm_expr.args[1], _BoolVarImpl)
+                    )
+                ) or (
+                    cpm_expr.name == "->" and isinstance(cpm_expr.args[1], Comparison)
+                )
+        )  # anything that can create full reif should go above...
+
         # gurobi does not round towards zero, so no 'div' in supported set: https://github.com/CPMpy/cpmpy/pull/593#issuecomment-2786707188
-        cpm_cons = linearize_constraint(cpm_cons, supported=frozenset({"sum", "wsum","->","sub","min","max","mul","abs","pow"}), csemap=self._csemap)  # the core of the MIP-linearization
+        cpm_cons = linearize_constraint(cpm_cons, supported=frozenset({"-","sum", "wsum","->","sub","min","max","mul","abs","pow","and","or"}), csemap=self._csemap)  # the core of the MIP-linearization
+
         cpm_cons = only_positive_bv(cpm_cons, csemap=self._csemap)  # after linearization, rewrite ~bv into 1-bv
+        # TODO don't rewrite ~p -> LinCons
         return cpm_cons
 
     def add(self, cpm_expr_orig):
@@ -392,13 +416,102 @@ def add(self, cpm_expr_orig):
 
         :return: self
       """
-      from gurobipy import GRB
+      from gurobipy import GRB, nlfunc
+      import gurobipy as gp
+
+      def add(cpm_expr):
+          """Recursively create a Gurobi constraint from a CPMpy expression."""
+
+          def add_(cpm_expr, depth):
+              indent = "  " * depth
+              depth+=1
+              self.grb_model.update()
+              print(f"{indent}Con:", cpm_expr, type(cpm_expr))
+
+              if is_num(cpm_expr):
+                  return int(cpm_expr)
+              elif isinstance(cpm_expr, NegBoolView):
+                  return 1 - self.solver_var(cpm_expr._bv)
+              elif isinstance(cpm_expr, _NumVarImpl):
+                  return self.solver_var(cpm_expr)
+              elif isinstance(cpm_expr, Operator):
+                  match cpm_expr.name:
+                      case "or": # TODO only usefull if we can represent an or as a sum of binary vars, which is not generally correct
+                          return gp.or_(add_(arg, depth) for arg in cpm_expr.args)
+                      case "and":
+                          return gp.and_(add_(arg, depth) for arg in cpm_expr.args)
+                      case "->":  # Gurobi indicator constraint: (Var == 0|1) >> (LinExpr sense LinExpr)
+                          a, b = cpm_expr.args
+                          assert isinstance(a, _BoolVarImpl), f"Implication constraint {cpm_expr} must have BoolVar as lhs, but had {a}"
+                          assert isinstance(b, Comparison),  f"Implication constraint {cpm_expr} must have Comparison as rhs, but had {b}"
+                          is_pos = not isinstance(a, NegBoolView)
+                          return (add_(a if is_pos else a._bv, depth) == int(is_pos)) >> add_(b, depth)
+                      case "not":
+                          return 1 - add_(cpm_expr.args[0], depth)
+                      case "-":
+                          return -add_(cpm_expr.args[0], depth=depth)
+                      case "sum":
+                          return sum(add_(arg, depth) for arg in cpm_expr.args)
+                      case "wsum":
+                          return sum(weight * add_(arg, depth) for weight, arg in zip(cpm_expr.args[0], cpm_expr.args[1]))
+                      case "sub":
+                          return add_(cpm_expr.args[0], depth) - add_(cpm_expr.args[1], depth)
+              elif isinstance(cpm_expr, Comparison):
+                  a, b = add_(cpm_expr.args[0], depth), add_(cpm_expr.args[1], depth)
+                  match cpm_expr.name:
+                      case "==":
+                          if isinstance(a, gp.NLExpr):
+                              # if flattening led to a non-linear expression, then it has to be constraint `y == f(x)` with `y` a `Var`
+                              return self.grb_model.addVar(lb=b, ub=b) == a
+                          else:
+                              # Else, this is a function constraint 
+                              # Note: Gurobi functions are called by `y == f(x)`, like normalized CPMpy boolexprs, but CPMpy numexprs are normalized to `f(x) == y` (e.g. `abs(x) == IV0`), so they need to be flipped
+                              return a == b if cpm_expr.args[0].is_bool() else b == a
+                      case "<=":
+                          return a <= b
+                      case ">=":
+                          return a >= b
+                      case _:
+                          raise Exception(f"Expected comparator to be ==,<=,>= in Comparison expression {cpm_expr}, but was {cpm_expr.name}")
+              elif isinstance(cpm_expr, cp.expressions.globalfunctions.GlobalFunction):
+                  match cpm_expr.name:
+                      case "mul":
+                          return add_(cpm_expr.args[0], depth) * add_(cpm_expr.args[1], depth)
+                      case "pow":
+                          return add_(cpm_expr.args[0], depth) ** add_(cpm_expr.args[1], depth)
+                      case _:  # other global functions cannot be 
+                          match cpm_expr.name:
+                              case "abs":  # y = abs(x)
+                                  # TODO we could support this inside the expression tree with sqrt(pow(x,2))?
+                                  return gp.abs_(add_(cpm_expr.args[0], depth))
+                              case "min" | "max":
+                                  # TODO outdated ; should be no need to make our own var
+                                  y = add_(cp.intvar(lb=min(x.lb for x in cpm_expr.args), ub=max(x.ub for x in cpm_expr.args)), depth)
+                                  flatargs = [add_(arg, depth) for arg in cpm_expr.args]
+                                  self.native_model.addConstr(y == (gp.min_(flatargs) if cpm_expr.name == "min" else gp.max_(flatargs)))
+                                  return y
+              else:
+                  raise NotImplementedError(f"add_() not implemented for {cpm_expr}, {type(cpm_expr)}, {getattr(cpm_expr, 'name', None)}")
+
+          grb_expr = add_(cpm_expr, 0)
+          if isinstance(grb_expr, (gp.Var, gp.LinExpr)):
+              # If add() returned a Gurobi Var (not a constraint), wrap it as >= 1
+              return grb_expr >= 1
+          elif isinstance(grb_expr, gp.TempConstr):
+              return grb_expr
+          else:
+              return self.grb_model.addVar(lb=1, ub=1) == grb_expr
 
       # add new user vars to the set
       get_variables(cpm_expr_orig, collect=self.user_vars)
 
       # transform and post the constraints
       for cpm_expr in self.transform(cpm_expr_orig):
+        con = add(cpm_expr)
+        self.grb_model.update()
+        print("out", con)
+        self.grb_model.addConstr(con)
+        continue
 
         # Comparisons: only numeric ones as 'only_implies()' has removed the '==' reification for Boolean expressions
         # numexpr `comp` bvar|const
@@ -487,8 +600,11 @@ def add(self, cpm_expr_orig):
         elif isinstance(cpm_expr, DirectConstraint):
             cpm_expr.callSolver(self, self.grb_model)
 
+        elif isinstance(cpm_expr, _BoolVarImpl):
+            self.grb_model.addConstr(self.solver_var(cpm_expr) >= 1)
+
         else:
-            raise NotImplementedError(cpm_expr)  # if you reach this... please report on github
+            raise NotImplementedError(f"Please report unsupported constraint in Gurobi interface: {cpm_expr} of type {type(cpm_expr)}")  # if you reach this... please report on github
 
       return self
     __add__ = add  # avoid redirect in superclass