perf: default integer arrays to int32 for ~25% memory reduction#566
perf: default integer arrays to int32 for ~25% memory reduction#566FBumann wants to merge 9 commits intoPyPSA:masterfrom
Conversation
linopy/constants.py — Added DEFAULT_LABEL_DTYPE = np.int32 linopy/model.py — Variable and constraint label assignment now uses np.arange(..., dtype=DEFAULT_LABEL_DTYPE) with overflow guards that raise ValueError if labels exceed int32 max. linopy/expressions.py — _term coord assignment and all .astype(int) for vars arrays now use DEFAULT_LABEL_DTYPE (int32). linopy/common.py — fill_missing_coords uses np.arange(..., dtype=DEFAULT_LABEL_DTYPE). Polars schema inference now checks array.dtype.itemsize instead of the old OS/numpy-version hack. test/test_constraints.py — Updated 2 dtype assertions to use np.issubdtype instead of == int. test/test_dtypes.py (new) — 7 tests covering int32 labels, expression vars, solve correctness, and overflow guards.
…k to int64 via astype(int), now use DEFAULT_LABEL_DTYPE. Also Variables.to_dataframe arange for map_labels. - linopy/constraints.py: Constraints.to_dataframe arange for map_labels. - linopy/common.py: save_join outer-join fallback was casting to int64.
…ords. Here's what changed: - test_linear_expression_sum / test_linear_expression_sum_with_const: v.loc[:9].add(v.loc[10:], join="override") → v.loc[:9] + v.loc[10:].assign_coords(dim_2=v.loc[:9].coords["dim_2"]) - test_add_join_override → test_add_positional_assign_coords: uses v + disjoint.assign_coords(...) - test_add_constant_join_override → test_add_constant_positional: now uses different coords [5,6,7] + assign_coords to make the test meaningful - test_same_shape_add_join_override → test_same_shape_add_assign_coords: uses + c.to_linexpr().assign_coords(...) - test_add_constant_override_positional → test_add_constant_positional_different_coords: expr + other.assign_coords(...) - test_sub_constant_override → test_sub_constant_positional: expr - other.assign_coords(...) - test_mul_constant_override_positional → test_mul_constant_positional: expr * other.assign_coords(...) - test_div_constant_override_positional → test_div_constant_positional: expr / other.assign_coords(...) - test_variable_mul_override → test_variable_mul_positional: a * other.assign_coords(...) - test_variable_div_override → test_variable_div_positional: a / other.assign_coords(...) - test_add_same_coords_all_joins: removed "override" from loop, added assign_coords variant - test_add_scalar_with_explicit_join → test_add_scalar: simplified to expr + 10
FBumann
left a comment
There was a problem hiding this comment.
Review
The guards and logical checks all look correct (see discussion in prior review). Two things to add before merge:
1. Release notes
Please add to doc/release_notes.rst under "Upcoming Version", something like:
* Default internal integer arrays (labels, variable indices, ``_term`` coordinates) to ``int32`` instead of ``int64``, reducing memory usage by ~25% and improving model build speed by 10-35%. The dtype is controlled by ``linopy.constants.DEFAULT_LABEL_DTYPE`` and can be changed back to ``np.int64`` before model construction if needed. An overflow guard raises ``ValueError`` if labels exceed the int32 maximum (~2.1 billion).2. Document how to override DEFAULT_LABEL_DTYPE
Since every module imports DEFAULT_LABEL_DTYPE by name at import time, simply assigning linopy.constants.DEFAULT_LABEL_DTYPE = np.int64 after import won't propagate. The constant should either:
(a) Be documented as a compile-time constant (users edit constants.py or monkey-patch before importing linopy), or
(b) Be read indirectly so runtime changes work. For example, change all usages to read from the module rather than a local binding:
# In constants.py — no change needed
DEFAULT_LABEL_DTYPE = np.int32
# In model.py, expressions.py, etc. — instead of:
from linopy.constants import DEFAULT_LABEL_DTYPE
# Use:
from linopy import constants
# ... then reference constants.DEFAULT_LABEL_DTYPE everywhereThis way linopy.constants.DEFAULT_LABEL_DTYPE = np.int64 at runtime would work. Option (b) is a small change and much more user-friendly. Up to you whether this is worth doing now or in a follow-up.
Minor items from prior review (still applicable)
- Add
-> Nonereturn type annotations to all test functions intest_dtypes.py(CI blocker) - Guard
test_solve_with_int32_labelswithpytest.importorskip("highspy") - Remove trailing spaces in the overflow error message strings in
model.py
|
Note on scipy compatibility: scipy sparse matrices ( |
- Move DEFAULT_LABEL_DTYPE from constants.py into options["label_dtype"] - Widen OptionSettings types from int to Any - Add validation: label_dtype only accepts np.int32 or np.int64 - Fix matrices.py empty clabels fallback to use configured dtype - Fix f-string quoting and trailing spaces in overflow error messages - Add -> None annotations and importorskip guard in test_dtypes.py - Add tests for int64 override and invalid dtype rejection - Add release notes entry Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Dimension coordinates (fill_missing_coords, _term coord) are small index arrays, not the large label/vars arrays that benefit from int32. xarray's index creation is slower with int32 than the default int64, causing a 13-38% build regression. Revert these to default int while keeping int32 for labels and vars where the memory savings matter. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Benchmark Results:
|
Changes proposed in this Pull Request
Cut memory for internal integer arrays (labels, vars indices,
_termcoords) by ~25% and improve build speed by ~10-35% by defaulting toint32instead ofint64.What changed
linopy/constants.py: AddedDEFAULT_LABEL_DTYPE = np.int32linopy/model.py: Variable and constraint label assignment usesnp.arange(..., dtype=DEFAULT_LABEL_DTYPE)with overflow guard that raisesValueErrorif labels exceed int32 max (~2.1 billion)linopy/expressions.py:_termcoord assignment and.astype(int)for vars arrays now useDEFAULT_LABEL_DTYPElinopy/variables.py:ffill,bfill,sanitizeuseDEFAULT_LABEL_DTYPEinstead ofastype(int)(which widened labels back to int64);Variables.to_dataframearange uses int32linopy/constraints.py:Constraints.to_dataframearange usesDEFAULT_LABEL_DTYPElinopy/common.py:fill_missing_coordsuses int32 arange;save_joinouter-join fallback usesDEFAULT_LABEL_DTYPEinstead ofastype(int); polars schema infersInt32/Int64based on actual array dtypetest/test_constraints.py: Updated dtype assertions to usenp.issubdtype(compatible with both int32 and int64)test/test_dtypes.py(new): Tests for int32 labels, expression vars, solve correctness, and overflow guarddev-scripts/benchmark_lp_writer.py(new): Benchmark script supporting--phase memory|build|lp_writewith--plotcomparison modeBenchmark results
Reproduce with:
Memory (dataset
.nbytes)Consistent 1.25x reduction across all problem sizes (e.g. 640 MB → 512 MB at 8M vars). The
labelsandvarsarrays shrink 50% (int64 → int32) whilelower/upper/coeffs/rhsstay float64.Build speed
Consistently ~1.1-1.35x faster across all sizes (30 iterations with GC, tight error bars). 10-20% for large models (170ms → 153ms at 8M vars), and up to 35% for small/medium models where the fixed overhead of array allocation matters more relative to total time.
Similar results on real PyPSA model.
No influence on lp-write
Checklist
doc.doc/release_notes.rstof the upcoming release is included.