Pyomo: implement build_maxcut_model; add support for quadratic objectives

miplearn/problems/maxcut.py

@@ -8,13 +8,14 @@ from typing import List, Union, Optional, Any
 import gurobipy as gp
 import networkx as nx
 import numpy as np
-from gurobipy import quicksum
+import pyomo.environ as pe
 from networkx import Graph
 from scipy.stats.distributions import rv_frozen
 
 from miplearn.io import read_pkl_gz
-from miplearn.problems import _gurobipy_set_params
+from miplearn.problems import _gurobipy_set_params, _pyomo_set_params
 from miplearn.solvers.gurobi import GurobiModel
+from miplearn.solvers.pyomo import PyomoModel
 
 
 @dataclass
@@ -80,6 +81,7 @@ class MaxCutGenerator:
     def _generate_graph(self) -> Graph:
         return nx.generators.random_graphs.binomial_graph(self.n.rvs(), self.p.rvs())
 
+
 def build_maxcut_model_gurobipy(
     data: Union[str, MaxCutData],
     params: Optional[dict[str, Any]] = None,
@@ -97,13 +99,47 @@ def build_maxcut_model_gurobipy(
     x = model.addVars(nodes, vtype=gp.GRB.BINARY, name="x")
 
     # Add the objective function
-    model.setObjective(quicksum(
-        - data.weights[i] * x[e[0]] * (1 - x[e[1]]) for (i, e) in enumerate(edges)
-    ))
+    model.setObjective(
+        gp.quicksum(
+            -data.weights[i] * x[e[0]] * (1 - x[e[1]]) for (i, e) in enumerate(edges)
+        )
+    )
 
     model.update()
     return GurobiModel(model)
 
+
+def build_maxcut_model_pyomo(
+    data: Union[str, MaxCutData],
+    solver: str = "gurobi_persistent",
+    params: Optional[dict[str, Any]] = None,
+) -> PyomoModel:
+    # Initialize model
+    model = pe.ConcreteModel()
+
+    # Read data
+    data = _maxcut_read(data)
+    nodes = pe.Set(initialize=list(data.graph.nodes))
+    edges = list(data.graph.edges())
+
+    # Add decision variables
+    model.x = pe.Var(nodes, domain=pe.Binary, name="x")
+
+    # Add the objective function
+    model.obj = pe.Objective(
+        expr=pe.quicksum(
+            -data.weights[i] * model.x[e[0]]
+            + data.weights[i] * model.x[e[0]] * model.x[e[1]]
+            for (i, e) in enumerate(edges)
+        ),
+        sense=pe.minimize,
+    )
+    model.pprint()
+    pm = PyomoModel(model, solver)
+    _pyomo_set_params(model, params, solver)
+    return pm
+
+
 def _maxcut_read(data: Union[str, MaxCutData]) -> MaxCutData:
     if isinstance(data, str):
         data = read_pkl_gz(data)

miplearn/solvers/gurobi.py

@@ -264,6 +264,13 @@ class GurobiModel(AbstractModel):
             h5.put_array(
                 h5_field, np.array(self.inner.getAttr(gp_field, gp_vars), dtype=float)
             )
+        obj = self.inner.getObjective()
+        if isinstance(obj, gp.QuadExpr):
+            nvars = len(self.inner.getVars())
+            obj_q = np.zeros((nvars, nvars))
+            for i in range(obj.size()):
+                obj_q[obj.getVar1(i).index, obj.getVar2(i).index] = obj.getCoeff(i)
+            h5.put_array("static_var_obj_coeffs_quad", obj_q)
 
     def _extract_after_load_constrs(self, h5: H5File) -> None:
         gp_constrs = self.inner.getConstrs()

miplearn/solvers/pyomo.py

@@ -9,6 +9,7 @@ import pyomo
 import pyomo.environ as pe
 from pyomo.core import Objective, Var, Suffix
 from pyomo.core.base import VarData
+from pyomo.core.expr import ProductExpression
 from pyomo.core.expr.numeric_expr import SumExpression, MonomialTermExpression
 from scipy.sparse import coo_matrix
 
@@ -207,19 +208,23 @@ class PyomoModel(AbstractModel):
         lower_bounds: List[float] = []
         obj_coeffs: List[float] = []
 
-        obj = None
+        obj_quad, obj_linear = None, None
         obj_offset = 0.0
         obj_count = 0
         for obj in self.inner.component_objects(Objective):
-            obj, obj_offset = self._parse_pyomo_expr(obj.expr)
+            obj_quad, obj_linear, obj_offset = self._parse_obj_expr(obj.expr)
             obj_count += 1
         assert obj_count == 1, f"One objective function expected; found {obj_count}"
+        assert obj_quad is not None
+        assert obj_linear is not None
 
+        varname_to_idx: Dict[str, int] = {}
         for i, var in enumerate(self.inner.component_objects(pyomo.core.Var)):
             for idx in var:
                 v = var[idx]
 
                 # Variable name
+                varname_to_idx[v.name] = len(names)
                 if idx is None:
                     names.append(var.name)
                 else:
@@ -249,11 +254,22 @@ class PyomoModel(AbstractModel):
                 lower_bounds.append(float(lb))
 
                 # Objective coefficients
-                if v.name in obj:
-                    obj_coeffs.append(obj[v.name])
+                if v.name in obj_linear:
+                    obj_coeffs.append(obj_linear[v.name])
                 else:
                     obj_coeffs.append(0.0)
 
+        if len(obj_quad) > 0:
+            nvars = len(names)
+            matrix = np.zeros((nvars, nvars))
+            for ((left_varname, right_varname), coeff) in obj_quad.items():
+                assert left_varname in varname_to_idx
+                assert right_varname in varname_to_idx
+                left_idx = varname_to_idx[left_varname]
+                right_idx = varname_to_idx[right_varname]
+                matrix[left_idx, right_idx] = coeff
+                h5.put_array("static_var_obj_coeffs_quad", matrix)
+
         h5.put_array("static_var_names", np.array(names, dtype="S"))
         h5.put_array("static_var_types", np.array(types, dtype="S"))
         h5.put_array("static_var_lower_bounds", np.array(lower_bounds))
@@ -327,8 +343,9 @@ class PyomoModel(AbstractModel):
                 _parse_constraint(constr, curr_row)
                 curr_row += 1
 
-        lhs = coo_matrix((lhs_data, (lhs_row, lhs_col))).tocoo()
-        h5.put_sparse("static_constr_lhs", lhs)
+        if len(lhs_data) > 0:
+            lhs = coo_matrix((lhs_data, (lhs_row, lhs_col))).tocoo()
+            h5.put_sparse("static_constr_lhs", lhs)
         h5.put_array("static_constr_names", np.array(names, dtype="S"))
         h5.put_array("static_constr_rhs", np.array(rhs))
         h5.put_array("static_constr_sense", np.array(senses, dtype="S"))
@@ -372,24 +389,47 @@ class PyomoModel(AbstractModel):
                     slacks.append(abs(self.inner.slack[c]))
         h5.put_array("mip_constr_slacks", np.array(slacks))
 
-    def _parse_pyomo_expr(self, expr: Any) -> Tuple[Dict[str, float], float]:
-        lhs = {}
-        offset = 0.0
+    def _parse_term(self, t: Any) -> Tuple[str, float]:
+        if isinstance(t, MonomialTermExpression):
+            return t._args_[1].name, float(t._args_[0])
+        elif isinstance(t, VarData):
+            return t.name, 1.0
+        else:
+            raise Exception(f"Unknown term type: {t.__class__.__name__}")
+
+    def _parse_obj_expr(
+        self, expr: Any
+    ) -> Tuple[Dict[Tuple[str, str], float], Dict[str, float], float]:
+        obj_coeff_linear = {}
+        obj_coeff_quadratic = {}
+        obj_offset = 0.0
         if isinstance(expr, SumExpression):
             for term in expr._args_:
-                if isinstance(term, MonomialTermExpression):
-                    lhs[term._args_[1].name] = float(term._args_[0])
-                elif isinstance(term, VarData):
-                    lhs[term.name] = 1.0
-                elif isinstance(term, float):
-                    offset += term
+                if isinstance(term, (int, float)):
+                    # Constant term
+                    obj_offset += term
+                elif isinstance(term, (MonomialTermExpression, VarData)):
+                    # Linear term
+                    var_name, var_coeff = self._parse_term(term)
+                    if var_name not in obj_coeff_linear:
+                        obj_coeff_linear[var_name] = 0.0
+                    obj_coeff_linear[var_name] += var_coeff
+                elif isinstance(term, ProductExpression):
+                    # Quadratic terms
+                    left_var_nane, left_coeff = self._parse_term(term._args_[0])
+                    right_var_nane, right_coeff = self._parse_term(term._args_[1])
+                    if (left_var_nane, right_var_nane) not in obj_coeff_quadratic:
+                        obj_coeff_quadratic[(left_var_nane, right_var_nane)] = 0.0
+                    obj_coeff_quadratic[(left_var_nane, right_var_nane)] += (
+                        left_coeff * right_coeff
+                    )
                 else:
                     raise Exception(f"Unknown term type: {term.__class__.__name__}")
         elif isinstance(expr, VarData):
-            lhs[expr.name] = 1.0
+            obj_coeff_linear[expr.name] = 1.0
         else:
             raise Exception(f"Unknown expression type: {expr.__class__.__name__}")
-        return lhs, offset
+        return obj_coeff_quadratic, obj_coeff_linear, obj_offset
 
     def _gap(self, zp: float, zd: float, tol: float = 1e-6) -> float:
         # Reference: https://www.gurobi.com/documentation/9.5/refman/mipgap2.html

tests/problems/test_maxcut.py

@@ -1,17 +1,26 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2025, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+
 import random
+from tempfile import TemporaryDirectory
 
 import numpy as np
-
-from miplearn.problems.maxcut import MaxCutGenerator, build_maxcut_model_gurobipy
 from scipy.stats import randint, uniform
 
+from miplearn.h5 import H5File
+from miplearn.problems.maxcut import (
+    MaxCutGenerator,
+    build_maxcut_model_gurobipy,
+    build_maxcut_model_pyomo,
+)
+
+
 def _set_seed():
     random.seed(42)
     np.random.seed(42)
 
+
 def test_maxcut_generator_not_fixed() -> None:
     _set_seed()
     gen = MaxCutGenerator(
@@ -22,12 +31,20 @@ def test_maxcut_generator_not_fixed() -> None:
     data = gen.generate(3)
     assert len(data) == 3
     assert list(data[0].graph.nodes()) == [0, 1, 2, 3, 4]
-    assert list(data[0].graph.edges()) == [(0, 2), (0, 3), (0, 4), (2, 3), (2, 4), (3, 4)]
+    assert list(data[0].graph.edges()) == [
+        (0, 2),
+        (0, 3),
+        (0, 4),
+        (2, 3),
+        (2, 4),
+        (3, 4),
+    ]
     assert data[0].weights.tolist() == [-1, 1, -1, -1, -1, 1]
     assert list(data[1].graph.nodes()) == [0, 1, 2, 3, 4]
     assert list(data[1].graph.edges()) == [(0, 1), (0, 3), (0, 4), (1, 4), (3, 4)]
     assert data[1].weights.tolist() == [-1, -1, -1, 1, -1]
 
+
 def test_maxcut_generator_fixed() -> None:
     random.seed(42)
     np.random.seed(42)
@@ -39,19 +56,70 @@ def test_maxcut_generator_fixed() -> None:
     data = gen.generate(3)
     assert len(data) == 3
     assert list(data[0].graph.nodes()) == [0, 1, 2, 3, 4]
-    assert list(data[0].graph.edges()) == [(0, 2), (0, 3), (0, 4), (2, 3), (2, 4), (3, 4)]
+    assert list(data[0].graph.edges()) == [
+        (0, 2),
+        (0, 3),
+        (0, 4),
+        (2, 3),
+        (2, 4),
+        (3, 4),
+    ]
     assert data[0].weights.tolist() == [-1, 1, -1, -1, -1, 1]
     assert list(data[1].graph.nodes()) == [0, 1, 2, 3, 4]
-    assert list(data[1].graph.edges()) == [(0, 2), (0, 3), (0, 4), (2, 3), (2, 4), (3, 4)]
+    assert list(data[1].graph.edges()) == [
+        (0, 2),
+        (0, 3),
+        (0, 4),
+        (2, 3),
+        (2, 4),
+        (3, 4),
+    ]
     assert data[1].weights.tolist() == [-1, -1, -1, 1, -1, -1]
 
+
 def test_maxcut_model():
     _set_seed()
     data = MaxCutGenerator(
-        n=randint(low=20, high=21),
+        n=randint(low=10, high=11),
         p=uniform(loc=0.5, scale=0.0),
         fix_graph=True,
     ).generate(1)[0]
-    model = build_maxcut_model_gurobipy(data)
-    model.optimize()
-    assert model.inner.ObjVal == -26
+    for build_model in [
+        build_maxcut_model_gurobipy,
+        build_maxcut_model_pyomo,
+    ]:
+        model = build_model(data)
+        with TemporaryDirectory() as tempdir:
+            with H5File(f"{tempdir}/data.h5", "w") as h5:
+                model.extract_after_load(h5)
+                obj_lin = h5.get_array("static_var_obj_coeffs")
+                assert obj_lin is not None
+                assert obj_lin.tolist() == [
+                    3.0,
+                    1.0,
+                    3.0,
+                    1.0,
+                    -1.0,
+                    0.0,
+                    -1.0,
+                    0.0,
+                    -1.0,
+                    0.0,
+                ]
+                obj_quad = h5.get_array("static_var_obj_coeffs_quad")
+                assert obj_quad is not None
+                assert obj_quad.tolist() == [
+                    [0.0, 0.0, -1.0, 1.0, -1.0, 0.0, 0.0, 0.0, -1.0, -1.0],
+                    [0.0, 0.0, 1.0, -1.0, 0.0, -1.0, -1.0, 0.0, 0.0, 1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, -1.0, -1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, -1.0, 1.0, -1.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, -1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                    [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                ]
+                model.optimize()
+                model.extract_after_mip(h5)
+                assert h5.get_scalar("mip_obj_value") == -4