perf: speed up LP file writing (2.5-3.9x on large models, no regressions on small)#564
Conversation
Extract _format_and_write() helper that uses lazy().collect(engine="streaming") with automatic fallback, replacing 7 instances of df.select(concat_str(...)).write_csv(...).
Replace the vertical concat + sort approach in Constraint.to_polars() with an inner join, so every row has all columns populated. This removes the need for the group_by validation step in constraints_to_file() and simplifies the formatting expressions by eliminating null checks on coeffs/vars columns.
…r short DataFrame - Skip group_terms_polars when _term dim size is 1 (no duplicate vars) - Build the short DataFrame (labels, rhs, sign) directly with numpy instead of going through xarray.broadcast + to_polars - Add sign column via pl.lit when uniform (common case), avoiding costly numpy string array → polars conversion Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…e vars Check n_unique before running the expensive group_by+sum. When all variable references are unique (common case for objectives), this saves ~31ms per 320k terms. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Replace np.unique with faster numpy equality check for sign uniformity. Eliminate redundant filter_nulls_polars and check_has_nulls_polars on the short DataFrame by applying the labels mask directly during construction. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Guard against IndexError when sign_flat is empty (no valid labels) by checking len(sign_flat) > 0 before accessing sign_flat[0]. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…for duplicate (labels, vars) pairs before calling group_terms_polars. Use it in both Constraint.to_polars() and LinearExpression.to_polars() to avoid expensive group_by when terms already reference distinct variables
|
Wonderful @FBumann ! This is very much welcome! |
|
@FabianHofmann Should a fix the codecov stuff? |
FabianHofmann
left a comment
There was a problem hiding this comment.
nice one @FBumann ! happy to pull this in! should we remove the dev-script? thanks for the transparency on the benchmarks
dev-scripts/benchmark_lp_writer.py
Outdated
| @@ -0,0 +1,388 @@ | |||
| #!/usr/bin/env python3 | |||
There was a problem hiding this comment.
just came to my mind, that a uv runnable script for dev-scripts would actually be super nice. not needed at all, just for inspiration
There was a problem hiding this comment.
Ill keep that in mind next time. Im also using uv. If linopy is too, ill do it the next time ;)
linopy/io.py
Outdated
| kwargs: Any = dict( | ||
| separator=" ", null_value="", quote_style="never", include_header=False | ||
| ) | ||
| try: |
There was a problem hiding this comment.
do you think this is needed? is the lazy operation still a bit flaky? happy to keep it in if you think it is safer
There was a problem hiding this comment.
Im not sure if you meant the kwargs themselves: I moved them into the method call for readability
There was a problem hiding this comment.
If you meant the try Except:
The streaming engine is stable and recomended (see https://pola.rs/posts/polars-in-aggregate-dec25/)
If we raise the lb of the polars version, we can remove the fallback i think (>=1.31.0, where the old one was removed)
There was a problem hiding this comment.
Sorry, the selection was bad. I meant the whole fallback mechanism
There was a problem hiding this comment.
Further: it's officially recommended, has transparent fallback built-in (no exceptions), and is on track to become the default
There was a problem hiding this comment.
I pushed a commit that did that. I can revert if you want to
There was a problem hiding this comment.
Wonderful! Let's pull this in then
There was a problem hiding this comment.
No need to revert, but as you prefer
There was a problem hiding this comment.
Im sure this is stable with the polars >=1.31 !
Lets merge this
I removed it. I think it would be great to keep such benchmarks, but ill discuss it in #567 instead |
Changes proposed in this Pull Request
Speed up LP file writing by up to 3.9x on large models, with consistent improvements across all problem sizes. Includes a benchmark script for reproducibility.
Added a lp-file benchmarking script, which is a majority of the lines changed
Performance optimizations
concat_str+write_csvvia new_format_and_write()helper (with automatic fallback + warning)concat+sortwithjoinfor constraint assemblymaybe_group_terms_polars()to skip expensivegroup_bywhen terms already reference distinct variablesBug fixes
IndexErroron empty constraint slices insign_flatcheckBenchmark results
Reproduce with
python dev-scripts/benchmark_lp_writer.py --model basic -o results.json --label "my run".Synthetic model (2×N² vars, 2×N² constraints)
No regressions on small models, speedup grows with problem size up to 3.9x at 8M variables.
PyPSA SciGrid-DE (realistic power system model, 24–1000 snapshots)
Consistent 2.5–2.7x speedup across all sizes, reaching 7.0s → 2.7s at 2.5M variables / 6M constraints.
Checklist
doc.doc/release_notes.rstof the upcoming release is included.