diff --git a/src/lpspec/relational/sinks/solvers/gurobi.py b/src/lpspec/relational/sinks/solvers/gurobi.py index 85d101be2..b42a01f9b 100644 --- a/src/lpspec/relational/sinks/solvers/gurobi.py +++ b/src/lpspec/relational/sinks/solvers/gurobi.py @@ -1,8 +1,8 @@ """The ``gurobi`` solver: the model in two calls, straight into gurobipy. -Reads the same ``dense_columns`` as -:mod:`~lpspec.relational.sinks.solvers.highs`, so the two cannot disagree -about the model they load. Two things differ: +The same hand-off as :mod:`~lpspec.relational.sinks.solvers.highs`, reading the +same ``dense_columns``, ``dense_rows`` and ``row_blocks``, so the two cannot +disagree about the model they load. Two things differ: - **The matrix's currency.** HiGHS takes the three CSR arrays; gurobipy's matrix API takes a matrix *object*, so they are wrapped in a @@ -11,7 +11,7 @@ - **Nothing is batched.** The columns cannot be, since ``addMConstr`` writes into one ``MVar`` spanning the model — and the matrix *should* not be, which is where this sink parts company with the HiGHS one. See - :func:`_row_blocks`. + :meth:`~lpspec.relational.sinks.tables.ModelTables.row_blocks`. ``gurobipy`` and ``scipy`` are imported inside the functions, so importing this module stays free for a caller who never solves with it. @@ -25,10 +25,11 @@ import polars as pl from lpspec.errors import LpspecError +from lpspec.relational.sinks.tables import SENSE_CODES from lpspec.relational.status import SolveStatus if TYPE_CHECKING: - from collections.abc import Iterator, Mapping + from collections.abc import Mapping from lpspec.relational.sinks.tables import ModelTables @@ -79,7 +80,8 @@ def build_gurobi( :func:`~lpspec.relational.sinks.solvers.highs.build_highs`'s seam, drawn for its reason: the search is the same work whoever filled the model. ``batch_rows`` is a *nonzero* budget that splits the matrix across calls; - it defaults to one call, and see :func:`_row_blocks` for why. + it defaults to one call — see + :meth:`~lpspec.relational.sinks.tables.ModelTables.row_blocks` for why. **The caller owns the model, so the environment follows it.** gurobipy has no ``Model.getEnv()``, so a caller handed only the model could never @@ -157,15 +159,17 @@ def _load( discrete: dict[str, Any] = {'vtype': np.where(integral, 'I', 'C')} if integral.any() else {} x = m.addMVar(model.column_count, lb=lb, ub=ub, obj=cost, **discrete) - senses = np.array([gurobipy.GRB.LESS_EQUAL, gurobipy.GRB.GREATER_EQUAL, gurobipy.GRB.EQUAL], dtype=' Iterator[tuple[pl.DataFrame, pl.DataFrame]]: - """The matrix as ``(rows, entries)`` pairs — **one pair unless asked**. +#: Our spelling of a comparison against Gurobi's, by ``GRB`` attribute name — +#: a name rather than a value because ``gurobipy`` is an optional import and +#: this is module level. +_GUROBI_SENSE = {'<=': 'LESS_EQUAL', '>=': 'GREATER_EQUAL', '==': 'EQUAL'} - Where the HiGHS sink batches by default, this one does not, and the - difference is measured rather than inherited. ``addMConstr`` costs about - 42 ns per *model* column per call, whatever the block holds: at 400k - columns that is 17 ms of overhead per extra call, so splitting a 2M-nonzero - matrix into 20 blocks took 0.55 s against 0.23 s in one, and 40 blocks took - 0.89 s. The cost scales with the model's width, so the wider the model the - worse batching is. - Batching bought nothing back. It cannot bound the peak here — Gurobi - accumulates the whole model as it goes, so the only thing a block saves is - one block-sized temporary: 0.62 GB against 0.65 GB on that model, for twice - the hand-off time. +def _spelled(gurobipy: Any) -> Any: + """:data:`SENSE_CODES` as the characters ``addMConstr`` wants, by code. - ``batch_rows`` still splits when a caller asks, which is how the tests - exercise block boundaries — the arithmetic that puts a row's entries in the - right ``indptr`` slot is worth a test even when nothing runs it by default. + Built from the mapping rather than written out in its order, so the two + cannot come to disagree: a wrong order here is a model whose comparisons + are silently permuted, which every solver would answer confidently. A + sense added to :data:`SENSE_CODES` and not to :data:`_GUROBI_SENSE` raises + instead. """ - ordered_rows = model.rows.sort('row').select( - 'row', - pl.col('sense').replace_strict({'<=': 0, '>=': 1, '==': 2}, return_dtype=pl.UInt8).alias('op'), - 'rhs', - ) - ordered_matrix = model.matrix.sort('row') - if batch_rows is None: - yield ordered_rows, ordered_matrix - return - for lo, hi in model.row_chunks_by_nonzeros(batch_rows): - yield ( - ordered_rows.filter(pl.col('row').is_between(lo, hi, closed='left')), - ordered_matrix.filter(pl.col('row').is_between(lo, hi, closed='left')), - ) + import numpy as np + + spelling = np.empty(len(SENSE_CODES), dtype=' Any: diff --git a/src/lpspec/relational/sinks/solvers/highs.py b/src/lpspec/relational/sinks/solvers/highs.py index cd8e86e6e..9b7ec1835 100644 --- a/src/lpspec/relational/sinks/solvers/highs.py +++ b/src/lpspec/relational/sinks/solvers/highs.py @@ -6,8 +6,9 @@ numpy slices, in batches. **Nothing textual crosses into numpy.** A polars ``String`` column converts by -boxing every value as a Python object, so a comparison against ``'<='`` is made -in polars and only the bound it selects is handed over — the same rule +boxing every value as a Python object, so a row's ``'<='`` becomes a +:data:`~lpspec.relational.sinks.tables.SENSE_CODES` byte before it is read here +— the same rule :meth:`~lpspec.relational.sinks.tables.ModelTables.dense_columns` states for the column vectors, where it was measured. @@ -22,6 +23,7 @@ import polars as pl from lpspec.errors import LpspecError +from lpspec.relational.sinks.tables import SENSE_CODES from lpspec.relational.status import SolveStatus if TYPE_CHECKING: @@ -39,9 +41,9 @@ #: #: Deliberately small, and the opposite of what the same budget settled at on #: the duckdb engine. There a wider chunk was worth a third of the hand-off, -#: because every chunk re-ran an ordered query; here the columns are numpy -#: slices and the rows a filter over an already-sorted frame, so more chunks -#: cost almost nothing and only residency scales with the budget. Measured at +#: because every chunk re-ran an ordered query; here both columns and rows are +#: numpy slices, so more chunks cost almost nothing and only residency scales +#: with the budget. Measured at #: `l`: 2e6 against 1e5 is 0.50s against 0.59s on `dispatch` and 0.71 against #: 0.75 on `transport`, for 0.6 GB and 0.2 GB more peak. A tenth of a second #: on a hand-off that precedes a minute of simplex is not worth half a @@ -112,26 +114,18 @@ def build_highs( integrality = np.full(len(noncontinuous), int(highspy.HighsVarType.kInteger), dtype=np.uint8) h.changeColsIntegrality(len(noncontinuous), noncontinuous, integrality) - ordered_rows = model.rows.sort('row') - ordered_matrix = model.matrix.sort('row') - for lo, hi in model.row_chunks_by_nonzeros(batch): - rows = ordered_rows.filter(pl.col('row').is_between(lo, hi, closed='left')).select( - 'row', - pl.when(pl.col('sense') == '<=').then(-inf).otherwise(pl.col('rhs')).alias('rlb'), - pl.when(pl.col('sense') == '>=').then(inf).otherwise(pl.col('rhs')).alias('rub'), - ) - a = ordered_matrix.filter(pl.col('row').is_between(lo, hi, closed='left')) - rlb = rows['rlb'].to_numpy() - rub = rows['rub'].to_numpy() - starts = np.searchsorted(a['row'].to_numpy(), rows['row'].to_numpy()).astype(np.int32) + sense, rhs = model.dense_rows(inf) + rlb = np.where(sense == SENSE_CODES['<='], -inf, rhs) + rub = np.where(sense == SENSE_CODES['>='], inf, rhs) + for lo, hi, a, starts in model.row_blocks(batch): _loaded( h, h.addRows( - len(rlb), - rlb, - rub, + hi - lo, + rlb[lo:hi], + rub[lo:hi], a.height, - starts, + starts.astype(np.int32), a['col'].to_numpy().astype(np.int32), a['coeff'].to_numpy(), ), diff --git a/src/lpspec/relational/sinks/tables.py b/src/lpspec/relational/sinks/tables.py index 8bf0d80dd..98fc09337 100644 --- a/src/lpspec/relational/sinks/tables.py +++ b/src/lpspec/relational/sinks/tables.py @@ -3,10 +3,10 @@ Four frames plus the scalars a writer needs to size its batching. A sink that needs a fifth thing states it here, where both sides can see it. -Also the one *projection* of them more than one sink needs -(:meth:`ModelTables.dense_columns`). It belongs to the contract rather than to -either solver, because two sinks computing it separately could disagree about -the model they loaded — the one thing neither may do. +Also the *projections* of them more than one sink needs — the dense column and +row vectors, and the matrix a block at a time. They belong to the contract +rather than to either solver, because two sinks computing them separately could +disagree about the model they loaded — the one thing neither may do. """ from __future__ import annotations @@ -30,6 +30,19 @@ npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.bool_] ] + #: A sense code and a right-hand side, each as long as the model has rows. + DenseRows = tuple[npt.NDArray[np.uint8], npt.NDArray[np.float64]] + + #: One chunk of rows: the half-open label range, the matrix entries those + #: rows own, and the offset of each row's entries within them. + RowBlock = tuple[int, int, pl.DataFrame, npt.NDArray[np.int64]] + + +#: ``sense`` as a number, so a row's comparison crosses into numpy as one byte +#: rather than as a boxed Python string. The order is arbitrary and shared: a +#: solver indexes its own spelling with these, so the two cannot drift. +SENSE_CODES = {'<=': 0, '>=': 1, '==': 2} + @dataclass(frozen=True) class ModelTables: @@ -107,9 +120,73 @@ def dense_columns(self, infinity: float) -> DenseColumns: np.nan_to_num(ub, copy=False, neginf=-infinity, posinf=infinity) return lb, ub, cost, integral + def dense_rows(self, infinity: float) -> DenseRows: + """``(sense, rhs)`` as numpy vectors over the solver's row index. + + The row half of :meth:`dense_columns`, and for the same reason: ``row`` + is dense ``0..n-1``, so it *is* the position a value belongs at, and a + chunk of rows is then a slice rather than a search. Sorting the row + frame and filtering it once per chunk read the same 6M rows nine times + over on `fleet/l`, to hand each of them over once. + + It stops at the sense rather than at bounds because that is where the + two solvers part: HiGHS wants a row's ``lower``/``upper`` and Gurobi its + comparison and right-hand side, and both are this pair spelled + differently. :data:`SENSE_CODES` is the spelling neither owns. + + A row with no entry is left with a comparison nothing can fail — + ``>=`` against ``-infinity`` — rather than the ``== 0`` that would be a + real equality the model never stated. + """ + sided = self.rows.select( + 'row', + pl.col('sense').replace_strict(SENSE_CODES, return_dtype=pl.UInt8).alias('op'), + 'rhs', + ) + at = sided['row'].to_numpy() + sense = _scattered(self.row_count, at, sided['op'].to_numpy(), SENSE_CODES['>=']) + rhs = _scattered(self.row_count, at, sided['rhs'].to_numpy(), -infinity) + return sense, rhs + + def row_blocks(self, budget: int | None) -> Iterator[RowBlock]: + """Each chunk of rows with the matrix entries it owns. + + **``budget=None`` is one block, and that is a real answer rather than a + degenerate one.** Whether splitting pays is a property of the API being + fed, not of the model: HiGHS takes a chunk at a time and its budget + bounds the temporary, while Gurobi's ``addMConstr`` charges about 42 ns + per *model column* per call whatever the block holds — 0.23 s in one + call against 0.89 s in forty on the same matrix (#434). So the caller + says, and both answers come out of the same code. + + The matrix is ordered once, and a chunk is then a ``slice`` of it + located by binary search on the label column — the range is contiguous + because ``row`` is dense and the frame is sorted, so scanning for it + would re-read the whole model once per chunk. + + **Searched in polars rather than through numpy.** Pulling the label + column out to search it there is marginally faster and holds a second + copy of one column of the model for the whole loop, which is 0.11 GB at + `transport/l` — the wrong trade in a pass that exists to stay bounded. + + ``starts`` is each row's offset within the block, which is what both + solvers' matrix APIs ask for. A row with no entries takes the next + row's offset, and so occupies no span. + """ + import numpy as np + + ordered = self.matrix.sort('row') + label = ordered['row'] + spans = [(0, self.row_count)] if budget is None else self.row_chunks_by_nonzeros(budget) + for lo, hi in spans: + first = int(label.search_sorted(lo, 'left')) + last = int(label.search_sorted(hi, 'left')) + entries = ordered.slice(first, last - first) + yield lo, hi, entries, np.searchsorted(entries['row'].to_numpy(), np.arange(lo, hi)) + def _scattered(count: int, at: Any, values: Any, absent: Any) -> Any: - """*values* written at the column each one belongs to, *absent* elsewhere.""" + """*values* written at the label each one belongs to, *absent* elsewhere.""" import numpy as np dense = np.full(count, absent, dtype=values.dtype) diff --git a/tests/test_relational.py b/tests/test_relational.py index 56c95feec..8a2c32547 100644 --- a/tests/test_relational.py +++ b/tests/test_relational.py @@ -37,6 +37,7 @@ Variable, VariableDeclaration, ) +from lpspec.relational.sinks import SOLVERS from tests.conftest import by_coord, solve_lp_file from tests.differential import RTOL, differential from tests.oracle import linopy, pd, transport_eager_objective, xr @@ -715,6 +716,36 @@ def test_infinite_bounds_survive_the_handoff(dispatch_data): assert ex.solve().is_ok +@pytest.mark.parametrize('solver_name', sorted(SOLVERS)) +@pytest.mark.parametrize('batch_rows', [1, 2, 7, 100_000], ids=['one', 'two', 'odd', 'whole']) +def test_a_row_with_no_terms_keeps_its_seat_at_any_chunking(solver_name, batch_rows): + """Rows reach a solver by position, so an empty one still occupies one. + + `where: "t > 0"` leaves `balance` at `t = 0` with nothing to sum, and this + lane keeps the row: `0 == 5` is infeasible and says so. Both solvers now + take the row bounds from a dense vector and slice it per chunk, which is + the same hand-off only if every label in the range has a seat — a row that + fell out of the frame would take a comparison nothing can fail, leave the + constraint unenforced, and the model would come back solved against a + model that cannot be. + + Ragged batches because the range loop is where a seat would be lost, and a + round number is the one split that hides an off-by-one. Both solvers, + because the seating is now theirs jointly rather than either one's. + """ + model = { + 'dimensions': {'t': {'dtype': 'int', 'values': [0, 1, 2]}, 'g': {'values': ['a', 'b']}}, + 'parameters': {'load': {'dims': ['t']}}, + 'variables': {'p': {'foreach': ['t', 'g'], 'where': 't > 0', 'bounds': {'lower': 0, 'upper': 100}}}, + 'constraints': {'balance': {'foreach': ['t'], 'expression': 'sum(p, over=g) == load'}}, + 'objectives': {'o': {'sense': 'minimize', 'expression': 'sum(sum(p, over=g), over=t)'}}, + } + with lps.build(model, {'load': pl.DataFrame({'t': [0, 1, 2], 'value': [5.0, 4.0, 6.0]})}) as ex: + assert sorted(set(ex._tables().matrix['row'].to_list())) == [1, 2], 'row 0 is the orphan under test' + solution = ex.solve(batch_rows=batch_rows, solver_name=solver_name) + assert solution.termination_condition == 'infeasible' + + def test_row_chunks_are_bounded_by_nonzeros_not_by_rows(): """A chunk of rows is a chunk of *entries*, and only entries are residency.