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