pyepo 2.2.4__tar.gz → 2.2.5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.4
+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

{pyepo-2.2.4 → 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.4 → 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.4 → pyepo-2.2.5}/pyepo/model/opt.py

@@ -27,41 +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) -> dict:
-    """Flatten a constructor call into keyword arguments that rebuild the model."""
+
+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
-        # **kwargs: merge captured keywords in directly
+        # snapshot each argument; values that cannot be deep-copied keep a reference
         if kind is inspect.Parameter.VAR_KEYWORD:
-            config.update(deepcopy(value))
-        # nameless positionals cannot replay by keyword; override get_config to keep them
-        elif kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.POSITIONAL_ONLY):
-            continue
+            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] = deepcopy(value)
-    return config
+            config[name] = _snapshot(value)
+    return tuple(init_args), config
 
 
 class optModel(ABC):
@@ -93,8 +115,10 @@ class optModel(ABC):
 
     def __init_subclass__(cls, **kwargs) -> None:
         super().__init_subclass__(**kwargs)
-        # only wrap a subclass that defines its own __init__
-        if "__init__" not in cls.__dict__:
+        # 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__
 
@@ -102,7 +126,7 @@ class optModel(ABC):
         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_config = _capture_init_config(user_init, args, kwargs)
+                self._init_args, self._init_config = _capture_init_config(user_init, args, kwargs)
             user_init(self, *args, **kwargs)
 
         cls.__init__ = _init_capturing
@@ -124,13 +148,17 @@ class optModel(ABC):
         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.4 → 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.4 → 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.4 → 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.4 → 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.4 → pyepo-2.2.5}/pyepo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyepo
-Version: 2.2.4
+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

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyepo"
-version = "2.2.4"
+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" }

{pyepo-2.2.4 → 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)
@@ -61,6 +61,70 @@ class AutoConfigModel(optModel):
         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)
 # ============================================================
@@ -179,7 +243,7 @@ class TestCostToNumpy:
 
 
 # ============================================================
-# explicit model reconstruction protocol (pure)
+# model reconstruction config (pure)
 # ============================================================
 
 
@@ -262,6 +326,32 @@ class TestModelSpec:
         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.4 → 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)