Feature: Numba/JIT-compiled constraint matrix assembly for 3-4x faster model building

[MaykThewessen]

Describe the feature you'd like to see

JIT-compile the sparse matrix assembly in matrices.py using Numba to reduce model build time by 3-4x.

Profiling in #271 (PyOptInterface benchmarks) showed that matrix construction — converting labeled xarray variables/constraints into sparse COO matrices for solver handoff — accounts for 55% of total solve time (8s out of 14.5s). The solver itself is often faster than the formulation layer.

For production workloads (120+ buses, 8760h, scenario sweeps), the model building step is the dominant bottleneck. This is especially painful for iterative workflows like MGA, rolling horizon, and sensitivity analysis where the model is rebuilt many times.

Motivation

Model Size	Current Build Time	Estimated with Numba	Improvement
50 buses, 8760h	~5s	~1.5s	3x
120 buses, 8760h	~15s	~4s	3-4x
500 buses (PyPSA-Eur scale), 8760h	~60s+	~15s	4x

Estimates are based on typical Numba JIT speedups for sparse matrix assembly in similar scientific computing workloads (scipy sparse construction, finite element assembly).

Why this matters more now

• Linopy's solver interface is already fast (persistent solver support in Support iterative solve of the same solver_model (when using highs) #198, multiple solver backends)
• GPU solver support (cuPDLPx in v0.6.0) makes the pre-solve bottleneck even more pronounced
• Iterative workflows (MGA via PyPSA, rolling horizon) multiply the cost of slow model building
• The profiling data from Explore PytOptInferface for writing matrix over native C API #271 confirms this is the right target

Proposed implementation

Target: MatrixAccessor in matrices.py

The core hot path is the loop that constructs COO triplets (row, col, data) from labeled xarray DataArrays. This is pure numerical work operating on integer indices and float coefficients — ideal for Numba.

Steps

1. Profile matrices.py to isolate the specific functions where time is spent (COO triplet construction, index alignment, coefficient extraction)
2. Extract the inner numerical loops into standalone functions with pure NumPy array signatures (no xarray/pandas inside the hot path)
3. Apply @numba.njit(cache=True) to the extracted functions — this compiles them to machine code with zero Python interpreter overhead
4. Parallel variant using @numba.njit(parallel=True) with numba.prange for the outer loop over constraints (each constraint's matrix contribution is independent)
5. Fallback path for when Numba is not installed — keep current pure-Python implementation as default, use Numba when available
6. Benchmark on standard test cases (PyPSA-Eur 37-bus, 128-bus, 256-bus electricity models)

Example pattern

import numba

@numba.njit(cache=True)
def _build_coo_triplets(var_labels, coeff_data, row_offsets, col_offsets):
    """Assemble COO sparse matrix entries from constraint coefficients.
    
    Pure numerical loop — no Python objects, no xarray, no pandas.
    """
    n_entries = len(var_labels)
    rows = np.empty(n_entries, dtype=np.int64)
    cols = np.empty(n_entries, dtype=np.int64)
    data = np.empty(n_entries, dtype=np.float64)
    
    for i in range(n_entries):
        rows[i] = row_offsets[var_labels[i]]
        cols[i] = col_offsets[var_labels[i]]
        data[i] = coeff_data[i]
    
    return rows, cols, data

Dependency consideration

Numba is already standard in the scientific Python ecosystem (depends only on NumPy + LLVM). It could be an optional dependency:

try:
    from linopy._numba_matrices import _build_coo_triplets
except ImportError:
    from linopy._matrices_fallback import _build_coo_triplets

Related issues

• Explore PytOptInferface for writing matrix over native C API #271 — PyOptInterface benchmarks showing matrix construction as the bottleneck (55% of total time)
• Support for sparse xarrays #248 — Sparse xarray support (complementary: reduces memory, this reduces time)
• The vision and the future of Linopy #207 — Linopy vision document (mentions time profiling as a priority)
• Optionally release memory during solving process #219 — Memory release during solving (complementary)

[coroa]

Well, neither the lp-polars, nor the the new pyoptinterface (#597 ,implementation draft), do use the matrix accessors, so this would not accelerate anything there, but @CharlieFModo reports systematically faster build times with the direct api and that depends on the matrix accessors (by a new xpress direct api and for some solvers the direct api is faster for small instances, highs for example, #596 ).

[MaykThewessen]

Oh great, didn't know that, thanks for the context @coroa

So if I understand correctly, the matrix accessors are mainly relevant for the direct API path now, and both lp-polars and the upcoming pyoptinterface IO bypass them entirely. That makes the Numba proposal less impactful than I estimated.

A few follow-up questions:

• Which IO path do you recommend for HiGHS with ~120 buses, 8760h models? I'm currently using the default (lp-polars). Is pyoptinterface (feat: add pyoptinterface io_api #597) expected to be faster for this scale?
• Would it be more useful if I helped test feat: add pyoptinterface io_api #597 or feat(xpress): add direct io #596 instead? Happy to benchmark against my production workloads if that helps move those PRs forward.

If the direction is to move away from matrix accessors entirely, that is even better than my proposal

[MaykThewessen]

or my HiGHS-based setup (~200 buses, 8760h chunked), the solver itself dominates at ~85% of wall time, not the model building step.

Since I'm already on lp-polars, these IO improvements would be marginal for my use case.

My real bottlenecks are RAM pressure during solving and HiGHS solve time.

Are there linopy-side improvements planned for memory efficiency (related to #248 sparse xarray)?

Or is the recommendation to look at solver-level optimizations (bus aggregation, model reduction) instead?

• I have already started to group fossil plants together and merge more nodes

[FBumann]

@MaykThewessen
Regarding your questions:

1. feat: add pyoptinterface io_api #597 is not fully implemented, so well see on that as soon as we are ready.
2. This would be appreciated of course! Benchmarks of real use cases always help.

Regarding Memory usage:
Im working on #566 for reduced memory usage.
There is nothing else planned addressing memory efficiency as far as i know. Sparsity is not trivial to handle, as xarray (the foundation of linopy), doesnt handle it explicitly.

If you are willing to investigate how to improve memory usage for sparse models, this would be highly appreciated.
Also, i think adding a notebook explaining best practices regarding memory in our docs would also be a good addition. If you would like to start sth like that.

[MaykThewessen]

@FBumann — happy to do both. I've just done a systematic investigation from our production codebase (200-bus PyPSA network, ~1800 generators, 8760 h/year solved in 168 h chunks) and drafted a best-practices notebook. Here's what I found and what I'd propose to contribute.

---

Memory investigation findings

For a 168-snapshot chunk at our model scale, the steady-state memory is ~9 GB and the peak during matrices.A assembly is ~15 GB (3× spike due to COO→CSC conversion). The three biggest holders are:

1. xarray.Dataset per variable/constraint — label arrays + bounds + term data. At (168 × 1.4M), float64 label+bounds storage alone is ~3.4 GB for variables.
2. Sparse matrix assembly spike — the COO intermediate before CSC conversion causes a 3× peak. This can't be avoided structurally, but it matters for sizing RAM.
3. Solution + dual vectors — dense numpy arrays, unavoidable but relatively small (~1.6 GB combined).

The single highest-impact mitigation we found was explicitly calling del network.model after each chunk followed by malloc_trim(0) on Linux (or malloc_zone_pressure_relief on macOS). Without the OS-level trim, Python returns freed pages to its allocator pool but RSS stays elevated — the next chunk then allocates on top of the previous peak rather than reusing freed memory.

matrices.clean_cached_properties() before the delete frees an additional 2–3 GB that would otherwise be held by cached sparse matrix accessors.

---

Best-practices notebook draft

I've drafted a notebook covering:

1. Linopy memory layout — formula for estimating model size (variables, constraints, sparse A, assembly peak)
2. Chunked time-horizon solving — the basic pattern, why del network.model + OS trim matters
3. LP blueprint reuse — for models where the LP structure is identical across chunks, skip linopy rebuild entirely for chunks 2+: extract HiGHS object from chunk 1, call changeRowsBounds() / changeColsCost() on it for subsequent chunks. Saves ~50 s and ~9 GB per chunk at our scale.
4. clearSolver() footgun — when reusing the HiGHS object, clearSolver() is essential before re-solve so presolve runs on the new RHS. Without it, HiGHS warm-starts from a stale basis → wrong solutions without warning.
5. Efficient RHS/objective updates — numpy broadcast fill vs pandas fillna (avoids allocating a full (T×N) intermediate), sparse push via keys >= 0 masking
6. Memory profiling — tracemalloc, psutil.rss_gb() patterns for identifying peaks
7. Sparse-friendly formulation — filter_missings=True, groupby vs dense broadcasting

The notebook is based on real production code patterns and includes working code examples. I'm happy to adapt it to linopy's docs style (Sphinx/MyST?) and open a PR.

---

On sparse xarray

Agree that native sparse xarray is hard — xarray delegates to numpy for most operations and sparse arrays have very uneven support. The most practical near-term path I see is what's already being done in #566: reduce the number of arrays rather than making individual arrays sparse. One concrete opportunity: for constraints where the RHS is static (doesn't vary by snapshot), storing it as a 1-D array instead of (T × M) would cut that component by a factor of T=168.

Happy to look at #566 in more detail if a second pair of eyes would help.