diff --git a/doc/api.rst b/doc/api.rst index 0a56c6d3..0f22b881 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -644,6 +644,8 @@ Utilities align options ABSENT + from_long + Relation ``options`` holds these keys: diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 55f78057..88583246 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -58,6 +58,15 @@ Version 1.0.0 * ``linopy.testing.assert_conequal(strict=False)`` compares the lhs with ``assert_equivalent``, independently of its store and term layout, instead of comparing the dense datasets. (`#1047 `__) * ``Constraint.update`` raises a ``ValueError`` when the new ``lhs`` or ``rhs`` expression adds dimensions to the constraint; broadcasting a constraint to a larger grid is not supported. Before, the update was accepted and the labels stayed on the old dimensions, which left an inconsistent constraint. +**Long-form inputs** + +.. note:: + + The following content was generated by AI. + +* ``linopy.from_long(frame, value, coords, default=nan)`` pivots a long-form pandas or polars frame, key columns plus one value column, into a DataArray to pass wherever a constant is accepted, e.g. bounds, coefficients, right-hand sides and penalties. Each key column names a dimension of ``coords``; the result spans the key dimensions only, with their full coordinates, so it is the full cube over the key dimensions: key combinations not listed take ``default``, and linopy broadcasts the result over the dimensions without a key column. Missing value or key columns, key columns not naming a dimension, null or duplicate keys, keys outside ``coords``, a non-numeric value column and key dtypes (numeric, temporal, string) differing from the coordinates raise. (`#1049 `__) +* ``expr.groupby(linopy.Relation(frame, member, group)).sum()`` sums over a many-to-many relation, a pandas or polars frame of member and group columns whose rows may repeat a member: every row adds its member's terms to its group, e.g. links to the buses of both their ports, without an incidence matrix and ``@``. Members not listed drop out; null values and members outside the dimension raise. It stays sparse on a sparse expression; a quadratic expression raises ``NotImplementedError``. (`#1049 `__) + **Deprecations** * The ``active_fill`` parameter of ``add_piecewise_formulation`` is deprecated and emits a ``FutureWarning``; it is removed in 1.1. Pad a partial gate explicitly instead, e.g. ``active=status.reindex(gen=gens).fillna(1)``, which gives the same model. A partial ``active`` without padding still raises. diff --git a/linopy/__init__.py b/linopy/__init__.py index 2122b52d..229957a6 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -24,6 +24,7 @@ from linopy.constraints import Constraint, Constraints from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf +from linopy.long import Relation, from_long from linopy.model import Model, Variable, Variables from linopy.objective import Objective from linopy.piecewise import ( @@ -60,6 +61,7 @@ "QuadraticExpression", "RemoteHandler", "Slopes", + "Relation", "SolverFeature", "Variable", "Variables", @@ -67,6 +69,7 @@ "available_solvers", "licensed_solvers", "breakpoints", + "from_long", "merge", "options", "read_netcdf", diff --git a/linopy/expressions.py b/linopy/expressions.py index 483c6c59..d07e7cdc 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -106,6 +106,7 @@ TERM_DIM, ) from linopy.csr import CSRStore, Grid, _densify_notice +from linopy.long import Relation from linopy.semantics import ( AbsentType, FillValueLike, @@ -218,6 +219,13 @@ def _resolve_group(group: Any, data: Dataset) -> Any: like ``groupby(data["name"])``, mirroring xarray. Other inputs (Series, DataFrame, DataArray, multi-key lists) are returned unchanged. """ + if isinstance(group, pl.DataFrame) or ( + isinstance(group, pd.DataFrame) and group.index.name is None + ): + raise TypeError( + "Grouping by a frame needs a pandas index named after a dimension, " + "use `linopy.Relation(frame, member, group)` for member-group pairs." + ) if isinstance(group, (list, tuple)) and len(group) == 1: group = group[0] if isinstance(group, str) and group in data.coords: @@ -433,7 +441,7 @@ class LinearExpressionGroupby: """ obj: xr.Dataset | BaseExpression - group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame + group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame | Relation model: Any kwargs: Mapping[str, Any] = field(default_factory=dict) @@ -550,13 +558,26 @@ def sum( "`observed=True` is not supported with `use_fallback=True`." ) + if isinstance(self.group, Relation): + if isinstance(self.obj, QuadraticExpression): + raise NotImplementedError( + "Grouping a quadratic expression by a relation is not supported." + ) + if use_fallback or observed: + raise ValueError( + "`use_fallback` and `observed` are not supported when " + "grouping by a relation." + ) + return self._relation_sum(self._store, self.group) + + store = self._store + labels = self._labels group = _resolve_group(self.group, labels) _check_grouper_alignment(group, labels) multikey_frame = None if use_fallback else _multikey_value_frame(group, labels) - store = self._store store_cls = result_store_cls(store, self.model) if multikey_frame is not None and not observed: _warn_dense_grid(multikey_frame) @@ -649,6 +670,37 @@ def func(ds: Dataset) -> Dataset: ds = _restore_group_dim_position(ds, original_dims, grouped_dims) return LinearExpression(ds, self.model) + def _relation_sum(self, store: Store, relation: Relation) -> LinearExpression: + """ + Sum over a relation: every row adds its member's terms to its group, + so a member listed in several rows lands in several groups and one + not listed in none. The member cells are gathered once per row onto a + positional member dim, then scattered by + :meth:`~linopy.store.Store.from_grouper`. + """ + member, group = relation.member, relation.group + if group in store.dims: + raise ValueError( + f"relation group column {group!r} must not name an " + f"existing dimension {list(store.dims)}." + ) + if member not in store.dims: + raise ValueError( + f"relation member column {member!r} must name a dimension " + f"of {list(store.dims)}." + ) + positions = relation.positions(store.grid.indexes[member]) + rows = pd.RangeIndex(len(positions), name=member) + cells = store.grid.dataarray(np.arange(store.grid.size)) + cells = cells.isel({member: positions}).assign_coords({member: rows}) + gathered = store.take_positions(cells) + grouper = pd.Series(relation.groups(), index=rows, name=group) + store_cls = result_store_cls(store, self.model) + res = store_cls.from_grouper( + gathered, self.model, grouper, group, True, store.dims + ) + return LinearExpression._from_store(res, self.model) + def _grouped_sum(self, group: pd.Series, data: Dataset | None = None) -> Dataset: """ Sum groups by scattering all terms directly into the final padded arrays. @@ -1830,7 +1882,7 @@ def diff(self, dim: str, n: int = 1) -> Self: def groupby( self, - group: DataFrame | Series | DataArray, + group: DataFrame | Series | DataArray | Relation, restore_coord_dims: bool | None = None, **kwargs: Any, ) -> LinearExpressionGroupby: @@ -1842,9 +1894,12 @@ def groupby( Parameters ---------- - group : str, DataArray or IndexVariable + group : str, DataArray, IndexVariable, pandas.Series, pandas.DataFrame or linopy.Relation Array whose unique values should be used to group this array. If a string, must be the name of a variable contained in this dataset. + A :class:`~linopy.long.Relation` maps members to groups, possibly + several groups per member: ``sum`` adds each member to each of + its groups. restore_coord_dims : bool, optional If True, also restore the dimension order of multi-dimensional coordinates. diff --git a/linopy/long.py b/linopy/long.py new file mode 100644 index 00000000..08fdff00 --- /dev/null +++ b/linopy/long.py @@ -0,0 +1,284 @@ +""" +Long-form inputs. + +:func:`from_long` pivots a long-form frame of key columns plus one value +column into a DataArray, :class:`Relation` maps the members of a dimension +to groups for :meth:`~linopy.expressions.LinearExpressionGroupby.sum`. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeAlias + +import numpy as np +import pandas as pd +import polars as pl +from xarray import DataArray + +from linopy.alignment import _as_index, _coords_to_dict +from linopy.types import CoordsLike + +Frame: TypeAlias = pl.DataFrame | pd.DataFrame +Column: TypeAlias = pl.Series | pd.Series + +KIND_FAMILIES = { + "b": "boolean", + "i": "numeric", + "u": "numeric", + "f": "numeric", + "M": "datetime", + "m": "duration", + "O": "string", + "U": "string", + "S": "string", +} + + +def _family(column: Column | pd.Index) -> str: + """ + The dtype family of ``column``: numeric, datetime with its time zone, + duration, string or boolean. + """ + dtype = column.dtype + if isinstance(dtype, pd.DatetimeTZDtype): + return f"datetime[{dtype.tz}]" + if not isinstance(dtype, pl.DataType): + return KIND_FAMILIES[dtype.kind] + if isinstance(dtype, pl.Datetime) and dtype.time_zone is not None: + return f"datetime[{dtype.time_zone}]" + if isinstance(dtype, pl.Date | pl.Datetime): + return "datetime" + if isinstance(dtype, pl.Duration): + return "duration" + if dtype.is_numeric(): + return "numeric" + if dtype == pl.Boolean: + return "boolean" + if dtype == pl.String or isinstance(dtype, pl.Categorical | pl.Enum): + return "string" + raise TypeError(f"Column {column.name!r} has unsupported dtype {dtype}.") + + +def _as_polars(index: pd.Index) -> pl.Series: + if isinstance(index, pd.DatetimeIndex) and index.tz is not None: + utc = pl.Series(index.asi8).cast(pl.Datetime(index.unit)) + return utc.dt.replace_time_zone("UTC").dt.convert_time_zone(str(index.tz)) + return pl.Series(index.to_numpy()) + + +def _has_nulls(column: Column) -> bool: + if isinstance(column, pl.Series): + return column.null_count() > 0 + return bool(column.isna().any()) + + +def _is_float_numpy(column: Column) -> bool: + return isinstance(column.dtype, np.dtype) and column.dtype.kind == "f" + + +def _check_frame(frame: Frame) -> None: + if not isinstance(frame, Frame): + raise TypeError( + f"Expected a polars or pandas DataFrame, got {type(frame).__name__}." + ) + if isinstance(frame, pd.DataFrame) and ( + frame.columns.nlevels > 1 or frame.index.names != [None] + ): + raise TypeError( + "Expected flat columns and an unnamed index, call `reset_index()` " + "to turn index levels into key columns." + ) + + +def _positions(index: pd.Index, column: Column) -> pl.Series: + """ + The positions of the values of ``column`` in ``index``. Raises for null + values, a dtype family other than the one of ``index`` and values not in + ``index``. + """ + if _has_nulls(column): + raise ValueError(f"Column {column.name!r} has null values.") + if _family(column) != _family(index): + raise TypeError( + f"Column {column.name!r} of dtype {column.dtype} does not match " + f"the coordinates of dimension {index.name!r} of dtype {index.dtype}." + ) + if isinstance(column, pd.Series): + pos = pl.Series(index.get_indexer(column)) + pos = pos.set(pos == -1, None) + elif ( + isinstance(index, pd.RangeIndex) + and index.step == 1 + and column.dtype.is_integer() + ): + pos = column.cast(pl.Int64, strict=False) - index.start + pos = pos.set((pos < 0) | (pos >= len(index)), None) + else: + lookup = _as_polars(index) + key = column.cast(lookup.dtype, strict=False) + if lookup.dtype != column.dtype: + key = key.set(key.cast(column.dtype, strict=False).ne_missing(column), None) + table = pl.DataFrame({"key": lookup, "pos": np.arange(len(index))}) + pos = key.to_frame("key").join( + table, on="key", how="left", maintain_order="left" + )["pos"] + if pos.null_count(): + missing = pos.is_null() + outside = ( + column.filter(missing).unique(maintain_order=True).to_list() + if isinstance(column, pl.Series) + else column[missing.to_numpy()].unique().tolist() + ) + raise ValueError( + f"Values {outside[:5]} of column {column.name!r} are " + f"not in the coordinates of dimension {index.name!r}." + ) + return pos.cast(pl.Int64) + + +def from_long( + frame: Frame, value: str, coords: CoordsLike, default: float = np.nan +) -> DataArray: + """ + Pivot a long-form frame onto ``coords``. + + Every column of ``frame`` but ``value`` is a key column named after a + dimension of ``coords``. The result spans the key dimensions only, in + column order, with the full coordinates of ``coords``: key combinations + not listed take ``default``. Dimensions of ``coords`` without a key column + are not part of the result, linopy broadcasts it over them where it is + used. The result is the full cube over the key dimensions, a sparsely + filled frame costs the memory of its cube. + + Parameters + ---------- + frame : polars.DataFrame or pandas.DataFrame + The long-form data, one row per key combination. A pandas frame needs + an unnamed index, its keys are columns. + value : str + Name of the value column, numeric without null values; NaN values + are kept. + coords : list/dict/xarray.Coordinates + The coordinates the keys are positioned in, e.g. the coordinates of + the variable the result is used with. + default : float + Value of key combinations missing from ``frame``. Defaults to NaN. + + Returns + ------- + xarray.DataArray + Float values over the key dimensions. + + Raises + ------ + ValueError + For a missing value column, no key column, a key column not naming a + dimension of ``coords``, null keys or values, duplicate keys and keys + outside ``coords``. + TypeError + For a frame type other than a polars or pandas DataFrame, a non-numeric + value column and key dtypes (numeric, datetime with its time zone, + duration, string) differing from the ones of ``coords``. + + Examples + -------- + >>> import pandas as pd + >>> import linopy + >>> m = linopy.Model() + >>> i = pd.RangeIndex(3, name="i") + >>> ub = linopy.from_long( + ... pd.DataFrame({"i": [0, 2], "ub": [5.0, 7.0]}), "ub", coords=[i], default=0 + ... ) + >>> m.add_variables(upper=ub, coords=[i], name="x").upper.values + array([5., 0., 7.]) + """ + _check_frame(frame) + if value not in frame.columns: + raise ValueError(f"Frame has no value column {value!r}.") + keys = [k for k in frame.columns if k != value] + if not keys: + raise ValueError("Frame has no key column.") + target = {d: _as_index(v).rename(d) for d, v in _coords_to_dict(coords).items()} + unknown = [k for k in keys if k not in target] + if unknown: + raise ValueError( + f"Key columns {unknown} do not name a dimension of {list(target)}." + ) + if _family(frame[value]) != "numeric": + raise TypeError( + f"Value column {value!r} has non-numeric dtype {frame[value].dtype}." + ) + if _has_nulls(frame[value]) and not _is_float_numpy(frame[value]): + raise ValueError(f"Value column {value!r} has null values.") + flat = pl.Series(np.zeros(frame.shape[0], dtype=np.int64)) + for k in keys: + flat = flat * len(target[k]) + _positions(target[k], frame[k]) + if flat.n_unique() != len(flat): + raise ValueError(f"Frame has duplicate keys over {keys}.") + shape = [len(target[k]) for k in keys] + size = int(np.prod(shape)) + values = frame[value].to_numpy().astype(float) + if len(flat) == size and flat.is_sorted(): + out = values + else: + out = np.full(size, default, dtype=float) + out[flat.to_numpy()] = values + return DataArray(out.reshape(shape), coords=[target[k] for k in keys], dims=keys) + + +@dataclass(frozen=True, eq=False) +class Relation: + """ + A many-to-many map from the members of a dimension to groups. + + Each row of ``frame`` relates the member in column ``member``, a + coordinate of the dimension of the same name, to the group in column + ``group``; a member may be listed in several rows. + ``expr.groupby(relation).sum()`` adds each row's member to its group, so a + member listed in several rows lands in several groups and one not listed + in none. + + Parameters + ---------- + frame : polars.DataFrame or pandas.DataFrame + The member-group pairs. + member : str + Name of the member column, a dimension of the grouped expression. + group : str + Name of the group column, the dimension of the result. + + Examples + -------- + >>> import pandas as pd + >>> import linopy + >>> m = linopy.Model() + >>> link = pd.Index(["l0", "l1"], name="link") + >>> p = m.add_variables(coords=[link], name="p") + >>> ports = pd.DataFrame({"link": ["l0", "l0", "l1"], "bus": ["b0", "b1", "b1"]}) + >>> p.groupby(linopy.Relation(ports, "link", "bus")).sum().nterm + 2 + """ + + frame: Frame + member: str + group: str + + def __post_init__(self) -> None: + _check_frame(self.frame) + missing = {self.member, self.group} - set(self.frame.columns) + if missing or self.member == self.group: + raise ValueError( + f"Relation needs distinct columns {self.member!r} and " + f"{self.group!r}, missing {sorted(missing)}." + ) + if _has_nulls(self.frame[self.group]): + raise ValueError(f"Relation group column {self.group!r} has null values.") + + def positions(self, index: pd.Index) -> np.ndarray: + """The positions of the members in ``index``.""" + return _positions(index, self.frame[self.member]).to_numpy() + + def groups(self) -> np.ndarray: + """The group of each row.""" + return self.frame[self.group].to_numpy() diff --git a/linopy/store.py b/linopy/store.py index f91d977a..bcb676a9 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -276,7 +276,11 @@ def take(self, rows: np.ndarray, grid: Grid) -> Self: absent without terms. """ - def _take_positions(self, positions: DataArray) -> Self: + def take_positions(self, positions: DataArray) -> Self: + """ + Gather cell ``positions[...]`` (flat positions into this grid) into + each cell of a grid over the dims and coords of ``positions``. + """ from linopy.csr import Grid dims = tuple(map(str, positions.dims)) @@ -285,14 +289,14 @@ def _take_positions(self, positions: DataArray) -> Self: def isel(self, indexers: Mapping[Hashable, slice]) -> Self: """The cells at the positional slices ``indexers`` of grid dims.""" grid = self.grid - return self._take_positions(grid.dataarray(np.arange(grid.size)).isel(indexers)) + return self.take_positions(grid.dataarray(np.arange(grid.size)).isel(indexers)) def transpose(self, dims: Sequence[str]) -> Self: """The cells on the grid dims reordered to ``dims``.""" if tuple(dims) == self.dims: return self grid = self.grid - return self._take_positions( + return self.take_positions( grid.dataarray(np.arange(grid.size)).transpose(*dims) ) diff --git a/linopy/variables.py b/linopy/variables.py index 794de240..0e93ae28 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -61,6 +61,7 @@ STASHED_UPPER, ) from linopy.containers import LabelContainer +from linopy.long import Relation from linopy.scaling import ensure_scaling, validate_scaling from linopy.semantics import ( FillValueLike, @@ -764,7 +765,7 @@ def dot(self, other: ndarray | Variable) -> QuadraticExpression | LinearExpressi def groupby( self, - group: DataArray, + group: DataFrame | pd.Series | DataArray | Relation, restore_coord_dims: bool | None = None, ) -> LinearExpressionGroupby: """ @@ -775,9 +776,12 @@ def groupby( Parameters ---------- - group : str, DataArray or IndexVariable + group : str, DataArray, IndexVariable, pandas.Series, pandas.DataFrame or linopy.Relation Array whose unique values should be used to group this array. If a string, must be the name of a variable contained in this dataset. + A :class:`~linopy.long.Relation` maps members to groups, possibly + several groups per member: ``sum`` adds each member to each of + its groups. restore_coord_dims : bool, optional If True, also restore the dimension order of multi-dimensional coordinates. diff --git a/test/test_long.py b/test/test_long.py new file mode 100644 index 00000000..014006ab --- /dev/null +++ b/test/test_long.py @@ -0,0 +1,390 @@ +"""Long-form inputs: ``linopy.from_long`` constants and ``linopy.Relation`` group-by.""" + +from __future__ import annotations + +import datetime as dt +from collections.abc import Callable +from typing import Any + +import numpy as np +import pandas as pd +import polars as pl +import pytest +import xarray as xr +from test_store_conformance import RELATION, RELATION_FRAME, A, B, assert_same + +from linopy import ( + LinearExpression, + Model, + QuadraticExpression, + Relation, + Variable, + from_long, +) +from linopy.long import Column +from linopy.testing import assert_equivalent + + +def to_polars(frame: pd.DataFrame) -> pl.DataFrame: + return pl.from_dict(frame.astype(object).where(frame.notna(), None).to_dict("list")) + + +FRAME_LIBS: dict[str, Callable[[pd.DataFrame], Any]] = { + "pandas": lambda f: f, + "polars": to_polars, +} +FRAME_AB = pd.DataFrame( + {"a": [0, 2, 1, 0], "b": ["p", "p", "q", "q"], "v": [1.0, 2, 3, 4]} +) +FRAME_A = pd.DataFrame({"a": [2, 0, 1], "v": [1.0, 2, 3]}) + +CONSTANTS: dict[str, tuple[pd.DataFrame, float, xr.DataArray]] = { + "all-keys-default": ( + FRAME_AB, + 5.0, + xr.DataArray([[1.0, 4], [5, 3], [2, 5]], coords=[A, B]), + ), + "missing-key-column": (FRAME_A, np.nan, xr.DataArray([2.0, 3, 1], coords=[A])), + "column-order": ( + pd.DataFrame({"b": ["q", "p"], "a": [1, 1], "v": [3.0, 4]}), + 1.0, + xr.DataArray([[1.0, 4, 1], [1, 3, 1]], coords=[B, A]), + ), +} + + +@pytest.mark.parametrize("lib", list(FRAME_LIBS)) +@pytest.mark.parametrize("case", list(CONSTANTS)) +def test_from_long(case: str, lib: str) -> None: + frame, default, want = CONSTANTS[case] + res = from_long(FRAME_LIBS[lib](frame), "v", coords=[A, B], default=default) + xr.testing.assert_identical(res, want) + + +SNAPSHOTS = pd.date_range("2030-01-01", periods=3, freq="h", name="k") +KEY_COORDS: dict[str, pd.Index] = { + "int": pd.Index([10, 20, 30], name="k"), + "range": pd.RangeIndex(5, 8, name="k"), + "str": pd.Index(["x", "y", "z"], name="k"), + "datetime": SNAPSHOTS, +} + + +@pytest.mark.parametrize("lib", list(FRAME_LIBS)) +@pytest.mark.parametrize("dtype", list(KEY_COORDS)) +def test_from_long_key_dtypes(dtype: str, lib: str) -> None: + coord = KEY_COORDS[dtype] + frame = pd.DataFrame({"k": coord[[2, 0]], "v": [1.0, 2.0]}) + res = from_long(FRAME_LIBS[lib](frame), "v", coords=[coord], default=0) + xr.testing.assert_identical(res, xr.DataArray([2.0, 0, 1], coords=[coord])) + + +@pytest.mark.parametrize("fill", [1.0, 0.5], ids=["full", "partial"]) +def test_from_long_sorted_full_cube_matches_scatter(fill: float) -> None: + coords = [pd.RangeIndex(40, name="g"), pd.date_range("2030", periods=30, name="t")] + cube = pd.MultiIndex.from_product(coords).to_frame(index=False) + cube["v"] = np.arange(len(cube), dtype=float) + rows = cube.sample(frac=fill, random_state=0).sort_index() + frame = to_polars(rows) + shuffled = frame.sample(fraction=1.0, shuffle=True, seed=0) + res = from_long(frame, "v", coords=coords, default=-1) + xr.testing.assert_identical( + res, from_long(shuffled, "v", coords=coords, default=-1) + ) + want = cube.v.where(cube.index.isin(rows.index), -1.0).to_numpy() + np.testing.assert_array_equal(res.values, want.reshape(40, 30)) + + +def x(m: Model) -> Variable: + return m.add_variables(coords=[A, B], name="x") + + +def softened(m: Model, penalty: Any) -> LinearExpression | QuadraticExpression: + v = x(m) + m.add_objective(v.sum()) + m.add_constraints(v >= 0, penalty=penalty) + return m.objective.expression + + +POSITIONS: dict[str, Callable[[Model, Any], Any]] = { + "lower": lambda m, c: m.add_variables(lower=c, coords=[A, B], name="v").lower, + "upper": lambda m, c: m.add_variables(upper=c, coords=[A, B], name="v").upper, + "coefficient": lambda m, c: c * x(m), + "to_linexpr": lambda m, c: x(m).to_linexpr(c), + "add-constant": lambda m, c: 2 * x(m) + c, + "sub-constant": lambda m, c: 2 * x(m) - c, + "rsub-constant": lambda m, c: c - 2 * x(m), + "rhs-add_constraints": lambda m, c: m.add_constraints(2 * x(m), "<=", c), + "rhs-le": lambda m, c: 2 * x(m) <= c, + "reflected-ge": lambda m, c: c >= 2 * x(m), + "div": lambda m, c: x(m) / c, + "quad-mul": lambda m, c: c * x(m) ** 2, + "penalty": softened, +} + + +@pytest.mark.parametrize("case", ["all-keys-default", "missing-key-column"]) +@pytest.mark.parametrize("position", list(POSITIONS)) +def test_from_long_positions(backend: Model, position: str, case: str) -> None: + frame, default, want = CONSTANTS[case] + const = from_long(frame, "v", coords=[A, B], default=default) + if position == "penalty": + const = const.fillna(1.0) + want = want.fillna(1.0) + expected = POSITIONS[position](Model(sparse=backend.sparse), want) + assert_same(POSITIONS[position](backend, const), expected) + + +INVALID: dict[str, tuple[Any, str, type[Exception], str]] = { + "no-value-column": (pd.DataFrame({"a": [0]}), "v", ValueError, "no value column"), + "no-key-column": (pd.DataFrame({"v": [1.0]}), "v", ValueError, "no key column"), + "unknown-key": ( + pd.DataFrame({"a": [0], "z": [0], "v": [1.0]}), + "v", + ValueError, + r"\['z'\] do not name a dimension", + ), + "null-key": ( + pd.DataFrame({"b": ["p", None], "v": [1.0, 2]}), + "v", + ValueError, + "'b' has null values", + ), + "non-numeric-value": ( + pd.DataFrame({"a": [0], "v": ["1"]}), + "v", + TypeError, + "non-numeric", + ), + "duplicate-keys": ( + pd.DataFrame({"a": [0, 0], "v": [1.0, 2]}), + "v", + ValueError, + "duplicate keys", + ), + "key-outside-coords": ( + pd.DataFrame({"a": [7, 1, 9], "v": [1.0, 2, 3]}), + "v", + ValueError, + r"Values \[7, 9\] of column 'a' are not in", + ), + "null-value": ( + pd.DataFrame({"a": [0, 1], "v": pd.array([1, None], dtype="Int64")}), + "v", + ValueError, + "'v' has null values", + ), + "key-dtype-family": ( + pd.DataFrame({"a": ["0"], "v": [1.0]}), + "v", + TypeError, + "does not match", + ), +} + + +@pytest.mark.parametrize("lib", list(FRAME_LIBS)) +@pytest.mark.parametrize("case", list(INVALID)) +def test_from_long_invalid_raises(case: str, lib: str) -> None: + frame, value, error, match = INVALID[case] + with pytest.raises(error, match=match): + from_long(FRAME_LIBS[lib](frame), value, coords=[A, B]) + + +NOON = pd.DatetimeIndex(["2030-01-01 12:00", "2030-01-02 12:00"], name="k") +MIDNIGHT = pd.date_range("2030-01-01", periods=2, name="k") +SUB_MS = pd.DatetimeIndex(["2030-01-01 00:00:00.0004", "2030-01-01"], name="k") +BERLIN = pd.date_range("2030-01-01", periods=2, freq="h", name="k", tz="Europe/Berlin") +DAY = pl.Series([dt.date(2030, 1, 1)]) +MS = pl.Series([dt.datetime(2030, 1, 1)], dtype=pl.Datetime("ms")) +NOT_IN = (ValueError, "not in the coordinates") +NO_MATCH = (TypeError, "does not match") + +KEY_DTYPES: dict[str, tuple[Column, pd.Index, list[float] | tuple[type, str]]] = { + "int32": (pl.Series([300], dtype=pl.Int32), pd.Index([10, 300]), [np.nan, 1]), + "uint8-range": ( + pl.Series([2], dtype=pl.UInt8), + pd.RangeIndex(3), + [np.nan, np.nan, 1], + ), + "uint64-range-overflow": ( + pl.Series([2**63 + 5], dtype=pl.UInt64), + pd.RangeIndex(3), + NOT_IN, + ), + "float-range": (pl.Series([2.5]), pd.RangeIndex(3), NOT_IN), + "int-float-lossy": (pl.Series([10]), pd.Index([10.7, 20.0]), NOT_IN), + "int-float-lossy-pandas": (pd.Series([10]), pd.Index([10.7, 20.0]), NOT_IN), + "int-float": (pl.Series([1]), pd.Index([1.5, 1.0]), [np.nan, 1]), + "date-noon": (DAY, NOON, NOT_IN), + "date-midnight": (DAY, MIDNIGHT, [1, np.nan]), + "ms-sub-ms": (MS, SUB_MS, [np.nan, 1]), + "ms-us": (MS, MIDNIGHT, [1, np.nan]), + "tz": (MS.dt.replace_time_zone("Europe/Berlin"), BERLIN, [1, np.nan]), + "tz-pandas": (pd.Series(BERLIN[[1]]), BERLIN, [np.nan, 1]), + "tz-naive": (MS.dt.replace_time_zone("UTC"), MIDNIGHT, NO_MATCH), + "naive-tz-pandas": (pd.Series(MIDNIGHT[[0]]), BERLIN, NO_MATCH), + "duration": (pl.Series([1], dtype=pl.Duration("ns")), MIDNIGHT, NO_MATCH), + "categorical": (pl.Series(["q"], dtype=pl.Categorical), B, [np.nan, 1]), +} + + +@pytest.mark.parametrize("case", list(KEY_DTYPES)) +def test_from_long_column_dtypes(case: str) -> None: + column, index, want = KEY_DTYPES[case] + index = index.rename("k") + data = {"k": column.rename("k"), "v": [1.0] * len(column)} + frame = pl.DataFrame(data) if isinstance(column, pl.Series) else pd.DataFrame(data) + if isinstance(want, tuple): + with pytest.raises(want[0], match=want[1]): + from_long(frame, "v", coords=[index]) + else: + res = from_long(frame, "v", coords=[index]) + xr.testing.assert_identical(res, xr.DataArray(want, coords=[index])) + + +@pytest.mark.parametrize( + "frame", + [ + to_polars(FRAME_A).lazy(), + FRAME_A.set_index("a").v, + FRAME_A.set_index("a"), + FRAME_AB.set_index(["a", "b"]), + ], + ids=["lazyframe", "series", "named-index", "multiindex"], +) +def test_from_long_rejects_frame_forms(frame: Any) -> None: + with pytest.raises(TypeError): + from_long(frame, "v", coords=[A, B]) + + +def incidence(relation: pd.DataFrame, weights: Any = 1.0) -> xr.DataArray: + member, group = relation.columns + weighted = relation.assign(w=weights) + matrix = weighted.pivot_table("w", member, group, aggfunc="sum", fill_value=0.0) + return xr.DataArray(matrix) + + +@pytest.mark.parametrize("lib", list(FRAME_LIBS)) +@pytest.mark.parametrize("unlisted", [False, True], ids=["all", "unlisted"]) +def test_relation_groupby_matches_incidence_matmul( + backend: Model, unlisted: bool, lib: str +) -> None: + frame = RELATION_FRAME[RELATION_FRAME.a != 1] if unlisted else RELATION_FRAME + expr = 2 * backend.add_variables(coords=[B, A], name="x") + 1 + res = expr.groupby(Relation(FRAME_LIBS[lib](frame), "a", "g")).sum() + assert res.is_sparse == backend.sparse + assert_equivalent(res, expr @ incidence(frame).reindex(a=A, fill_value=0.0)) + + +@pytest.mark.parametrize("lib", list(FRAME_LIBS)) +def test_groupby_long_frame_raises(backend: Model, lib: str) -> None: + with pytest.raises(TypeError, match="linopy.Relation"): + (1 * x(backend)).groupby(FRAME_LIBS[lib](RELATION_FRAME)).sum() + + +def test_relation_groupby_quadratic_raises(backend: Model) -> None: + with pytest.raises(NotImplementedError, match="quadratic"): + (x(backend) ** 2).groupby(RELATION).sum() + + +@pytest.mark.parametrize( + ("relation", "kwargs", "error", "match"), + [ + ( + Relation(pd.DataFrame({"a": [0, 9], "g": ["u", "v"]}), "a", "g"), + {}, + ValueError, + r"Values \[9\] of column 'a' are not in", + ), + ( + Relation(pd.DataFrame({"a": [0, None], "g": ["u", "v"]}), "a", "g"), + {}, + ValueError, + "'a' has null values", + ), + ( + Relation(pd.DataFrame({"a": ["0"], "g": ["u"]}), "a", "g"), + {}, + TypeError, + "does not match", + ), + ( + Relation(pd.DataFrame({"a": [0, 1], "b": ["p", "q"]}), "a", "b"), + {}, + ValueError, + "must not name", + ), + ( + Relation(pd.DataFrame({"c": [0], "g": ["u"]}), "c", "g"), + {}, + ValueError, + "must name a dimension", + ), + (RELATION, {"use_fallback": True}, ValueError, "not supported"), + (RELATION, {"observed": True}, ValueError, "not supported"), + ], + ids=[ + "outside", + "null-member", + "member-dtype", + "group-is-dim", + "member-not-dim", + "fallback", + "observed", + ], +) +def test_invalid_relation_groupby_raises( + backend: Model, + relation: Relation, + kwargs: dict[str, bool], + error: type[Exception], + match: str, +) -> None: + with pytest.raises(error, match=match): + (1 * x(backend)).groupby(relation).sum(**kwargs) + + +@pytest.mark.parametrize( + ("frame", "match"), + [ + (pd.DataFrame({"a": [0]}), "missing \\['g'\\]"), + (pd.DataFrame({"a": [0], "g": [None]}), "'g' has null values"), + ], + ids=["missing-column", "null-group"], +) +def test_invalid_relation_raises(frame: pd.DataFrame, match: str) -> None: + with pytest.raises(ValueError, match=match): + Relation(frame, "a", "g") + + +LINKS = pd.Index(["l0", "l1", "l2", "l3"], name="link") +TIMES = pd.RangeIndex(2, name="snapshot") +PORTS = pd.DataFrame( + { + "link": list(LINKS) * 2, + "port": [0] * 4 + [1] * 4, + "bus": ["b0", "b0", "b1", "b2", "b1", "b2", "b2", "b0"], + } +) +EFFICIENCY = pd.DataFrame({"link": LINKS, "eff": [0.9, 0.8, 0.95, 0.7]}) + + +def port(i: int) -> Relation: + return Relation(PORTS[PORTS.port == i], "link", "bus") + + +def test_link_ports_to_bus(backend: Model) -> None: + p = backend.add_variables(coords=[TIMES, LINKS], name="p") + eff = from_long(EFFICIENCY, "eff", coords=[TIMES, LINKS]) + touching = p.groupby(Relation(PORTS, "link", "bus")).sum() + balance = (-1 * p).groupby(port(0)).sum() + (eff * p).groupby(port(1)).sum() + + weights = np.where(PORTS.port == 0, -1.0, np.tile(EFFICIENCY.eff, 2)) + for res, matrix in [ + (touching, incidence(PORTS[["link", "bus"]])), + (balance, incidence(PORTS[["link", "bus"]], weights)), + ]: + assert isinstance(res, LinearExpression) + assert res.is_sparse == backend.sparse + assert_equivalent(res, 1 * p @ matrix) diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 306aeaeb..cfcd226d 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -50,6 +50,8 @@ COEFFS = OVER_AB.expand_dims({TERM_DIM: 2}, axis=-1) SERIES = pd.Series([0, 1, 0], index=A, name="g") FRAME = pd.DataFrame({"k1": [0, 1, 0], "k2": [1, 1, 2]}, index=A) +RELATION_FRAME = pd.DataFrame({"a": [0, 0, 1, 2, 2], "g": ["u", "v", "v", "u", "w"]}) +RELATION = linopy.Relation(RELATION_FRAME, "a", "g") def var(e: LinearExpression, name: str) -> Variable: @@ -428,6 +430,7 @@ class Case(NamedTuple): "const": Case(grid, lambda e: (e + OVER_A).const, KEPT), "const-assign-new-dim": Case(grid, const_assigned(OVER_C), ASSIGNMENT), "coeffs-assign": Case(grid, coeffs_assigned, ASSIGN_COEFFS), + "groupby-relation": Case(grid, lambda e: e.groupby(RELATION).sum(), KEPT), } CONST_REJECTED: dict[str, tuple[Any, str]] = {