pyepo 2.2.4__tar.gz → 2.2.6__tar.gz

{pyepo-2.2.4 → pyepo-2.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.4
+Version: 2.2.6
 Summary: PyTorch/JAX-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT
@@ -16,6 +16,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Operating System :: OS Independent
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -29,7 +30,7 @@ Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: pathos
 Requires-Dist: tqdm
-Requires-Dist: scikit_learn
+Requires-Dist: scikit-learn
 Requires-Dist: torch>=1.13.1
 Provides-Extra: pyomo
 Requires-Dist: pyomo>=6.1.2; extra == "pyomo"

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/data/_validation.py

@@ -6,8 +6,7 @@ from numbers import Real
 
 def validate_degree(deg: int) -> None:
     """Validate a positive integer polynomial degree."""
-    if not isinstance(deg, int) or isinstance(deg, bool) or deg <= 0:
-        raise ValueError(f"deg = {deg} should be a positive integer.")
+    validate_positive_int(deg, "deg")
 
 
 def validate_nonnegative(value: float, name: str) -> None:

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/data/dataset.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, cast
 import numpy as np
 import torch
 from scipy.spatial import distance
-from torch.utils.data import Dataset
+from torch.utils.data import DataLoader, Dataset
 from tqdm import tqdm
 
 from pyepo import EPO
@@ -33,14 +33,20 @@ def _validate_inputs(
     model: optModel,
     feats: np.ndarray | torch.Tensor,
     costs: np.ndarray | torch.Tensor,
-) -> None:
-    """Validate the common constructor contract for optimization datasets."""
+) -> tuple[np.ndarray | torch.Tensor, np.ndarray | torch.Tensor]:
+    """Validate and normalize the common constructor contract for optimization datasets."""
     if not isinstance(model, optModel):
         raise TypeError("arg model is not an optModel")
+    # array-likes become numpy arrays
+    if not isinstance(feats, (np.ndarray, torch.Tensor)):
+        feats = np.asarray(feats)
+    if not isinstance(costs, (np.ndarray, torch.Tensor)):
+        costs = np.asarray(costs)
     if len(feats) != len(costs):
         raise ValueError(
             f"feats and costs must have the same number of instances: {len(feats)} vs {len(costs)}."
         )
+    return feats, costs
 
 
 def _as_float_tensor(data) -> torch.Tensor:
@@ -91,7 +97,7 @@ class optDataset(Dataset):
             feats: data features
             costs: costs of objective function
         """
-        _validate_inputs(model, feats, costs)
+        feats, costs = _validate_inputs(model, feats, costs)
         self.model = model
         # data
         self.feats = feats
@@ -229,7 +235,7 @@ class optDatasetKNN(optDataset):
             k: number of nearest neighbours selected
             weight: self-weight in the kNN convex combination (1.0 = no smoothing)
         """
-        _validate_inputs(model, feats, costs)
+        feats, costs = _validate_inputs(model, feats, costs)
         self.model = model
         # at most num_data-1 neighbours exist (self excluded), so k must stay below it
         num_data = len(feats)
@@ -315,7 +321,8 @@ class optDatasetConstrs(optDataset):
     currently requires a Gurobi-backed ``optModel``.
 
     Per-instance row counts differ (different constraints bind at different
-    vertices), so batches must be assembled with ``collate_tight_constraints``.
+    vertices), so batch with ``optDataLoader`` or pass
+    ``collate_tight_constraints`` to a PyTorch ``DataLoader``.
 
     Reference: Tang & Khalil (2024)
     `<https://link.springer.com/chapter/10.1007/978-3-031-60599-4_12>`_
@@ -346,7 +353,7 @@ class optDatasetConstrs(optDataset):
             costs: costs of objective function
             skip_infeas: if True, drop infeasible instances instead of raising
         """
-        _validate_inputs(model, feats, costs)
+        feats, costs = _validate_inputs(model, feats, costs)
         self.model = model
         self.skip_infeas = skip_infeas
         # data
@@ -455,6 +462,29 @@ def collate_tight_constraints(batch):
     )
 
 
+# optDatasetConstrs yields ragged binding-constraint matrices; pad them when batching
+optDatasetConstrs.collate_fn = staticmethod(collate_tight_constraints)
+
+
+class optDataLoader(DataLoader):
+    """
+    ``DataLoader`` that applies a dataset's own ``collate_fn`` when present.
+
+    Datasets with ragged samples (e.g. ``optDatasetConstrs``, whose
+    binding-constraint matrices vary in row count) carry a ``collate_fn``; this
+    loader uses it so the caller never passes one explicitly. Plain datasets
+    fall back to the default PyTorch collation.
+    """
+
+    def __init__(self, dataset, *args, **kwargs):
+        # use the dataset's own collate_fn unless the caller supplies one
+        if len(args) <= 5 and "collate_fn" not in kwargs:
+            collate_fn = getattr(dataset, "collate_fn", None)
+            if collate_fn is not None:
+                kwargs["collate_fn"] = collate_fn
+        super().__init__(dataset, *args, **kwargs)
+
+
 def _extract_tight_normals(
     model: optModel,
     sol: np.ndarray,

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/cave.py

@@ -59,12 +59,16 @@ class coneAlignedCosine(optModule):
     cutting the per-epoch cost without measurable regret loss.
 
     Training data must come from ``pyepo.data.dataset.optDatasetConstrs``
-    (Gurobi-backed) and be collated with ``collate_tight_constraints``.
+    (Gurobi-backed) and be batched with ``pyepo.data.dataset.optDataLoader``
+    or a ``DataLoader`` using ``collate_tight_constraints``.
 
     Reference: Tang & Khalil (2024)
     `<https://link.springer.com/chapter/10.1007/978-3-031-60599-4_12>`_
     """
 
+    # Clarabel-only workers
+    _pool_needs_optmodel = False
+
     def __init__(
         self,
         optmodel: optModel,

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/jax/cave.py

@@ -34,6 +34,9 @@ class coneAlignedCosine(optModule):
     `<https://link.springer.com/chapter/10.1007/978-3-031-60599-4_12>`_
     """
 
+    # Clarabel-only workers
+    _pool_needs_optmodel = False
+
     def __init__(
         self,
         optmodel,

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/jax/regularized.py

@@ -139,9 +139,8 @@ def _away_step_frank_wolfe(theta, module, use_cache=False):
         w = w.at[bidx, match_idx].add(gamma_fw * (has_match & use_fw).astype(theta.dtype))
         vt = vt.at[bidx, free_idx].set(jnp.where(add_new[:, None], v, vt[bidx, free_idx]))
         vn = vn.at[bidx, free_idx].set(jnp.where(add_new, vnv, vn[bidx, free_idx]))
-        w = w.at[bidx, free_idx].set(
-            jnp.where(add_new, gamma_fw + w[bidx, free_idx], w[bidx, free_idx])
-        )
+        # displaced mass is dropped
+        w = w.at[bidx, free_idx].set(jnp.where(add_new, gamma_fw, w[bidx, free_idx]))
         # away subtract, then clear FP residue so dropped atoms leave the active set
         w = w.at[bidx, away_idx].add(-gamma_away)
         w = jnp.where(w < 1e-12, 0.0, jnp.maximum(w, 0.0))

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/jax/surrogate.py

@@ -14,7 +14,7 @@ import jax.numpy as jnp
 from pyepo.func._common import is_minimize, validate_positive
 from pyepo.func.jax.abcmodule import optModule
 from pyepo.func.jax.utils import _full_cost, _solve_or_cache
-from pyepo.utils import _EPS
+from pyepo.utils import _EPS, objective_offset
 
 if TYPE_CHECKING:
     from pyepo.func.runtime import Reduction
@@ -68,6 +68,11 @@ def _spoplus_value_and_grad(pred_cost, true_cost, true_sol, true_obj, module):
     # solve the perturbed problem
     sol, obj = _solve_or_cache(2.0 * pred_cost - true_cost, module)
     z = jnp.squeeze(true_obj, axis=-1) if true_obj.ndim > 1 else true_obj
+    # drop the bare objective constant
+    offset = objective_offset(module.optmodel)
+    if offset:
+        obj = obj - offset
+        z = z - offset
     inner = 2.0 * jnp.einsum("bi,bi->b", pred_cost, true_sol)
     # loss and subgradient
     if is_minimize(module.optmodel.modelSense):

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/jax/utils.py

@@ -12,6 +12,7 @@ import numpy as np
 from pyepo.func._common import is_minimize, solution_pool_tolerance
 from pyepo.func.utils import _solve_batch_np
 from pyepo.model.mpax import optMpaxModel
+from pyepo.utils import objective_offset
 
 try:
     from jax.extend.core import concrete_or_error as _concrete_or_error
@@ -83,6 +84,10 @@ def _solve_batch_mpax(cost, optmodel):
     # obj in true sense
     if not minimize:
         obj = -obj
+    # add the bare objective constant
+    offset = objective_offset(optmodel)
+    if offset:
+        obj = obj + offset
     return sol, obj
 
 
@@ -127,6 +132,10 @@ def _cache_in_pass(cost, optmodel, solpool):
     select = jnp.argmin if is_minimize(optmodel.modelSense) else jnp.argmax
     ind = select(solpool_obj, axis=1)
     obj = jnp.take_along_axis(solpool_obj, ind[:, None], axis=1).squeeze(1)
+    # add the bare objective constant
+    offset = objective_offset(optmodel)
+    if offset:
+        obj = obj + offset
     return solpool[ind], obj
 
 

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/perturbed.py

@@ -13,20 +13,16 @@ from torch.autograd import Function
 
 from pyepo.func._common import (
     is_minimize,
-    require_solution_pool,
     validate_positive,
     validate_positive_int,
 )
 from pyepo.func.abcmodule import optModule
 from pyepo.func.utils import (
     _mask_pred,
+    _solve_or_cache,
     _torch_generator,
-    _update_solution_pool,
     sumGammaDistribution,
 )
-from pyepo.func.utils import (
-    _solve_batch as _solve_batch_2d,
-)
 from pyepo.utils import _EPS
 
 if TYPE_CHECKING:
@@ -625,76 +621,12 @@ class adaptiveImplicitMLEFunc(implicitMLEFunc):
 
 def _solve_or_cache_3d(ptb_c: torch.Tensor, module: optModule) -> torch.Tensor:
     """
-    Solve or use cached solutions for perturbed costs (3D: n_samples × batch × vars).
-    Delegates to the shared 2D functions in utils after flattening.
-    """
-    optmodel = module.optmodel
-    processes = module.processes
-    pool = module.pool
-    solpool = module.solpool
-    if module._branch_rng.uniform() <= module.solve_ratio:
-        ptb_sols, solpool = _solve_in_pass_3d(ptb_c, optmodel, processes, pool, solpool)
-    else:
-        solpool = require_solution_pool(solpool)
-        ptb_sols, solpool = _cache_in_pass_3d(ptb_c, optmodel, solpool)
-    module.solpool = solpool
-    return ptb_sols
-
-
-def _solve_in_pass_3d(
-    ptb_c: torch.Tensor,
-    optmodel: optModel,
-    processes: int,
-    pool,
-    solpool: torch.Tensor | None = None,
-) -> tuple[torch.Tensor, torch.Tensor | None]:
-    """
-    Solve optimization for perturbed 3D costs and update solution pool.
-
-    Args:
-        ptb_c: perturbed costs, shape (batch, n_samples, vars)
-        optmodel: optimization model
-        processes: number of processors
-        pool: process pool
-        solpool: solution pool
-
-    Returns:
-        tuple: (solutions shape (batch, n_samples, vars), updated solpool)
-    """
-    ins_num, n_samples, num_vars = ptb_c.shape
-    # flatten (batch, n_samples, vars) → (batch * n_samples, vars)
-    flat_c = ptb_c.reshape(-1, num_vars)
-    # solve using shared 2D function
-    flat_sols, _ = _solve_batch_2d(flat_c, optmodel, processes, pool)
-    # update pool while flat_sols is still contiguous
-    if solpool is not None:
-        solpool = _update_solution_pool(flat_sols, solpool)
-        if solpool.device != ptb_c.device:
-            solpool = solpool.to(ptb_c.device)
-    # reshape (batch * n_samples, vars) → (batch, n_samples, vars)
-    ptb_sols = flat_sols.reshape(ins_num, n_samples, num_vars)
-    return ptb_sols, solpool
-
-
-def _cache_in_pass_3d(
-    ptb_c: torch.Tensor,
-    optmodel: optModel,
-    solpool: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Use solution pool for perturbed 3D costs (batch × n_samples × vars).
-    Unlike the 2D version in utils, this handles the extra sample dimension.
+    Solve or use cached solutions for perturbed costs (batch × n_samples × vars).
+    Flattens the sample axis and delegates to the shared 2D path in utils.
     """
-    # move solpool to the correct device
-    if solpool.device != ptb_c.device:
-        solpool = solpool.to(ptb_c.device)
-    # compute objective values: (batch, n_samples, pool_size)
-    solpool_obj = torch.einsum("bnd,sd->bns", ptb_c, solpool)
-    # best solution in pool
-    select = torch.argmin if is_minimize(optmodel.modelSense) else torch.argmax
-    best_inds = select(solpool_obj, dim=2)
-    ptb_sols = solpool[best_inds]
-    return ptb_sols, solpool
+    batch, n_samples, num_vars = ptb_c.shape
+    sol, _ = _solve_or_cache(ptb_c.reshape(-1, num_vars), module)
+    return sol.reshape(batch, n_samples, num_vars)
 
 
 # acronym aliases

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/regularized.py

@@ -120,9 +120,8 @@ def _away_step_frank_wolfe(
         vertex_norms[batch_idx, free_idx] = torch.where(
             add_new, v_norm_sq, vertex_norms[batch_idx, free_idx]
         )
-        weights[batch_idx, free_idx] = torch.where(
-            add_new, gamma_fw + weights[batch_idx, free_idx], weights[batch_idx, free_idx]
-        )
+        # displaced mass is dropped
+        weights[batch_idx, free_idx] = torch.where(add_new, gamma_fw, weights[batch_idx, free_idx])
         # away subtract, then clear FP residue so dropped atoms leave the active set
         weights[batch_idx, away_idx] = weights[batch_idx, away_idx] - gamma_away
         weights = weights.clamp(min=0.0)

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/runtime.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import itertools
 import multiprocessing as mp
 import weakref
 from dataclasses import dataclass
@@ -63,20 +64,27 @@ def normalize_processes(
     return cpu_count if processes == 0 else processes
 
 
+# unique pool ids
+_pool_ids = itertools.count()
+
+
 def create_solver_pool(
     optmodel: optModel,
     processes: int,
     *,
     owner=None,
+    with_solver: bool = True,
 ) -> ProcessingPool | None:
     """Create a worker pool, optionally tied to an owner's lifetime."""
     if processes == 1:
         return None
-    pool = ProcessingPool(
-        processes,
-        initializer=_init_worker_model,
-        initargs=(optmodel.to_spec(),),
+    # optional per-worker optmodel preload
+    init_kwargs = (
+        {"initializer": _init_worker_model, "initargs": (optmodel.to_spec(),)}
+        if with_solver
+        else {}
     )
+    pool = ProcessingPool(processes, id=f"pyepo-{next(_pool_ids)}", **init_kwargs)
     if owner is not None:
         weakref.finalize(owner, _close_pool, pool)
     return pool
@@ -117,7 +125,12 @@ def init_runtime(
         raise ValueError(f"No reduction '{reduction}'.")
 
     normalized_processes = normalize_processes(optmodel, processes, logger)
-    pool = create_solver_pool(optmodel, normalized_processes, owner=owner)
+    pool = create_solver_pool(
+        optmodel,
+        normalized_processes,
+        owner=owner,
+        with_solver=getattr(owner, "_pool_needs_optmodel", True),
+    )
     logger.info("Num of cores: %d", normalized_processes)
     return RuntimeState(
         optmodel=optmodel,

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/surrogate.py

@@ -13,7 +13,7 @@ from torch.autograd import Function
 from pyepo.func._common import is_minimize, validate_positive
 from pyepo.func.abcmodule import optModule
 from pyepo.func.utils import _solve_or_cache
-from pyepo.utils import _EPS
+from pyepo.utils import _EPS, objective_offset
 
 if TYPE_CHECKING:
     from pyepo.data.dataset import optDataset
@@ -115,6 +115,11 @@ class SPOPlusFunc(Function):
         # _check_sol(c, w, z)
         # solve
         sol, obj = _solve_or_cache(2 * cp - c, module)
+        # drop the bare objective constant
+        offset = objective_offset(module.optmodel)
+        if offset:
+            obj = obj - offset
+            z = z - offset
         # calculate loss
         if is_minimize(module.optmodel.modelSense):
             loss = -obj + 2 * torch.einsum("bi,bi->b", cp, w) - z.squeeze(dim=-1)

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/func/utils.py

@@ -20,7 +20,7 @@ from pyepo.func._common import (
     validate_positive_int,
 )
 from pyepo.model.mpax import optMpaxModel
-from pyepo.utils import costToNumpy
+from pyepo.utils import costToNumpy, objective_offset
 
 if TYPE_CHECKING:
     from pyepo.func.abcmodule import optModule
@@ -87,8 +87,9 @@ def _solve_batch(
     """
     A function to solve optimization in the forward/backward pass
     """
-    # get device
+    # get device and dtype
     device = cp.device if isinstance(cp, torch.Tensor) else torch.device("cpu")
+    dtype = cp.dtype if isinstance(cp, torch.Tensor) else torch.float32
     # MPAX batch solving
     if isinstance(optmodel, optMpaxModel):
         # get params
@@ -113,11 +114,17 @@ def _solve_batch(
         # obj sense
         if not is_minimize(optmodel.modelSense):
             obj = -obj
+        # add the bare objective constant
+        offset = objective_offset(optmodel)
+        if offset:
+            obj = obj + offset
+        # match input dtype
+        sol, obj = sol.to(dtype), obj.to(dtype)
     # host solving on numpy costs
     else:
         sol_np, obj_np = _solve_batch_np(costToNumpy(cp), optmodel, processes, pool)
-        sol = torch.as_tensor(sol_np).to(device)
-        obj = torch.as_tensor(obj_np).to(device)
+        sol = torch.as_tensor(sol_np).to(device=device, dtype=dtype)
+        obj = torch.as_tensor(obj_np).to(device=device, dtype=dtype)
     return sol, obj
 
 
@@ -130,6 +137,8 @@ def _solve_batch_np(
     """
     A function to solve a batch of numpy costs on the host, shared by both frontends
     """
+    # match input precision
+    out_dtype = np.result_type(cp.dtype, np.float32)
     # single-core
     if processes == 1:
         sol_list: list = []
@@ -140,13 +149,13 @@ def _solve_batch_np(
             sol_list.append(solp)
             obj_list.append(objp)
         # stack + dtype convert in a single call
-        sol = np.asarray(sol_list, dtype=np.float32)
-        obj = np.asarray(obj_list, dtype=np.float32)
+        sol = np.asarray(sol_list, dtype=out_dtype)
+        obj = np.asarray(obj_list, dtype=out_dtype)
     # multi-core (workers pre-loaded with optmodel via pool initializer)
     else:
         res = pool.amap(_solve_with_obj_in_worker, cp).get()
-        sol = np.stack([r[0] for r in res]).astype(np.float32)
-        obj = np.asarray([r[1] for r in res], dtype=np.float32)
+        sol = np.stack([r[0] for r in res]).astype(out_dtype)
+        obj = np.asarray([r[1] for r in res], dtype=out_dtype)
     return sol, obj
 
 
@@ -163,6 +172,9 @@ def _update_solution_pool(
         return torch.unique(sol, dim=0).clone()
     if sol.device != solpool.device:
         sol = sol.to(solpool.device)
+    # match pool dtype
+    if solpool.dtype != sol.dtype:
+        solpool = solpool.to(sol.dtype)
     sol_uniq = torch.unique(sol, dim=0)
     # capped L1-tolerance dedup: first-order solvers re-emit near-identical vertices
     tol = solution_pool_tolerance(sol.shape[-1])
@@ -181,15 +193,19 @@ def _cache_in_pass(
     """
     A function to use solution pool in the forward/backward pass
     """
-    # move solpool to the correct device
-    if solpool.device != cp.device:
-        solpool = solpool.to(cp.device)
+    # move solpool to the correct device and dtype
+    if solpool.device != cp.device or solpool.dtype != cp.dtype:
+        solpool = solpool.to(device=cp.device, dtype=cp.dtype)
     # best solution in pool
     solpool_obj = torch.matmul(cp, solpool.T)
     select = torch.argmin if is_minimize(optmodel.modelSense) else torch.argmax
     ind = select(solpool_obj, dim=1)
     obj = solpool_obj.gather(1, ind.view(-1, 1)).squeeze(1)
     sol = solpool[ind]
+    # add the bare objective constant
+    offset = objective_offset(optmodel)
+    if offset:
+        obj = obj + offset
     return sol, obj, solpool
 
 

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/metric/_common.py

@@ -30,12 +30,6 @@ def regret_from_objective(obj, true_obj, model_sense):
     raise ValueError("Invalid modelSense.")
 
 
-def objective_offset(optmodel: optModel) -> float:
-    """Return a compiled DSL problem's bare objective constant, if present."""
-    problem = getattr(optmodel, "problem", None)
-    return float(problem.obj_offset) if problem is not None else 0.0
-
-
 def require_linear_objective(optmodel: optModel) -> None:
     """Reject models carrying a quadratic objective term."""
     problem = getattr(optmodel, "problem", None)

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/metric/mse.py

@@ -5,15 +5,13 @@ Mean Squared Error
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 import torch
 
 from pyepo.metric._common import torch_evaluation, validate_prediction_batch
 
 if TYPE_CHECKING:
-    from collections.abc import Sized
-
     from torch import nn
     from torch.utils.data import DataLoader
 
@@ -24,19 +22,22 @@ def MSE(predmodel: nn.Module, dataloader: DataLoader) -> float:
 
     Args:
         predmodel: a regression neural network for cost prediction
-        dataloader: Torch dataloader from optDataSet
+        dataloader: Torch dataloader from optDataSet (fields beyond
+            ``(x, c, w, z)`` are ignored)
 
     Returns:
         float: MSE loss
     """
     loss = 0
+    total = 0
     with torch_evaluation(predmodel) as device, torch.no_grad():
         # load data
         for data in dataloader:
-            x, c, _, _ = data
+            x, c, _, _ = data[:4]
             x, c = x.to(device), c.to(device)
             # predict
             cp = predmodel(x)
             validate_prediction_batch(cp, c)
             loss += ((cp - c) ** 2).mean(dim=1).sum().item()
-    return loss / len(cast("Sized", dataloader.dataset))
+            total += x.shape[0]
+    return loss / total if total else 0.0

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/metric/regret.py

@@ -15,14 +15,13 @@ from pyepo.func.runtime import create_solver_pool, normalize_processes
 from pyepo.func.utils import _close_pool, _solve_batch
 from pyepo.metric._common import (
     normalize_regret,
-    objective_offset,
     regret_from_objective,
     require_linear_objective,
     torch_evaluation,
     validate_cost_vectors,
     validate_prediction_batch,
 )
-from pyepo.utils import costToNumpy
+from pyepo.utils import costToNumpy, objective_offset
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -65,7 +64,8 @@ def regret(
             JAX callable ``f(x_numpy) -> cost_array``
         optmodel: a PyEPO optimization model
         dataloader: PyTorch DataLoader over an ``optDataset`` (yielding
-            ``(x, c, w, z)`` tuples)
+            ``(x, c, w, z, ...)`` tuples; fields beyond the first four,
+            e.g. CaVE tight constraints, are ignored)
         processes: number of processors, 1 for single-core, 0 for all of
             cores; a fresh worker pool is spawned per call, each worker
             rebuilding the model from its constructor args
@@ -89,7 +89,7 @@ def regret(
         with torch_evaluation(torch_model) as device:
             # load data
             for data in dataloader:
-                x, c, _, z = data
+                x, c, _, z = data[:4]
                 if torch_model is not None:
                     x, c, z = x.to(device), c.to(device), z.to(device)
                     with torch.no_grad():

{pyepo-2.2.4 → pyepo-2.2.6}/pyepo/metric/unambregret.py

@@ -14,7 +14,6 @@ import torch
 from pyepo import EPO
 from pyepo.metric._common import (
     normalize_regret,
-    objective_offset,
     regret_from_objective,
     require_linear_objective,
     torch_evaluation,
@@ -23,7 +22,7 @@ from pyepo.metric._common import (
     validate_retry_count,
     validate_tolerance,
 )
-from pyepo.utils import costToNumpy
+from pyepo.utils import costToNumpy, objective_offset
 
 if TYPE_CHECKING:
     from torch import nn
@@ -57,7 +56,8 @@ def unambRegret(
     Args:
         predmodel: a regression neural network for cost prediction
         optmodel: a PyEPO optimization model
-        dataloader: PyTorch DataLoader over an ``optDataset``
+        dataloader: PyTorch DataLoader over an ``optDataset`` (fields beyond
+            ``(x, c, w, z)`` are ignored)
         tolerance: precision used when rounding predicted costs to find ties
         max_iter: maximum number of solve retries with relaxed tolerance
 
@@ -72,8 +72,8 @@ def unambRegret(
     with torch_evaluation(predmodel) as device:
         # load data
         for data in dataloader:
-            x, c, w, z = data
-            x, c, w, z = x.to(device), c.to(device), w.to(device), z.to(device)
+            x, c, _, z = data[:4]
+            x, c, z = x.to(device), c.to(device), z.to(device)
             # pred cost
             with torch.no_grad():
                 cp = costToNumpy(predmodel(x))