perf(persistent): store tz identity in coordinate snapshots

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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/pull/960>`__)
 
 **Bug fixes**
 

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)

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"

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