diff --git a/doc/release_notes.rst b/doc/release_notes.rst index dae3a6dd4..e6e1581a6 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -45,6 +45,7 @@ Upcoming Version * ``@``/``dot`` against a constant matrix that holds zeros no longer densifies the result to one term per contracted member. The zero-coefficient terms are dropped, so the term dimension shrinks to the widest non-zero cell. On PyPSA's Kirchhoff Voltage Law constraint (a cycle matrix with ~3 branches per cycle) this cuts the expression from 852 to 3 terms — 284x fewer cells — which in turn shrinks the downstream ``merge``. A constant without zeros is unaffected. (`#748 `__) * ``densify_terms`` (used by ``sum(drop_zeros=True)`` and the sparse ``@`` path) is now fully vectorised. It previously counted the non-zero positions with a Python loop that scaled quadratically in the number of non-zero terms — 127 s for a (2000 x 60) expression, now 3 ms — and allocated the compacted output at the full original term width. It now allocates only the compacted width and returns the expression unchanged when it holds no zeros. +* Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 `__) **Bug fixes** diff --git a/linopy/persistent/diff.py b/linopy/persistent/diff.py index 38173d12a..52058d433 100644 --- a/linopy/persistent/diff.py +++ b/linopy/persistent/diff.py @@ -63,12 +63,21 @@ def _same(a: np.ndarray, b: np.ndarray) -> bool: def _coords_equal( - a: dict[str, np.ndarray], b: dict[str, np.ndarray], ignored: frozenset[str] + a: dict[str, tuple[str | None, np.ndarray]], + b: dict[str, tuple[str | None, np.ndarray]], + ignored: frozenset[str], ) -> bool: keys = a.keys() - ignored if keys != b.keys() - ignored: return False - return all(np.array_equal(a[k], b[k]) for k in keys) + for k in keys: + tz_a, arr_a = a[k] + tz_b, arr_b = b[k] + # the tz key carries tz identity: a naive index never equals a + # tz-aware one, and differently-zoned indexes always reindex + if tz_a != tz_b or not np.array_equal(arr_a, arr_b): + return False + return True def _structural_reason(base: StructuralKey, model: Model) -> RebuildReason | None: @@ -359,8 +368,8 @@ def __init__( # Target-state material for the snapshot assembled in finalize(). self.var_buffers: dict[str, ContainerVarBuffers] = {} self.con_buffers: dict[str, ContainerConBuffers] = {} - self.var_coords: dict[str, dict[str, np.ndarray]] = {} - self.con_coords: dict[str, dict[str, np.ndarray]] = {} + self.var_coords: dict[str, dict[str, tuple[str | None, np.ndarray]]] = {} + self.con_coords: dict[str, dict[str, tuple[str | None, np.ndarray]]] = {} self._snap_obj_c: np.ndarray | None = None self._snap_obj_sense: str | None = None @@ -395,7 +404,7 @@ def diff_var( name: str, var: Variable, base_buf: ContainerVarBuffers, - base_coords: dict[str, np.ndarray], + base_coords: dict[str, tuple[str | None, np.ndarray]], ) -> RebuildReason | None: new_buf = _extract_var_buffers(var) new_coords = _coord_snapshot(var) @@ -442,7 +451,7 @@ def diff_con( name: str, con: ConstraintBase, base_buf: ContainerConBuffers, - base_coords: dict[str, np.ndarray], + base_coords: dict[str, tuple[str | None, np.ndarray]], skip_coef_compare: bool, ) -> RebuildReason | None: new_buf = _extract_con_buffers(con, self.var_label_index) diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index 187a4bb04..76396d5ad 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING import numpy as np +import pandas as pd from linopy import expressions from linopy.constraints import Constraint @@ -130,8 +131,28 @@ class ContainerConBuffers: active_labels: np.ndarray -def _coord_snapshot(obj: Variable | ConstraintBase) -> dict[str, np.ndarray]: - return {str(name): np.asarray(idx) for name, idx in obj.indexes.items()} +def _coord_snapshot( + obj: Variable | ConstraintBase, +) -> dict[str, tuple[str | None, np.ndarray]]: + """ + Snapshot a container's coordinate indexes. + + Each coordinate is stored as a ``(tz_key, values)`` pair: the raw values as + a numpy array plus the tz identity for equality purposes. For tz-aware + ``DatetimeIndex`` coordinates the values are stored as UTC-ns + ``datetime64`` — ``np.asarray`` would instead materialize an object array + of Timestamps (O(n) Python objects per container and per diff) — and the + tz string is carried so tz identity stays part of the equality: a naive + index never equals a tz-aware one, matching pandas' own naive/aware + comparison semantics. + """ + out: dict[str, tuple[str | None, np.ndarray]] = {} + for name, idx in obj.indexes.items(): + if isinstance(idx, pd.DatetimeIndex) and idx.tz is not None: + out[str(name)] = (str(idx.tz), idx.tz_convert(None).to_numpy()) + else: + out[str(name)] = (None, np.asarray(idx)) + return out def clear_coef_dirty(model: Model) -> None: @@ -152,8 +173,12 @@ class ModelSnapshot: structural_key: StructuralKey var_buffers: dict[str, ContainerVarBuffers] = field(default_factory=dict) con_buffers: dict[str, ContainerConBuffers] = field(default_factory=dict) - var_coords: dict[str, dict[str, np.ndarray]] = field(default_factory=dict) - con_coords: dict[str, dict[str, np.ndarray]] = field(default_factory=dict) + var_coords: dict[str, dict[str, tuple[str | None, np.ndarray]]] = field( + default_factory=dict + ) + con_coords: dict[str, dict[str, tuple[str | None, np.ndarray]]] = field( + default_factory=dict + ) obj_c: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=np.float64)) obj_quad_present: bool = False obj_sense: str = "min" diff --git a/test/test_persistent_snapshot_diff.py b/test/test_persistent_snapshot_diff.py index 8730bf340..f9135cb76 100644 --- a/test/test_persistent_snapshot_diff.py +++ b/test/test_persistent_snapshot_diff.py @@ -208,6 +208,75 @@ def test_ignore_dims_detects_coord_change() -> None: assert isinstance(ModelDiff.from_snapshot(snap, m2, ignore_dims={"t"}), ModelDiff) +def _model_with_time_coords(index: pd.Index) -> Model: + m = Model() + m.add_variables(0, 10, coords=[index], name="x") + m.add_constraints(m.variables["x"] >= 0, name="c1") + m.add_objective(m.variables["x"].sum()) + return m + + +def test_tz_aware_coords_roundtrip_no_rebuild() -> None: + """ + A tz-aware index snapshots to UTC-ns arrays; identical coordinates must + diff clean (the snapshot conversion must not change equality semantics). + """ + idx = pd.date_range("2025-01-01", periods=3, freq="h", tz="Etc/GMT-10", name="t") + snap = ModelSnapshot.capture(_model_with_time_coords(idx)) + assert isinstance( + ModelDiff.from_snapshot(snap, _model_with_time_coords(idx)), ModelDiff + ) + + +def test_tz_aware_coords_survive_dst_instants() -> None: + """ + Coordinates around a DST transition compare by instant, and equal + instants across the boundary stay equal (no re-localization happens). + """ + idx = pd.date_range( + "2025-03-30 00:00", periods=6, freq="30min", tz="Australia/Sydney", name="t" + ) + snap = ModelSnapshot.capture(_model_with_time_coords(idx)) + assert isinstance( + ModelDiff.from_snapshot(snap, _model_with_time_coords(idx)), ModelDiff + ) + + +def test_naive_coords_never_equal_tz_aware_coords() -> None: + """ + A naive index and a tz-aware index with identical wall labels must + rebuild: pandas never equates naive and aware timestamps, and the + snapshot's tz key preserves that. + """ + naive = pd.date_range("2025-01-01", periods=3, freq="h", name="t") + aware = pd.date_range("2025-01-01", periods=3, freq="h", tz="UTC", name="t") + snap = ModelSnapshot.capture(_model_with_time_coords(aware)) + assert ( + ModelDiff.from_snapshot(snap, _model_with_time_coords(naive)) + is RebuildReason.COORD_REINDEX + ) + + +def test_different_tz_same_instants_reindex() -> None: + """ + Differently-zoned indexes never diff clean, even with equal instants: + a tz identity change is a coordinate change and triggers a rebuild (the + conservative direction - a rebuild is always safe, an in-place update + against re-labeled coordinates would not be). + """ + utc = pd.date_range("2025-01-01", periods=3, freq="h", tz="UTC", name="t") + gmt10 = pd.date_range( + "2025-01-01 10:00", periods=3, freq="h", tz="Etc/GMT-10", name="t" + ) # Etc/GMT-10 is UTC+10: wall 10:00 == 00:00 UTC + assert (gmt10.tz_convert(None) == utc.tz_convert(None)).all() # same instants + + snap = ModelSnapshot.capture(_model_with_time_coords(utc)) + assert ( + ModelDiff.from_snapshot(snap, _model_with_time_coords(gmt10)) + is RebuildReason.COORD_REINDEX + ) + + def _assert_snapshot_equal(a: ModelSnapshot, b: ModelSnapshot) -> None: assert a.structural_key == b.structural_key assert a.var_buffers.keys() == b.var_buffers.keys() @@ -230,7 +299,10 @@ def _assert_snapshot_equal(a: ModelSnapshot, b: ModelSnapshot) -> None: for name in coords_a: assert coords_a[name].keys() == coords_b[name].keys() for dim in coords_a[name]: - np.testing.assert_array_equal(coords_a[name][dim], coords_b[name][dim]) + tz_a, arr_a = coords_a[name][dim] + tz_b, arr_b = coords_b[name][dim] + assert tz_a == tz_b + np.testing.assert_array_equal(arr_a, arr_b) np.testing.assert_array_equal(a.obj_c, b.obj_c) assert a.obj_quad_present == b.obj_quad_present assert a.obj_sense == b.obj_sense