perf(persistent): store tz identity in coordinate snapshots

[olivier-lacroix]

This PR improves speed when using tz-aware indexes and persistent solve machinery

Note

The following content was generated by AI.

Changes proposed in this Pull Request

_coord_snapshot snapshotted coordinates with np.asarray, which on tz-aware
DatetimeIndex coordinates triggers pandas' object-array fallback: one Timestamp
Python object per coordinate per container — on every snapshot capture, and again on
every warm-start diff (the diff re-snapshots each container). On a 35,040-step,
8-container model that is ~280,000 Timestamps per solver build.

Snapshots are only consumed by _coords_equal (array equality), so each coordinate is
now stored as a (tz_key, UTC-ns datetime64 array) pair:

• the array converts vectorized (idx.tz_convert(None).to_numpy()), and
• the tz key keeps equality semantics exact — a naive index never equals a tz-aware
  one (matching pandas), and a tz identity change now triggers COORD_REINDEX: the
  conservative direction, since a rebuild is always safe while an in-place update
  against re-labeled coordinates would not be.

Coordinate-snapshot equality semantics

comparison	master	this branch
same tz, same instants	equal	equal
same tz, different instants (incl. across DST)	unequal	unequal
different tz, same instants	equal (instant-based)	reindex (conservative)
naive vs tz-aware	unequal	unequal (tz key)

DST needs no special handling: pandas stores tz-aware indexes as UTC internally and
tz_convert(None) drops the tz without re-localizing, so ambiguous/nonexistent wall
times — resolved at index construction — play no role at snapshot time.

Micro-benchmark

The script at the bottom of this description times the
snapshot path directly — no solver involved: 500,000 steps, 8 containers (4 variables

• 4 constraints sharing one DatetimeIndex), tz-aware vs naive index, interleaved
  A/B (2 rounds per side). The naive pass is the control: np.asarray was already
  vectorized there, so its timings should be flat across versions — and they are.

phase (mean of 2)	master	branch	delta
capture, tz-aware	1,242 ms	52 ms	~24x
diff, tz-aware	1,823 ms	55 ms	~33x
capture, naive (control)	45 ms	45 ms	flat
diff, naive (control)	52 ms	54 ms	flat

Repro script

#!/usr/bin/env python
"""
Repro: tz-aware coordinate snapshots materialize Timestamp object arrays.

Times the persistent snapshot path on a model whose 8 containers share an
N-step DatetimeIndex. No solver is involved: ModelSnapshot.capture and
ModelDiff.from_snapshot (which re-snapshots every container) are pure
Python-side, and each pays the coordinate snapshot once per container.

On master, _coord_snapshot calls np.asarray on each index; for tz-aware
DatetimeIndex coordinates pandas answers with an object array of Timestamps
(one Python object per coordinate per container). The branch stores a
(tz_key, UTC-ns datetime64 array) pair, converted vectorized.

The naive-index pass is the control: np.asarray was already vectorized
there, so its timings should be flat across versions.

Usage: python repro_tz_coord_snapshot.py [n_steps]
"""

import os
import sys
import time

import pandas as pd

from linopy import Model
from linopy.persistent.diff import ModelDiff
from linopy.persistent.snapshot import ModelSnapshot


def build_model(n: int, tz: str | None) -> Model:
    idx = pd.date_range("2025-01-01", periods=n, freq="s", tz=tz, name="t")
    m = Model()
    for k in range(4):
        m.add_variables(0, 1, coords=[idx], name=f"x{k}")
        m.add_constraints(m.variables[f"x{k}"] >= 0, name=f"c{k}")
    m.add_objective(m.variables["x0"].sum())
    return m


def timed(n: int, tz: str | None) -> None:
    label = "tz-aware (UTC)" if tz else "naive (control)"
    m = build_model(n, tz)
    t0 = time.perf_counter()
    snap = ModelSnapshot.capture(m)
    t1 = time.perf_counter()
    result = ModelDiff.from_snapshot(snap, m)
    t2 = time.perf_counter()
    print(f"{label}")
    print(f"  capture: {1e3 * (t1 - t0):9.1f} ms")
    print(f"  diff:    {1e3 * (t2 - t1):9.1f} ms   -> {type(result).__name__}")


def main(n: int = 500_000) -> None:
    import linopy

    print(f"linopy from: {os.path.realpath(linopy.__file__)}")
    print(f"n={n:,}   8 containers\n")
    timed(n, "UTC")
    timed(n, None)


if __name__ == "__main__":
    main(int(sys.argv[1]) if len(sys.argv) > 1 else 500_000)

Checklist

• AI-generated content is marked (see AGENTS.md).
• Code changes are sufficiently documented; i.e. new functions contain docstrings and further explanations may be given in doc.
• Unit tests for new features were added (if applicable).
• A note for the release notes doc/release_notes.rst of the upcoming release is included.
• I consent to the release of this PR's code under the MIT license.

[codspeed]

Merging this PR will degrade performance by 19.49%

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

❌ 4 regressed benchmarks
✅ 177 untouched benchmarks
⏩ 181 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
❌	test_to_lp[qp-n=1000]	1.9 MB	2.6 MB	-24.24%
❌	test_to_lp[nodal_balance_sparse-severity=50]	2.8 MB	3.7 MB	-24.11%
❌	test_to_lp[merge_balance-severity=0]	2.7 MB	3.3 MB	-18.26%
❌	test_to_lp[nodal_balance-severity=50]	3.3 MB	3.7 MB	-10.59%

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing olivier-lacroix:perf/tz-aware-coord-snapshot (2dfc3e6) with master (718c0c1)}

Footnotes

1. 181 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. ↩