pyepo 2.2.0__tar.gz → 2.2.1__tar.gz

{pyepo-2.2.0 → pyepo-2.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.0
+Version: 2.2.1
 Summary: PyTorch-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT
@@ -31,7 +31,6 @@ Requires-Dist: pathos
 Requires-Dist: tqdm
 Requires-Dist: scikit_learn
 Requires-Dist: torch>=1.13.1
-Requires-Dist: clarabel
 Provides-Extra: pyomo
 Requires-Dist: pyomo>=6.1.2; extra == "pyomo"
 Provides-Extra: gurobi
@@ -40,6 +39,8 @@ Provides-Extra: copt
 Requires-Dist: coptpy; extra == "copt"
 Provides-Extra: ortools
 Requires-Dist: ortools>=9.6; extra == "ortools"
+Provides-Extra: cave
+Requires-Dist: clarabel; extra == "cave"
 Provides-Extra: mpax
 Requires-Dist: mpax; extra == "mpax"
 Requires-Dist: jax>=0.4.1; extra == "mpax"
@@ -51,6 +52,7 @@ Provides-Extra: autosklearn
 Requires-Dist: auto-sklearn; extra == "autosklearn"
 Requires-Dist: configspace; extra == "autosklearn"
 Provides-Extra: all
+Requires-Dist: clarabel; extra == "all"
 Requires-Dist: pyomo>=6.1.2; extra == "all"
 Requires-Dist: gurobipy>=9.1.2; extra == "all"
 Requires-Dist: coptpy; extra == "all"

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/data/dataset.py

@@ -64,6 +64,11 @@ class optDataset(Dataset):
         """
         if not isinstance(model, optModel):
             raise TypeError("arg model is not an optModel")
+        if len(feats) != len(costs):
+            raise ValueError(
+                f"feats and costs must have the same number of instances: "
+                f"{len(feats)} vs {len(costs)}."
+            )
         self.model = model
         # data
         self.feats = feats
@@ -86,15 +91,10 @@ class optDataset(Dataset):
         objs = []
         logger.info("Optimizing for optDataset...")
         for c in tqdm(self.costs):
-            try:
-                sol, obj = self._solve(c)
-                # to numpy
-                if isinstance(sol, torch.Tensor):
-                    sol = sol.detach().cpu().numpy()
-            except Exception as e:
-                raise ValueError(
-                    "For optModel, the method 'solve' should return solution vector and objective value."
-                ) from e
+            sol, obj = self._solve(c)
+            # to numpy
+            if isinstance(sol, torch.Tensor):
+                sol = sol.detach().cpu().numpy()
             sols.append(np.asarray(sol))
             objs.append(obj)
         return np.stack(sols), np.asarray(objs).reshape(-1, 1)
@@ -107,15 +107,19 @@ class optDataset(Dataset):
         from pyepo.model.mpax.mpaxmodel import _warn_if_not_optimal
 
         model = cast("_optMpaxModelT", self.model)
-        model.setObj(model._fullCost(self.costs))
+        model._setFullObj(model._fullCost(self.costs))
         sols, objs, status = model.batch_optimize(model.c)
         _warn_if_not_optimal(status)
         # writable copy; torch.as_tensor warns on JAX read-only buffers
         sols_np = np.array(sols, dtype=np.float32)
         objs_np = np.array(objs, dtype=np.float32)
-        # jitted_solve returns c·sol where setObj already negated c for MAX
+        # jitted_solve returns c·sol where the objective write already negated c for MAX
         if self.model.modelSense == EPO.MAXIMIZE:
             objs_np = -objs_np
+        # compiled DSL problems carry bare objective constants outside the solver model
+        problem = getattr(self.model, "problem", None)
+        if problem is not None:
+            objs_np += problem.obj_offset
         return sols_np, objs_np.reshape(-1, 1)
 
     def _solve(
@@ -131,7 +135,7 @@ class optDataset(Dataset):
         Returns:
             tuple: optimal solution (np.ndarray) and objective value (float)
         """
-        self.model.setObj(self.model._fullCost(cost))
+        self.model._setFullObj(self.model._fullCost(cost))
         sol, obj = self.model.solve()
         return sol, obj
 
@@ -207,6 +211,11 @@ class optDatasetKNN(optDataset):
         """
         if not isinstance(model, optModel):
             raise TypeError("arg model is not an optModel")
+        if len(feats) != len(costs):
+            raise ValueError(
+                f"feats and costs must have the same number of instances: "
+                f"{len(feats)} vs {len(costs)}."
+            )
         self.model = model
         # at most num_data-1 neighbours exist (self excluded), so k must stay below it
         num_data = len(feats)
@@ -241,14 +250,9 @@ class optDatasetKNN(optDataset):
             sol_knn = np.zeros((self.costs.shape[1], self.k))
             obj_knn = np.zeros(self.k)
             for i, c in enumerate(c_knn.T):
-                try:
-                    sol_i, obj_i = self._solve(c)
-                    if isinstance(sol_i, torch.Tensor):
-                        sol_i = sol_i.detach().cpu().numpy()
-                except Exception as e:
-                    raise ValueError(
-                        "For optModel, the method 'solve' should return solution vector and objective value."
-                    ) from e
+                sol_i, obj_i = self._solve(c)
+                if isinstance(sol_i, torch.Tensor):
+                    sol_i = sol_i.detach().cpu().numpy()
                 sol_knn[:, i] = sol_i
                 obj_knn[i] = obj_i
             # get average
@@ -264,6 +268,11 @@ class optDatasetKNN(optDataset):
         """
         A method to get kNN costs
         """
+        # scipy needs host numpy arrays
+        if isinstance(self.feats, torch.Tensor):
+            self.feats = self.feats.detach().cpu().numpy()
+        if isinstance(self.costs, torch.Tensor):
+            self.costs = self.costs.detach().cpu().numpy()
         # calculate distances between features
         distances = distance.cdist(self.feats, self.feats, "euclidean")
         # exclude self (diagonal) to get true nearest neighbours
@@ -281,9 +290,9 @@ class optDatasetConstrs(optDataset):
     """
     PyTorch ``Dataset`` for the CaVE cone-aligned loss.
 
-    Stores features and cost coefficients, solves each instance with a fresh
-    copy of the Gurobi model, and extracts the **normals of the binding
-    constraints at the optimal vertex** in canonical ``<=`` orientation.
+    Stores features and cost coefficients, solves each instance, and extracts
+    the **normals of the binding constraints at the optimal vertex** in
+    canonical ``<=`` orientation.
     These normals span the polyhedral cone onto which ``coneAlignedCosine``
     projects the predicted cost vector during training.
 
@@ -327,6 +336,11 @@ class optDatasetConstrs(optDataset):
         """
         if not isinstance(model, optModel):
             raise TypeError("arg model is not an optModel")
+        if len(feats) != len(costs):
+            raise ValueError(
+                f"feats and costs must have the same number of instances: "
+                f"{len(feats)} vs {len(costs)}."
+            )
         self.model = model
         self.skip_infeas = skip_infeas
         # data
@@ -355,11 +369,16 @@ class optDatasetConstrs(optDataset):
         ctrs: list[np.ndarray] = []
         valid: list[int] = []
         logger.info("Optimizing for optDatasetConstrs...")
+        model = self.model
         for i, c in enumerate(tqdm(self.costs)):
-            # fresh per-instance copy keeps the lazy-constraint buffer clean
-            model = self.model.copy()
-            model.setObj(model._fullCost(c))
-            sol, obj = model.solve()
+            model._setFullObj(model._fullCost(c))
+            try:
+                sol, obj = model.solve()
+            except RuntimeError as e:
+                if self.skip_infeas:
+                    logger.warning("Instance %d had no solution, skipping: %s", i, e)
+                    continue
+                raise
             # infeasibility check
             if model._model.Status != GRB.OPTIMAL:
                 if self.skip_infeas:
@@ -375,7 +394,13 @@ class optDatasetConstrs(optDataset):
                 )
             # binary-vertex check: CaVE is defined for binary linear programs
             sol_arr = np.asarray(sol, dtype=np.float64)
-            if not np.all((sol_arr < 1e-5) | (sol_arr > 1 - 1e-5)):
+            is_binary = np.all(
+                np.isclose(sol_arr, 0.0, atol=1e-5) | np.isclose(sol_arr, 1.0, atol=1e-5)
+            )
+            if not is_binary:
+                if self.skip_infeas:
+                    logger.warning("Instance %d optimal vertex is not binary, skipping.", i)
+                    continue
                 raise ValueError(
                     f"Instance {i} optimal vertex is not binary; "
                     "CaVE requires binary linear programs."
@@ -384,6 +409,8 @@ class optDatasetConstrs(optDataset):
             objs.append([float(obj)])
             ctrs.append(_extract_tight_normals(model, sol_arr))
             valid.append(i)
+        if not valid:
+            raise ValueError("No valid instances (all skipped or empty input).")
         return np.stack(sols), np.asarray(objs), ctrs, valid
 
     def __len__(self) -> int:
@@ -455,11 +482,14 @@ def _extract_tight_normals(
         senses_arr = np.asarray(grb.getAttr("Sense", constrs))
         tight_mask = np.abs(slacks) < tol
         if tight_mask.any():
-            # project the constraint matrix onto cost-variable columns
-            cost_col_idx = np.asarray([v.index for v in cost_vars])
-            A = grb.getA().tocsr()
+            # cost-column slice cached across solves; the row count keys invalidation
+            cache = getattr(grb, "_cave_A_cost", None)
+            if cache is None or cache[0] != grb.NumConstrs:
+                cost_col_idx = np.asarray([v.index for v in cost_vars])
+                cache = (grb.NumConstrs, grb.getA().tocsr()[:, cost_col_idx])
+                grb._cave_A_cost = cache
             # extract all tight rows in a single sparse-to-dense conversion
-            A_tight = A[:, cost_col_idx][tight_mask].toarray()
+            A_tight = cache[1][tight_mask].toarray()
             tight_senses = senses_arr[tight_mask]
             is_le = tight_senses == GRB.LESS_EQUAL
             is_ge = tight_senses == GRB.GREATER_EQUAL

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/data/knapsack.py

@@ -44,6 +44,8 @@ def genData(
         raise ValueError(f"deg = {deg} should be int.")
     if deg <= 0:
         raise ValueError(f"deg = {deg} should be positive.")
+    if noise_width < 0:
+        raise ValueError(f"noise_width = {noise_width} should be non-negative.")
     # set seed
     rnd = np.random.RandomState(seed)
     # number of data points

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/data/shortestpath.py

@@ -41,6 +41,8 @@ def genData(
         raise ValueError(f"deg = {deg} should be int.")
     if deg <= 0:
         raise ValueError(f"deg = {deg} should be positive.")
+    if noise_width < 0:
+        raise ValueError(f"noise_width = {noise_width} should be non-negative.")
     # set seed
     rnd = np.random.RandomState(seed)
     # number of data points

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/data/tsp.py

@@ -43,6 +43,8 @@ def genData(
         raise ValueError(f"deg = {deg} should be int.")
     if deg <= 0:
         raise ValueError(f"deg = {deg} should be positive.")
+    if noise_width < 0:
+        raise ValueError(f"noise_width = {noise_width} should be non-negative.")
     # set seed
     rnd = np.random.RandomState(seed)
     # number of data points

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/dsl/compiled.py

@@ -42,13 +42,22 @@ class compiledBase(optModel):
 
     def setObj(self, c):
         """Set the objective from a predicted cost of length ``num_cost``, scattered onto the known fixed costs."""
-        # scatter onto fixed costs
         prob = self.problem
         coef = costToNumpy(c)
-        if coef.shape[-1] != prob.num_vars:
+        # scatter onto fixed costs; an unambiguous full-length vector passes through
+        if coef.shape[-1] == prob.num_cost:
             full = prob.fixed_cost.copy()
             full[prob.c_pred_index] += coef
             coef = full
+        elif coef.shape[-1] != prob.num_vars:
+            raise ValueError("Size of cost vector does not match number of cost variables.")
+        self._write_obj(coef)
+
+    def _setFullObj(self, c):
+        """Set the objective from full-space coefficients (length ``num_vars``), bypassing the predicted-cost scatter."""
+        coef = costToNumpy(c)
+        if coef.shape[-1] != self.problem.num_vars:
+            raise ValueError("Size of cost vector does not match number of variables.")
         self._write_obj(coef)
 
     def _fullCost(self, pred_cost):
@@ -57,8 +66,12 @@ class compiledBase(optModel):
         idx = prob.c_pred_index
         if isinstance(pred_cost, torch.Tensor):
             index = torch.as_tensor(idx, dtype=torch.long, device=pred_cost.device)
-            scattered = pred_cost.new_zeros((*pred_cost.shape[:-1], prob.num_vars)).index_add(-1, index, pred_cost)
-            return scattered + torch.as_tensor(prob.fixed_cost, dtype=pred_cost.dtype, device=pred_cost.device)
+            scattered = pred_cost.new_zeros((*pred_cost.shape[:-1], prob.num_vars)).index_add(
+                -1, index, pred_cost
+            )
+            return scattered + torch.as_tensor(
+                prob.fixed_cost, dtype=pred_cost.dtype, device=pred_cost.device
+            )
         arr = np.asarray(pred_cost, dtype=float)
         full = np.broadcast_to(prob.fixed_cost, (*arr.shape[:-1], prob.num_vars)).copy()
         full[..., idx] += arr
@@ -67,17 +80,26 @@ class compiledBase(optModel):
     def solve(self):
         """Solve and return the full decision-vector solution (length ``num_vars``) with its objective value."""
         sol, obj = self._read_sol()
-        return np.asarray(sol), obj
+        # bare objective constants live outside the solver model
+        return np.asarray(sol), obj + self.problem.obj_offset
 
     def addConstr(self, coefs, rhs):
         # add a cut coefs @ x <= rhs over the full variable vector
-        return self._add_cut(np.asarray(coefs, dtype=float).reshape(-1), rhs)
+        coefs = np.asarray(coefs, dtype=float).reshape(-1)
+        new_model = self._add_cut(coefs, rhs)
+        # track for replay on relax
+        new_model._extra_constrs = [*self._extra_constrs, (coefs, float(rhs))]
+        return new_model
 
     def relax(self):
         # recompile the relaxed problem, preserving backend kwargs
         kwargs = getArgs(self)
         kwargs["problem"] = self.problem.relax()
-        return type(self)(**kwargs)
+        model_rel = type(self)(**kwargs)
+        # replay user cuts on the relaxation
+        for coefs, rhs in self._extra_constrs:
+            model_rel = model_rel.addConstr(coefs, rhs)
+        return model_rel
 
     def _apply_params(self):
         # push self.params to the solver

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/dsl/expression.py

@@ -39,7 +39,7 @@ class Variable:
     """
 
     def __init__(self, shape, *, vtype=EPO.CONTINUOUS, lb=None, ub=None, name=None):
-        self.shape = (int(shape),) if np.isscalar(shape) else tuple(int(s) for s in shape)
+        self.shape = tuple(int(s) for s in np.atleast_1d(shape))
         self.size = int(np.prod(self.shape)) if self.shape else 1
         # per-entry type: a scalar EPO.VarType broadcast, or a per-entry array
         self.vtype = np.broadcast_to(np.asarray(vtype, dtype=object), self.shape).reshape(-1).copy()
@@ -86,6 +86,9 @@ class Variable:
     def __rmul__(self, o):
         return self._to_affine() * o
 
+    def __truediv__(self, o):
+        return self._to_affine() / o
+
     def __matmul__(self, o):
         return self._to_affine() @ o
 
@@ -150,7 +153,11 @@ class Affine:
                 self._add_block(blocks, v, b)
             return Affine(blocks, self.const + o.const, self.shape)
         if _is_num(o):
-            return Affine(self.blocks, self.const + np.asarray(o, dtype=float).reshape(-1), self.shape)
+            # broadcast the constant over the logical shape; a flat full-length vector is kept as-is
+            a = np.asarray(o, dtype=float)
+            if not (a.ndim == 1 and a.size == self.size):
+                a = np.broadcast_to(a, self.shape or (self.size,))
+            return Affine(self.blocks, self.const + a.reshape(-1), self.shape)
         return NotImplemented
 
     def __radd__(self, o):
@@ -186,6 +193,12 @@ class Affine:
         # reflected scale
         return self * o
 
+    def __truediv__(self, o):
+        # elementwise divide by a scalar or shape-broadcast array
+        if _is_num(o):
+            return self * (1.0 / np.asarray(o, dtype=float))
+        return NotImplemented
+
     def __matmul__(self, o):
         # Affine @ (ndarray)  -> Affine ;  Affine @ (Variable | Affine) -> Quadratic
         if isinstance(o, Variable):
@@ -216,22 +229,29 @@ class Affine:
         # reduce along axis (or fully) via a 0/1 summation matrix
         if axis is None:
             return self._row_op(np.ones((1, self.size)), ())
+        # normalize a negative axis (numpy semantics)
+        if not -len(self.shape) <= axis < len(self.shape):
+            raise ValueError(f"axis {axis} is out of bounds for shape {self.shape}.")
+        axis %= len(self.shape)
         idx = np.arange(self.size).reshape(self.shape)
         kept = tuple(d for a, d in enumerate(self.shape) if a != axis)
         out_m = int(np.prod(kept)) if kept else 1
         # build summation matrix: out row r sums the flat cols collapsed into it
         keep_idx = np.moveaxis(idx, axis, -1).reshape(out_m, self.shape[axis])
-        S = sp.lil_matrix((out_m, self.size))
-        for r in range(out_m):
-            S[r, keep_idx[r]] = 1.0
-        return self._row_op(S.tocsr(), kept)
+        rows = np.repeat(np.arange(out_m), self.shape[axis])
+        S = sp.csr_matrix(
+            (np.ones(keep_idx.size), (rows, keep_idx.reshape(-1))), shape=(out_m, self.size)
+        )
+        return self._row_op(S, kept)
 
     def __getitem__(self, idx):
         # numpy-style indexing -> row selection
         sel = np.arange(self.size).reshape(self.shape)[idx]
         out_shape = sel.shape if hasattr(sel, "shape") else ()
         sel = np.atleast_1d(sel).reshape(-1)
-        S = sp.csr_matrix((np.ones(sel.size), (np.arange(sel.size), sel)), shape=(sel.size, self.size))
+        S = sp.csr_matrix(
+            (np.ones(sel.size), (np.arange(sel.size), sel)), shape=(sel.size, self.size)
+        )
         return self._row_op(S, out_shape)
 
     # ---- constraints ----
@@ -290,8 +310,10 @@ class Quadratic:
         # Quadratic + (Affine | Variable | Quadratic | const) -> Quadratic
         if isinstance(o, Variable):
             o = o._to_affine()
-        # affine / constant folds into the linear part
+        # affine / constant folds into the scalar linear part
         if isinstance(o, Affine):
+            if o.size != 1:
+                raise TypeError("A Quadratic is scalar; reduce the added term with .sum().")
             return Quadratic(self.quad, self.affine + o)
         # merge the two quadratic block dicts
         if isinstance(o, Quadratic):
@@ -307,6 +329,14 @@ class Quadratic:
         # reflected add
         return self + o
 
+    def __sub__(self, o):
+        # subtract via negated add
+        return self + (-o)
+
+    def __rsub__(self, o):
+        # reflected subtract
+        return (-self) + o
+
     def __mul__(self, o):
         # scale the quadratic and its affine part by a scalar
         if _is_num(o):
@@ -318,6 +348,16 @@ class Quadratic:
         # reflected scale
         return self * o
 
+    def __truediv__(self, o):
+        # scale by the reciprocal scalar
+        if _is_num(o):
+            return self * (1.0 / float(o))
+        return NotImplemented
+
+    def __neg__(self):
+        # negate the quadratic and its affine part
+        return self * -1.0
+
     # ---- quadratic constraints ----
     def __le__(self, rhs):
         return Constraint(self, "<=", rhs)
@@ -339,9 +379,12 @@ class Quadratic:
 
 def _is_row_selection(S, rows):
     # one entry of value 1 per row, distinct columns (a plain index / slice pick)
-    return (S.nnz == rows and (S.data == 1).all()
-            and np.array_equal(S.indptr, np.arange(rows + 1))
-            and len(set(S.indices)) == S.nnz)
+    return (
+        S.nnz == rows
+        and (S.data == 1).all()
+        and np.array_equal(S.indptr, np.arange(rows + 1))
+        and len(set(S.indices)) == S.nnz
+    )
 
 
 def _as_selection(operand):
@@ -349,12 +392,14 @@ def _as_selection(operand):
     if isinstance(operand, Variable):
         return operand, None
     if isinstance(operand, Affine) and len(operand.blocks) == 1 and not operand.const.any():
-        (var, S), = operand.blocks.items()
+        ((var, S),) = operand.blocks.items()
         S = S.tocsr()
         if _is_row_selection(S, operand.size):
-            return var, S.indices.copy()                # selected columns of var
-    raise TypeError("A Parameter may only multiply a Variable or a plain slice / index "
-                    "of one (e.g. x, x[:k], x[idx]).")
+            return var, S.indices.copy()  # selected columns of var
+    raise TypeError(
+        "A Parameter may only multiply a Variable or a plain slice / index "
+        "of one (e.g. x, x[:k], x[idx])."
+    )
 
 
 def _make_objective(param, base, sel, inner=False):
@@ -362,10 +407,14 @@ def _make_objective(param, base, sel, inner=False):
     var, local_idx = _as_selection(sel)
     # `@` is a 1-D inner product; a multi-dimensional cost must go through (c * x).sum()
     if inner and len(getattr(sel, "shape", ())) > 1:
-        raise TypeError("c @ x is a 1-D inner product; for a multi-dimensional cost write (c * x).sum().")
+        raise TypeError(
+            "c @ x is a 1-D inner product; for a multi-dimensional cost write (c * x).sum()."
+        )
     n = var.size if local_idx is None else len(local_idx)
     if param.size != n:
-        raise TypeError(f"Parameter size {param.size} does not match the selected variable size {n}.")
+        raise TypeError(
+            f"Parameter size {param.size} does not match the selected variable size {n}."
+        )
     return ParametricObjective(param, var, local_idx, base=base)
 
 
@@ -385,7 +434,7 @@ class Parameter:
     """
 
     def __init__(self, shape, *, name=None):
-        self.shape = (int(shape),) if np.isscalar(shape) else tuple(int(s) for s in shape)
+        self.shape = tuple(int(s) for s in np.atleast_1d(shape))
         self.size = int(np.prod(self.shape)) if self.shape else 1
         self.name = name
 
@@ -418,8 +467,10 @@ class Parameter:
         return self._forbid()
 
     def _forbid(self, *a, **k):
-        raise TypeError("Unsupported operation on Parameter (only c @ var / c * var, "
-                        "optionally with a known base: (d + c) @ var).")
+        raise TypeError(
+            "Unsupported operation on Parameter (only c @ var / c * var, "
+            "optionally with a known base: (d + c) @ var)."
+        )
 
     __rsub__ = __neg__ = __getitem__ = __le__ = __ge__ = __eq__ = _forbid
 
@@ -437,7 +488,10 @@ class ParametricCoef:
 
     def __init__(self, param, base):
         self.param = param
-        self.base = np.broadcast_to(np.asarray(base, dtype=float), (param.size,)).copy()
+        base = np.asarray(base, dtype=float)
+        # multi-dim bases broadcast over the parameter's logical shape, flat ones over its size
+        shape = param.shape if base.ndim > 1 else (param.size,)
+        self.base = np.broadcast_to(base, shape).reshape(-1).astype(float)
 
     __array_ufunc__ = None
 
@@ -470,7 +524,9 @@ class ParametricVector:
         var, local_idx = _as_selection(sel)
         n = var.size if local_idx is None else len(local_idx)
         if param.size != n:
-            raise TypeError(f"Parameter size {param.size} does not match the selected variable size {n}.")
+            raise TypeError(
+                f"Parameter size {param.size} does not match the selected variable size {n}."
+            )
         self.param = param
         self.base = base
         self.sel = sel
@@ -481,12 +537,16 @@ class ParametricVector:
     def sum(self, axis=None):
         # reduce to the scalar predicted objective term
         if axis is not None:
-            raise TypeError("A predicted objective must reduce to a scalar; use .sum() with no axis.")
+            raise TypeError(
+                "A predicted objective must reduce to a scalar; use .sum() with no axis."
+            )
         return _make_objective(self.param, self.base, self.sel)
 
     def __add__(self, o):
-        raise TypeError("c * x is an elementwise vector; reduce it with .sum() before adding "
-                        "other terms, e.g. (c * x).sum() + d @ y.")
+        raise TypeError(
+            "c * x is an elementwise vector; reduce it with .sum() before adding "
+            "other terms, e.g. (c * x).sum() + d @ y."
+        )
 
     __radd__ = __add__
 
@@ -519,9 +579,14 @@ class ParametricObjective:
 
     def _with(self, fixed=None, quad=None):
         # copy with a replaced fixed / quad term
-        return ParametricObjective(self.cost_param, self.cost_var, self.local_idx, self.base,
-                                  self.fixed if fixed is None else fixed,
-                                  self.quad if quad is None else quad)
+        return ParametricObjective(
+            self.cost_param,
+            self.cost_var,
+            self.local_idx,
+            self.base,
+            self.fixed if fixed is None else fixed,
+            self.quad if quad is None else quad,
+        )
 
     def sum(self, axis=None):
         # already a scalar; sum is the identity here
@@ -530,12 +595,17 @@ class ParametricObjective:
     def __add__(self, o):
         # attach a known linear (Affine / Variable) or quadratic (Quadratic) term
         if isinstance(o, ParametricObjective):
-            raise TypeError("Cannot add two predicted cost terms; a problem has exactly one predicted cost.")
+            raise TypeError(
+                "Cannot add two predicted cost terms; a problem has exactly one predicted cost."
+            )
         if isinstance(o, Quadratic):
             return self._with(quad=o if self.quad is None else self.quad + o)
         if isinstance(o, Variable):
             o = o._to_affine()
         if isinstance(o, Affine):
+            # a vector objective term is meaningless; require a scalar
+            if o.size != 1:
+                raise TypeError("A known objective term must be scalar; reduce it with .sum().")
             return self._with(fixed=o if self.fixed is None else self.fixed + o)
         if _is_num(o):
             raise TypeError("A constant objective term has no effect; omit it.")
@@ -552,7 +622,7 @@ class Constraint:
     Attributes:
         lhs (Affine | Quadratic): left-hand side expression
         sense (str): one of ``"<="`` / ``">="`` / ``"=="``
-        rhs (np.ndarray): right-hand side, broadcast to the LHS length
+        rhs (np.ndarray): right-hand side, broadcast over the LHS shape
     """
 
     def __init__(self, lhs, sense, rhs):
@@ -599,5 +669,9 @@ class Constraint:
         return None, A, self.sense, self._rhs_vec(aff.size) - aff.const
 
     def _rhs_vec(self, m):
-        # broadcast the rhs to a flat (m,) vector
-        return np.broadcast_to(np.asarray(self.rhs, dtype=float), (m,)).astype(float)
+        # broadcast the rhs over the LHS logical shape; a flat full-length vector is kept as-is
+        rhs = np.asarray(self.rhs, dtype=float)
+        if rhs.ndim == 1 and rhs.size == m:
+            return rhs.astype(float)
+        shape = getattr(self.lhs, "shape", ()) or (m,)
+        return np.broadcast_to(rhs, shape).reshape(-1).astype(float)

{pyepo-2.2.0 → pyepo-2.2.1}/pyepo/dsl/objective.py

@@ -25,14 +25,19 @@ class Objective:
 
     def __init__(self, expr):
         from pyepo.dsl.expression import ParametricObjective, ParametricVector
+
         # an elementwise c * x is a vector, not a scalar objective
         if isinstance(expr, ParametricVector):
-            raise TypeError("c * x is an elementwise vector, not a scalar objective; "
-                            "write (c * x).sum() or c @ x.")
+            raise TypeError(
+                "c * x is an elementwise vector, not a scalar objective; "
+                "write (c * x).sum() or c @ x."
+            )
         # the objective must carry the predicted cost (a ParametricObjective)
         if not isinstance(expr, ParametricObjective):
-            raise TypeError("The objective must be a predicted cost term like c @ x, "
-                            "optionally plus a known quadratic term.")
+            raise TypeError(
+                "The objective must be a predicted cost term like c @ x, "
+                "optionally plus a known quadratic term."
+            )
         self.expr = expr
 
     @property
@@ -46,9 +51,11 @@ class Objective:
 
 class Minimize(Objective):
     """Minimization objective."""
+
     modelSense = EPO.MINIMIZE
 
 
 class Maximize(Objective):
     """Maximization objective."""
+
     modelSense = EPO.MAXIMIZE