[FBumann]

Prompt: Let's review #687 — Can we make it more performant?

Stacked on #689, which is stacked on #688, which is stacked on #687 — review those first; this diff is only what it adds. Rebased onto 0429cc8; the numbers below were taken at e2c76e1, which differs from it only by a commit touching examples/README.md.

Note

The following content was generated by AI.

The build side is ~4× faster and the reformulation ~1.6×. On examples/sos.yaml at 20 000 snapshots × 20 generators × 5 breakpoints (2 000 000 members, 2 800 000 columns, 1 220 000 rows), against the stack's head:

	e2c76e1	this branch
_build_sos — placing every set	66–75 ms	16–19 ms
sos.reformulated(tables)	246–293 ms	167–178 ms
structure digest (a rebind's keep-or-reload question)	254–259 ms	233–251 ms

Ranges are the min of 15 reps in each of three interleaved rounds, on a shared 4-core box; medians move ±25% and are not worth quoting. A model that declares no set still pays nothing: 971 ms against 988 ms over four alternating rounds of the same build with the sos: block dropped, the sign flipping between them.

Where the time was

Placing the sets read every member's coordinate back into an ordinal. row_major(…, ordinal_of) costs a replace_strict against a per-label dictionary for every dim binding did not encode as an Enum — 30 ms for snapshot, 23 ms for bp, plus 6 ms building the 20 000-entry dict.

But a set and a weight are the two halves of a coordinate's position in the declared product, split at over: with c its cardinality and s the product of the dims after it, a position reads hcs + ws + l, the weight is w, and the set is hs + l — the position the same coordinate would have in the product without over. And the position is the label, wherever the variable dropped nothing, so an unmasked variable's sets are two divisions on int_range and its frame is never read. A masked one still reads its ordinals and still renumbers densely; that renumbering turned out to be load-bearing, for the reason _DTYPES already gives about col — a pre-mask position can pass 2^31 while every survivor fits, so a gapped set number is one an Int32 cannot hold. This is the split labels.frame already makes ("nothing renumbers unless a row was dropped", #520), in the one other place that numbers by position.

SosDeclaration.dims goes with it. It was the variable's foreach minus over — a second home for a fact, and the engine now needs where over sits in the full foreach anyway, so it asks Program.variable.

The reformulation built a 14-column × 2M frame to answer questions that are all scans. One of those columns, #cardinality row, was computed with a cum_sum and read by nothing at all. filter(holds) copied all 14 to use 2, n_unique() hashed a column already sorted, and the two refusals were two passes. sinks/sos.py is numpy end to end now — the same departure #687 already made for _linking_matrix, and the module docstring says so where the departure is.

The linking block also got simpler on the way. np.repeat(binary, entries) already names the right binary in every slot except a member's own and the segment closing into it, so three masked scatters became one broadcast, one scatter and one -= 1:

col = np.repeat(members.binary, members.entries)
coeff = np.repeat(members.magnitude, members.entries)
col[at] = members.col
coeff[at] = 1.0
col[(at + 1)[members.closes]] -= 1

set and weight were Int64 in the one frame with a row per member. Both are bounded by a count that already has to fit a 32-bit index — a set holds at least one column, a weight counts along one dim of one variable — so they are Int32 for the reason col is. 16 MB less carried, and the digest reads it.

Verified

Full suite green on the rebase (1393 passed, 6 skipped), ruff format --check, ruff check and pyrefly clean. The gurobi and LP-file arms of test_sos.py both run here, so the native lane is covered by the same change as the reformulating one — and #688's piecewise: {method: sos2} expands into a sos: block, so its cases go down this path too.

Two guards are new, because the placement now has two paths where it had one:

• a mask that drops nothing must place the sets where the arithmetic does — the same model down both paths has to produce the same frame, parametrized over both foreach orders, since the split is about where over sits in it;
• a mask that empties the first set must leave the numbering dense — the one arrangement where a hole and no hole are different numbers, and what pins the renumbering that the Int32 bound needs.

Not checked: nothing was measured on a quiet machine, and no solver other than HiGHS and Gurobi has read a file this writes — both inherited from #687.

Mutation table — every guard the change rests on, deleted in turn

Re-run after the rebase. uv run pytest -q -x tests/test_sos.py, falling back to the whole suite where that was green.

guard deleted
the sort when a set's members interleave	caught
the set keeps the dims after over (the stride)	caught
the set drops the over dim (the span)	caught
the weight is the over ordinal	caught
a weight is 1-based	caught
the positional path is only for an unmasked variable	caught¹
the counted path renumbers densely	caught¹
a SOS2 singleton is dropped	caught
nothing is dropped unless a set is a singleton	caught
a member's own column opens its row	caught
the member's own coefficient is 1	caught
the coefficient is broadcast over the whole span	caught
the span is as wide as the member's entries	caught
the closing segment is the binary before	caught
only a SOS2 member closes a segment	caught
a last member opens no binary	caught
a cardinality row owns its set's binaries	caught
a big-M caps the bound	caught
a member is non-negative	caught
a big-M is finite	caught

¹ survived the first pass and got a purpose-built probe before merge — the density one above. Both mutations leave a gapped set numbering, which every sink downstream tolerates (they read edges off !=, and _n_sets only offsets the next declaration), so nothing observable broke until the numbering itself was asserted.

Method

Base e2c76e1, both arms in separate processes off one venv (PYTHONPATH at a worktree of the base), alternating base/head, model built once per process and the two functions timed in isolation — a whole build is ~1 s on this box and the difference of two of those is noise. structure is re-read off a fresh _tables() each rep so the cached_property is not being timed.

The end-to-end marginal (build with the sos: block against build without) moved 115–146 ms → 4–28 ms, consistent with the isolated numbers plus the narrower stream, but it is a difference of two ~1 s builds and is quoted only as corroboration. That pair was taken against 94d8948, an earlier head of #687; the sos paths are byte-identical across every rebase this branch has been through, so it was not re-run.

Deliberately not done

• _add_sets in the gurobi sink is untouched, and it is the next thing that will hurt. It calls x.tolist(), which materialises a Python Var per model column rather than per member — the opposite of what its docstring claims — and then loops partition_by('set'), one frame per set. At 400 000 sets that lane is not usable. It is the native path rather than the one this PR made fast, and measuring it needs a licence past the wheel's size limit, so it is a finding and not a change.
• type and big_m still ride per member, and are still hashed per member by the digest, though both are per-declaration constants. Dropping them from structure would be reasoning from the declarations where that docstring deliberately reasons from the data, and it would trade a correctness margin for ~30 ms.
• _row_starts still allocates two temporaries it could write through. ~7 ms, against four lines where there are two.
• Departed from no Part 2 default. The one worth naming is the module-level import numpy in sinks/sos.py, where the neighbouring tables.py defers it into each function: this module is numpy in seven of its eight functions now, it already imports polars at module level, and it is reachable only from the engine-side sink registry.

[read-the-docs-community]

Documentation build overview

📚 lpspec | 🛠️ Build #34064657 | 📁 Comparing 195c561 against latest (8acf527)

  🔍 Preview build  

26 files changed · + 1 added · ± 25 modified

+ Added

• models/sos/index.html

± Modified

• ARCHITECTURE/index.html
• SPEC/index.html
• api/index.html
• benchmarks/index.html
• models/index.html
• design/ceiling/index.html
• models/data/index.html
• models/dispatch/index.html
• models/facility_location/index.html
• models/monthly_budget/index.html
• and 15 more...

[codspeed]

Merging this PR will not alter performance

✅ 14 untouched benchmarks
⏩ 38 skipped benchmarks¹

---

_{Comparing claude/pr-687-performance-n6rrs6 (195c561) with claude/diagnostics-what-the-sink-added (ff588cb)}

Footnotes

1. 38 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩