pyepo 2.2.3__tar.gz → 2.2.5__tar.gz

{pyepo-2.2.3 → pyepo-2.2.5}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.3
-Summary: PyTorch-based End-to-End Predict-then-Optimize Tool
+Version: 2.2.5
+Summary: PyTorch/JAX-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT
 Project-URL: Homepage, https://github.com/khalil-research/PyEPO
@@ -9,7 +9,7 @@ Project-URL: Documentation, https://khalil-research.github.io/PyEPO
 Project-URL: Repository, https://github.com/khalil-research/PyEPO
 Project-URL: Issues, https://github.com/khalil-research/PyEPO/issues
 Project-URL: Paper, https://link.springer.com/article/10.1007/s12532-024-00255-x
-Keywords: predict-then-optimize,end-to-end,decision-focused learning,optimization,deep learning,pytorch,linear programming,integer programming
+Keywords: predict-then-optimize,end-to-end,decision-focused learning,optimization,deep learning,pytorch,jax,linear programming,integer programming
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10

{pyepo-2.2.3 → pyepo-2.2.5}/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
@@ -315,7 +315,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>`_
@@ -455,6 +456,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.3 → pyepo-2.2.5}/pyepo/func/cave.py

@@ -59,7 +59,8 @@ 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>`_

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/bases.py

@@ -16,7 +16,6 @@ from __future__ import annotations
 
 import math
 from collections import defaultdict
-from copy import deepcopy
 from itertools import combinations
 from numbers import Integral, Real
 from typing import TYPE_CHECKING
@@ -105,9 +104,6 @@ class shortestPathBase(optModel):
         self.arcs = _get_grid_arcs(self.grid)
         super().__init__(*args, **kwargs)
 
-    def get_config(self) -> dict:
-        return {**super().get_config(), "grid": self.grid}
-
     @property
     def num_cost(self) -> int:
         return len(self.arcs)
@@ -208,14 +204,6 @@ class portfolioBase(optModel):
             raise ValueError("gamma must be greater than or equal to zero.")
         super().__init__(*args, **kwargs)
 
-    def get_config(self) -> dict:
-        return {
-            **super().get_config(),
-            "num_assets": self.num_assets,
-            "covariance": self.covariance.copy(),
-            "gamma": self.gamma,
-        }
-
     @property
     def num_cost(self) -> int:
         return self.num_assets
@@ -263,9 +251,6 @@ class tspABBase(optModel):
         self._extra_constrs: list = []
         super().__init__(*args, **kwargs)
 
-    def get_config(self) -> dict:
-        return {**super().get_config(), "num_nodes": self.num_nodes}
-
     @property
     def num_cost(self) -> int:
         # use edges; backend's self.x has 2*num_edges directed Vars
@@ -381,15 +366,6 @@ class vrpABBase(optModel):
         self._extra_constrs: list = []
         super().__init__(*args, **kwargs)
 
-    def get_config(self) -> dict:
-        return {
-            **super().get_config(),
-            "num_nodes": self.num_nodes,
-            "demands": deepcopy(self.demands),
-            "capacity": self.capacity,
-            "num_vehicle": self.num_vehicle,
-        }
-
     @property
     def num_cost(self) -> int:
         # one predicted cost per undirected edge

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/grb/grbmodel.py

@@ -171,7 +171,8 @@ class optGrbModel(optModel):
         else:
             # LinExpr(coeffs, vars) builds the affine expression in one C call
             vars_list = new_model._vars_list
-            assert vars_list is not None
+            if vars_list is None:
+                raise RuntimeError("Gurobi variable list is unavailable.")
             expr = gp.LinExpr(coefs_np.tolist(), vars_list) <= rhs
             new_model._model.addConstr(expr)
         # track for replay on relax

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/grb/tsp.py

@@ -178,9 +178,6 @@ class tspDFJModel(tspABModel):
         self._recycled_keys: set = set()
         super().__init__(num_nodes, *args, **kwargs)
 
-    def get_config(self) -> dict:
-        return {**super().get_config(), "recycle_cuts": self.recycle_cuts}
-
     def _getModel(self) -> tuple:
         """
         A method to build Gurobi model

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/grb/vrp.py

@@ -70,9 +70,6 @@ class vrpRCIModel(vrpABModel):
         self._recycled_keys: set = set()
         super().__init__(num_nodes, demands, capacity, num_vehicle)
 
-    def get_config(self) -> dict:
-        return {**super().get_config(), "recycle_cuts": self.recycle_cuts}
-
     def _getModel(self) -> tuple:
         """
         A method to build Gurobi model

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/opt.py

@@ -5,6 +5,8 @@ Abstract optimization model
 
 from __future__ import annotations
 
+import functools
+import inspect
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
@@ -25,20 +27,63 @@ class ModelSpec:
     """Serializable recipe for building a fresh optimization model."""
 
     model_type: type[optModel]
+    _args: tuple
     _config: dict
 
-    def __init__(self, model_type: type[optModel], config: dict) -> None:
+    def __init__(
+        self,
+        model_type: type[optModel],
+        config: dict,
+        args: tuple = (),
+    ) -> None:
         object.__setattr__(self, "model_type", model_type)
+        object.__setattr__(self, "_args", deepcopy(args))
         object.__setattr__(self, "_config", deepcopy(config))
 
+    @property
+    def args(self) -> tuple:
+        """Return an independent copy of positional constructor arguments."""
+        return deepcopy(self._args)
+
     @property
     def config(self) -> dict:
-        """Return an independent copy of the constructor configuration."""
+        """Return an independent copy of keyword constructor arguments."""
         return deepcopy(self._config)
 
     def build(self) -> optModel:
         """Build a fresh model without sharing mutable configuration values."""
-        return self.model_type.from_config(self._config)
+        return self.model_type.from_config(self._config, self._args)
+
+
+def _snapshot(value):
+    """Deep-copy a constructor argument, keeping the reference if it cannot be copied."""
+    try:
+        return deepcopy(value)
+    except Exception:  # noqa: BLE001 -- any copy failure falls back to a reference
+        return value
+
+
+def _capture_init_config(init, args, kwargs) -> tuple[tuple, dict]:
+    """Flatten a constructor call into arguments that rebuild the model."""
+    sig = inspect.signature(init)
+    bound = sig.bind(None, *args, **kwargs)
+    init_args = []
+    config = {}
+    for i, (name, value) in enumerate(bound.arguments.items()):
+        # the first bound argument is self
+        if i == 0:
+            continue
+        kind = sig.parameters[name].kind
+        # snapshot each argument; values that cannot be deep-copied keep a reference
+        if kind is inspect.Parameter.VAR_KEYWORD:
+            config.update({k: _snapshot(v) for k, v in value.items()})
+        elif kind is inspect.Parameter.VAR_POSITIONAL:
+            init_args.extend(_snapshot(v) for v in value)
+        elif kind is inspect.Parameter.POSITIONAL_ONLY:
+            init_args.append(_snapshot(value))
+        else:
+            config[name] = _snapshot(value)
+    return tuple(init_args), config
 
 
 class optModel(ABC):
@@ -53,11 +98,6 @@ class optModel(ABC):
     and MPAX (``optMpaxModel``); subclass ``optModel`` directly to integrate
     any other solver or algorithm.
 
-    Models that take constructor arguments should override ``get_config`` and
-    cooperatively merge ``super().get_config()``. The resulting configuration
-    powers ``rebuild()``, multiprocessing workers, and sklearn scorers without
-    inspecting constructor signatures or runtime solver state.
-
     The default objective sense is minimization; set
     ``self.modelSense = EPO.MAXIMIZE`` in ``_getModel`` or ``__init__`` for
     maximization problems (some backends, e.g. Gurobi and COPT, detect this
@@ -73,6 +113,24 @@ class optModel(ABC):
     arcs: list
     _cost_vars: list
 
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        # Only wrap subclasses that define their own __init__ and use the
+        # default reconstruction config. Custom get_config implementations may
+        # intentionally accept objects that are not deepcopyable/rebuildable.
+        if "__init__" not in cls.__dict__ or "get_config" in cls.__dict__:
+            return
+        user_init = cls.__init__
+
+        @functools.wraps(user_init)
+        def _init_capturing(self, *args, **kwargs):
+            # record only the outermost call; nested super().__init__ leaves it intact
+            if "_init_config" not in self.__dict__:
+                self._init_args, self._init_config = _capture_init_config(user_init, args, kwargs)
+            user_init(self, *args, **kwargs)
+
+        cls.__init__ = _init_capturing
+
     def __init__(self) -> None:
         # Cache for models whose solver variables do not map one-to-one to
         # predicted costs (for example directed TSP/VRP formulations).
@@ -86,17 +144,21 @@ class optModel(ABC):
         return "optModel " + self.__class__.__name__
 
     def get_config(self) -> dict:
-        """Return the explicit constructor configuration for this model."""
-        return {}
+        """Return the constructor configuration for this model."""
+        return deepcopy(self.__dict__.get("_init_config", {}))
 
     @classmethod
-    def from_config(cls, config: dict) -> Self:
+    def from_config(cls, config: dict, args: tuple = ()) -> Self:
         """Build a model from a configuration produced by ``get_config``."""
-        return cls(**deepcopy(config))
+        return cls(*deepcopy(args), **deepcopy(config))
 
     def to_spec(self) -> ModelSpec:
         """Return a serializable, immutable-snapshot rebuild recipe."""
-        return ModelSpec(type(self), self.get_config())
+        return ModelSpec(
+            type(self),
+            self.get_config(),
+            self.__dict__.get("_init_args", ()),
+        )
 
     def rebuild(self) -> Self:
         """Build a structurally equivalent model with clean runtime state."""

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/ort/compile.py

@@ -40,7 +40,6 @@ class compiledOrtProblem(compiledBase, optOrtModel):
         self.problem = deepcopy(problem)
         self.params = dict(params) if params else {}
         self.solver = solver
-        self._extra_constrs = []  # (coef, rhs) cuts replayed on copy
         optModel.__init__(self)  # builds the model via _getModel
         self._model.SuppressOutput()
         self._set_obj_sense()

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/ort/ortcpmodel.py

@@ -53,7 +53,6 @@ class optOrtCpModel(optModel):
             raise ImportError(
                 "OR-Tools is not installed. Please install ortools to use this feature."
             )
-        self._extra_constrs = []
         self._objective_coefs: list[int] | None = None
         super().__init__()
         # cache ordered Var list for batched weighted_sum / per-Var Value() loop

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/model/ort/ortmodel.py

@@ -54,7 +54,6 @@ class optOrtModel(optModel):
                 "OR-Tools is not installed. Please install ortools to use this feature."
             )
         self.solver = solver
-        self._extra_constrs = []
         super().__init__()
         # suppress output
         self._model.SuppressOutput()

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo/utils.py

@@ -21,7 +21,7 @@ _EPS: float = 1e-8
 
 def getArgs(model: optModel) -> dict:
     """
-    Compatibility wrapper for the explicit model configuration protocol.
+    Compatibility wrapper for model reconstruction configuration.
 
     Args:
         model: optimization model

{pyepo-2.2.3 → pyepo-2.2.5}/pyepo.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.3
-Summary: PyTorch-based End-to-End Predict-then-Optimize Tool
+Version: 2.2.5
+Summary: PyTorch/JAX-based End-to-End Predict-then-Optimize Tool
 Author-email: Bo Tang <bolucas.tang@mail.utoronto.ca>
 License: MIT
 Project-URL: Homepage, https://github.com/khalil-research/PyEPO
@@ -9,7 +9,7 @@ Project-URL: Documentation, https://khalil-research.github.io/PyEPO
 Project-URL: Repository, https://github.com/khalil-research/PyEPO
 Project-URL: Issues, https://github.com/khalil-research/PyEPO/issues
 Project-URL: Paper, https://link.springer.com/article/10.1007/s12532-024-00255-x
-Keywords: predict-then-optimize,end-to-end,decision-focused learning,optimization,deep learning,pytorch,linear programming,integer programming
+Keywords: predict-then-optimize,end-to-end,decision-focused learning,optimization,deep learning,pytorch,jax,linear programming,integer programming
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10

{pyepo-2.2.3 → pyepo-2.2.5}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyepo"
-version = "2.2.3"
-description = "PyTorch-based End-to-End Predict-then-Optimize Tool"
+version = "2.2.5"
+description = "PyTorch/JAX-based End-to-End Predict-then-Optimize Tool"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { text = "MIT" }
 authors = [{ name = "Bo Tang", email = "bolucas.tang@mail.utoronto.ca" }]
@@ -17,6 +17,7 @@ keywords = [
     "optimization",
     "deep learning",
     "pytorch",
+    "jax",
     "linear programming",
     "integer programming",
 ]

{pyepo-2.2.3 → pyepo-2.2.5}/tests/test_10_utils.py

@@ -19,7 +19,7 @@ from .conftest import requires_gurobi
 
 
 class ConfigModel(optModel):
-    """Solver-free model for the explicit reconstruction protocol."""
+    """Solver-free model with custom reconstruction config."""
 
     def __init__(self, values, label="default"):
         self.values = np.asarray(values)
@@ -43,6 +43,88 @@ class ConfigModel(optModel):
         return np.zeros(len(self.values)), 0.0
 
 
+class AutoConfigModel(optModel):
+    """Solver-free model that relies on optModel's automatic config capture."""
+
+    def __init__(self, values, **kwargs):
+        self.values = values
+        self.kwargs = kwargs
+        super().__init__()
+
+    def _getModel(self):
+        return None, list(range(len(self.values)))
+
+    def setObj(self, c):
+        self.cost = np.asarray(c)
+
+    def solve(self):
+        return np.zeros(len(self.values)), 0.0
+
+
+class VarArgsConfigModel(optModel):
+    """Solver-free model whose constructor needs positional replay."""
+
+    def __init__(self, *values, label="default"):
+        self.values = values
+        self.label = label
+        super().__init__()
+
+    def _getModel(self):
+        return None, list(range(len(self.values)))
+
+    def setObj(self, c):
+        self.cost = np.asarray(c)
+
+    def solve(self):
+        return np.zeros(len(self.values)), 0.0
+
+
+class PosOnlyConfigModel(optModel):
+    """Solver-free model with a positional-only constructor argument."""
+
+    def __init__(self, values, /, label="default"):
+        self.values = values
+        self.label = label
+        super().__init__()
+
+    def _getModel(self):
+        return None, list(range(len(self.values)))
+
+    def setObj(self, c):
+        self.cost = np.asarray(c)
+
+    def solve(self):
+        return np.zeros(len(self.values)), 0.0
+
+
+class _NoDeepcopy:
+    def __init__(self, values):
+        self.values = values
+
+    def __deepcopy__(self, memo):
+        raise TypeError("not deepcopyable")
+
+
+class CustomConfigModel(optModel):
+    """Custom config should control reconstruction for unusual constructor inputs."""
+
+    def __init__(self, resource=None, values=None):
+        self.values = list(values if values is not None else resource.values)
+        super().__init__()
+
+    def get_config(self):
+        return {"values": self.values.copy()}
+
+    def _getModel(self):
+        return None, list(range(len(self.values)))
+
+    def setObj(self, c):
+        self.cost = np.asarray(c)
+
+    def solve(self):
+        return np.zeros(len(self.values)), 0.0
+
+
 # ============================================================
 # unionFind (pure)
 # ============================================================
@@ -161,7 +243,7 @@ class TestCostToNumpy:
 
 
 # ============================================================
-# explicit model reconstruction protocol (pure)
+# model reconstruction config (pure)
 # ============================================================
 
 
@@ -222,6 +304,54 @@ class TestModelSpec:
         np.testing.assert_array_equal(sol, [0.0, 0.0, 0.0])
         assert obj == 0.0
 
+    def test_auto_config_snapshots_constructor_inputs(self):
+        values = [1, 2, 3]
+        nested = {"tag": ["x"]}
+        model = AutoConfigModel(values, nested=nested)
+        values[0] = 99
+        nested["tag"][0] = "changed"
+
+        rebuilt = model.rebuild()
+
+        assert rebuilt.values == [1, 2, 3]
+        assert rebuilt.kwargs == {"nested": {"tag": ["x"]}}
+
+    def test_auto_config_export_is_independent(self):
+        model = AutoConfigModel([1, 2, 3], nested={"tag": ["x"]})
+        config = model.get_config()
+        config["values"][0] = 99
+        config["nested"]["tag"][0] = "changed"
+
+        assert model.get_config()["values"] == [1, 2, 3]
+        assert model.get_config()["nested"] == {"tag": ["x"]}
+        assert model.rebuild().kwargs == {"nested": {"tag": ["x"]}}
+
+    def test_auto_config_replays_varargs(self):
+        model = VarArgsConfigModel(1, 2, 3, label="x")
+        spec = model.to_spec()
+
+        assert spec.args == (1, 2, 3)
+        assert spec.config == {"label": "x"}
+        rebuilt = model.rebuild()
+        assert rebuilt.values == (1, 2, 3)
+        assert rebuilt.label == "x"
+
+    def test_auto_config_replays_positional_only_args(self):
+        model = PosOnlyConfigModel([1, 2, 3], label="x")
+        spec = model.to_spec()
+
+        assert spec.args == ([1, 2, 3],)
+        assert spec.config == {"label": "x"}
+        rebuilt = model.rebuild()
+        assert rebuilt.values == [1, 2, 3]
+        assert rebuilt.label == "x"
+
+    def test_custom_get_config_accepts_uncopyable_constructor_input(self):
+        model = CustomConfigModel(_NoDeepcopy([1, 2, 3]))
+        rebuilt = model.rebuild()
+
+        assert rebuilt.values == [1, 2, 3]
+
 
 # ============================================================
 # getArgs (needs a real optModel)

{pyepo-2.2.3 → pyepo-2.2.5}/tests/test_40_dataset.py

@@ -14,6 +14,7 @@ import torch
 from pyepo.data import shortestpath
 from pyepo.data.dataset import (
     collate_tight_constraints,
+    optDataLoader,
     optDataset,
     optDatasetConstrs,
     optDatasetKNN,
@@ -299,3 +300,51 @@ class TestCollateTightConstraints:
         assert padded.shape == (2, 5, 4)
         # the shorter matrix is zero-padded at the tail
         assert torch.allclose(padded[0, 2:], torch.zeros(3, 4))
+
+
+class TestOptDataLoader:
+    """Pure: optDataLoader applies a dataset's collate_fn automatically."""
+
+    def test_constrs_collate_fn_wired(self):
+        assert optDatasetConstrs.collate_fn is collate_tight_constraints
+
+    def test_uses_dataset_collate_fn(self):
+        class _Ragged(list):
+            collate_fn = staticmethod(lambda batch: ("auto", len(batch)))
+
+        loader = optDataLoader(_Ragged([0, 1, 2, 3]), batch_size=2)
+        assert next(iter(loader)) == ("auto", 2)
+
+    def test_accepts_positional_dataloader_args(self):
+        loader = optDataLoader([0, 1, 2, 3], 2)
+        assert torch.equal(next(iter(loader)), torch.tensor([0, 1]))
+
+    def test_explicit_collate_fn_wins(self):
+        class _Ragged(list):
+            collate_fn = staticmethod(lambda batch: ("auto", len(batch)))
+
+        loader = optDataLoader(
+            _Ragged([0, 1, 2, 3]), batch_size=2, collate_fn=lambda b: ("explicit", len(b))
+        )
+        assert next(iter(loader)) == ("explicit", 2)
+
+    def test_positional_collate_fn_wins(self):
+        class _Ragged(list):
+            collate_fn = staticmethod(lambda batch: ("auto", len(batch)))
+
+        loader = optDataLoader(
+            _Ragged([0, 1, 2, 3]),
+            2,
+            None,
+            None,
+            None,
+            0,
+            lambda b: ("positional", len(b)),
+        )
+        assert next(iter(loader)) == ("positional", 2)
+
+    def test_plain_dataset_uses_default_collate(self):
+        data = [(torch.tensor([float(i)]),) for i in range(4)]
+        loader = optDataLoader(data, batch_size=2)
+        (batch,) = next(iter(loader))
+        assert batch.shape == (2, 1)