From 663a15c856b51208193acea43cc93982ddeee788 Mon Sep 17 00:00:00 2001 From: Olivier Lacroix Date: Sun, 11 Oct 2026 23:10:02 +1100 Subject: [PATCH] perf(model): relabel assign_coords containers without a label dict Model.assign_coords rebuilt a master->new dict per container and mapped each container's labels through Index.map. On a tz-aware DatetimeIndex that map falls into pandas' object-to-datetime inference, costing milliseconds per container. Map full-length containers straight to the new values and locate subset containers' labels with get_indexer/take, and build a reassigned Dataset once instead of twice. ~817 ms -> ~3 ms for 30 containers on a 4380-label tz-aware axis. Co-Authored-By: Claude Code --- doc/release_notes.rst | 4 ++++ linopy/common.py | 10 ++++++---- linopy/model.py | 15 +++++++++++---- test/test_assign_coords.py | 28 ++++++++++++++++++++++++++++ 4 files changed, 49 insertions(+), 8 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 0b5b93b2..5829c1bc 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -102,6 +102,10 @@ Version 0.10.1 * ``to_netcdf``/``read_netcdf`` keep the auxiliary (non-index) coordinates of dense variables, constraints, expressions and ``Model.parameters``, such as the levels on a flat dimension that v1 recommends instead of a ``MultiIndex``, and scalar coordinates left by a scalar selection. (`#1024 `__) * ``licensed_solvers`` and ``check_solver_licenses()`` no longer report Xpress as licensed when its licence is rejected, e.g. a licence that only covers an older Xpress release. Xpress now probes its licence with ``xpress.init()``. (`#994 `__) +**Performance** + +* ``Model.assign_coords`` no longer relabels containers through a Python dict. It maps full-length containers straight to the new values and locates subset containers' labels with a positional ``get_indexer`` / ``take``: on a tz-aware ``DatetimeIndex`` — the rolling-horizon case it is built for — the previous ``Index.map`` over a dict fell into pandas' object-to-datetime inference and cost milliseconds per container. Reassigning a dataset's coordinates also builds it once instead of twice. A model with ~30 containers on a 4380-label tz-aware axis reassigns in ~3 ms, down from ~817 ms. + Version 0.10.0 -------------- diff --git a/linopy/common.py b/linopy/common.py index 3fbe8247..47b260d1 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -383,10 +383,12 @@ def reassign_coords(ds: Dataset, **coords: Any) -> Dataset: """ sizes = {str(name): coord.size for name, coord in ds.coords.items()} new_indexes = validate_coords_reassignment(sizes, coords, "dataset") - new = ds.assign_coords(new_indexes) - ordered = {name: new[name] for name in ds.coords} - data_vars = {name: new[name].variable for name in new.data_vars} - return Dataset(data_vars, coords=ordered, attrs=new.attrs) + # Build the dataset once, keeping each reassigned coordinate in its original position. + ordered: dict[Any, Any] = {name: ds.coords[name].variable for name in ds.coords} + for name, index in new_indexes.items(): + ordered[name] = index + data_vars = {name: var.variable for name, var in ds.data_vars.items()} + return Dataset(data_vars, coords=ordered, attrs=ds.attrs) T = TypeVar("T", Dataset, DataArray, "Variable", "LinearExpression", "Constraint") diff --git a/linopy/model.py b/linopy/model.py index 9c1b59a6..2d983155 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -2253,16 +2253,23 @@ def assign_coords(self, **coords_kwargs: Any) -> Model: ) for name, index in carriers: - if not index.isin(master).all(): + if index.equals(master): + # A full-length container carries the master's labels, so its + # new values are simply ``new``. + mapped.setdefault(dim, {})[name] = new + continue + # A subset container: locate each of its labels in the master and + # take the new value positionally. ``get_indexer`` is a C-level + # lookup, far cheaper than ``Index.map`` over a Python dict. + positions = master.get_indexer(index) + if (positions < 0).any(): raise ValueError( f"Cannot assign coordinates to dimension '{dim}': " f"container '{name}' carries labels outside the " "dimension's index. Relabel model-wide first, or " "align the container with `.sel`." ) - mapped.setdefault(dim, {})[name] = index.map( - dict(zip(master, new)) - ).rename(dim) + mapped.setdefault(dim, {})[name] = new.take(positions).rename(dim) for container in (self.variables, self.constraints, self.expressions): for name, item in container.items(): diff --git a/test/test_assign_coords.py b/test/test_assign_coords.py index 2ab5013f..a9bae91f 100644 --- a/test/test_assign_coords.py +++ b/test/test_assign_coords.py @@ -302,6 +302,34 @@ def test_assign_coords_maps_subset_carriers(m: Model) -> None: m2.assign_coords(gen=["h1", "h2", "h3"]) +def test_assign_coords_tz_aware_datetime_index() -> None: + """ + A tz-aware DatetimeIndex relabels across every container, full-length and + subset alike. + + The tz-aware path is the expensive one: mapping labels through a Python + dict falls into pandas' object-to-datetime inference, so the relabel locates + labels positionally instead. This guards that both the full-length + (``equals``) and subset (``get_indexer``) branches keep the right labels. + """ + idx = pd.date_range( + "2026-01-01", periods=4, freq="h", tz="Australia/Melbourne", name="time" + ) + new = idx + pd.Timedelta("1h") + model = Model() + x = model.add_variables(coords=[idx], name="x") + model.add_variables(coords=[idx[:2]], name="u") + model.add_constraints(x >= 0, name="c") + model.add_constraints(x.isel(time=slice(0, 2)) >= 1, name="c_subset") + + model.assign_coords(time=new) + + assert model.variables["x"].coords["time"].to_index().equals(new) + assert model.variables["u"].coords["time"].to_index().equals(new[:2]) + assert model.constraints["c"].coords["time"].to_index().equals(new) + assert model.constraints["c_subset"].coords["time"].to_index().equals(new[:2]) + + def test_assign_coords_rejects_unmatched_master_length(m: Model) -> None: """New values matching no carrier's length raise, listing the lengths.""" with pytest.raises(ValueError, match="no container carries it"):