From 8b56afbd024c484cbfebc6557aabe5dbd2fa7147 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 24 Oct 2026 16:13:25 +0200 Subject: [PATCH 1/3] feat: linopy.Table long-form inputs and relation group-by (#1019 step 5) Table pivots pandas/polars long tables onto target coords for bounds, coefficients, constants, rhs and penalty. expr.groupby(relation_frame).sum() scatters repeated member rows through from_grouper on both stores; quadratic relation group-by raises. --- doc/api.rst | 1 + doc/release_notes.rst | 9 ++ linopy/__init__.py | 2 + linopy/alignment.py | 19 +++- linopy/constraints.py | 5 +- linopy/expressions.py | 77 +++++++++++++- linopy/model.py | 6 +- linopy/store.py | 10 +- linopy/table.py | 103 +++++++++++++++++++ linopy/types.py | 3 + linopy/variables.py | 10 +- test/test_store_conformance.py | 11 ++ test/test_table.py | 178 +++++++++++++++++++++++++++++++++ 13 files changed, 416 insertions(+), 18 deletions(-) create mode 100644 linopy/table.py create mode 100644 test/test_table.py diff --git a/doc/api.rst b/doc/api.rst index 0a56c6d3..0ac8100b 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -644,6 +644,7 @@ Utilities align options ABSENT + Table ``options`` holds these keys: diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 55f78057..9c771637 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.Table(frame, value, default=nan)`` passes a constant in long form, a pandas or polars frame of key columns plus one value column, wherever a constant is accepted: ``add_variables(lower=..., upper=...)``, coefficients (``table * x``, ``x.to_linexpr(table)``), arithmetic constants (``expr + table``), right-hand sides (``expr <= table``, ``add_constraints(..., rhs)``) and penalties (``add_constraints(..., penalty=table)``, ``Constraint.soften``). Each key column names a dimension; the table is pivoted onto the target coordinates, key combinations not listed take ``default`` and the values broadcast over dimensions without a key column. Duplicate keys, key columns without a target dimension and key values outside the target coordinates raise. The pivot builds the full cube over the key dimensions, which is what the stores hold. (`#1049 `__) +* ``expr.groupby(relation).sum()`` accepts a relation, a two-column pandas or polars frame of member dimension and group, 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 ``@``. 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..e3529204 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -36,6 +36,7 @@ from linopy.remote import RemoteHandler from linopy.semantics import ABSENT from linopy.solvers import SolverFeature, available_solvers, licensed_solvers +from linopy.table import Table try: from linopy.remote import OetcCredentials, OetcHandler, OetcSettings # noqa: F401 @@ -61,6 +62,7 @@ "RemoteHandler", "Slopes", "SolverFeature", + "Table", "Variable", "Variables", "align", diff --git a/linopy/alignment.py b/linopy/alignment.py index 597e4a05..e1e30d6e 100644 --- a/linopy/alignment.py +++ b/linopy/alignment.py @@ -48,6 +48,7 @@ enforce_no_multiindex, first_mismatched_dim, ) +from linopy.table import Table from linopy.types import UNLABELED_TYPES, CoordsLike, DimsLike @@ -360,8 +361,9 @@ def as_dataarray( Picks the right constructor for each supported input type (pandas, polars, numpy, scalar, DataArray) and labels positional axes with - ``dims`` / ``coords``. The result is not reshaped against ``coords``: - dims are neither expanded nor reordered. Use :func:`broadcast_to_coords` when + ``dims`` / ``coords``. A :class:`~linopy.table.Table` is pivoted onto + ``coords`` for its key columns. The result is not reshaped against + ``coords``: dims are neither expanded nor reordered. Use :func:`broadcast_to_coords` when ``coords`` should govern the result's shape. Parameters @@ -380,7 +382,12 @@ def as_dataarray( DataArray The converted input, dims and entries as ``arr`` provides them. """ - if isinstance(arr, pd.Series | pd.DataFrame): + if isinstance(arr, Table): + if coords is None: + raise ValueError("A Table needs target coordinates to pivot onto.") + target = _coords_to_dict(coords, dims=dims) + arr = arr.to_dataarray({d: _as_index(v) for d, v in target.items()}) + elif isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, dims=dims, **kwargs) elif isinstance(arr, np.ndarray): arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs) @@ -407,6 +414,7 @@ def as_dataarray( np.ndarray, DataArray, pl.Series, + Table, ] supported_types_str = ", ".join([t.__name__ for t in supported_types]) raise TypeError( @@ -417,6 +425,11 @@ def as_dataarray( return fill_missing_coords(arr) +def pivot_table(arr: Any, coords: CoordsLike) -> Any: + """``arr`` pivoted onto ``coords`` if it is a :class:`~linopy.table.Table`, else ``arr``.""" + return as_dataarray(arr, coords=coords) if isinstance(arr, Table) else arr + + def _pair_axes_by_size( shape: tuple[int, ...], sizes: dict[Hashable, int] ) -> tuple[list[Hashable] | None, str | None]: diff --git a/linopy/constraints.py b/linopy/constraints.py index 9762d468..5bdcbf7b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -35,7 +35,7 @@ from xarray.core.utils import Frozen from linopy import expressions, variables -from linopy.alignment import as_dataarray, broadcast_to_coords +from linopy.alignment import as_dataarray, broadcast_to_coords, pivot_table from linopy.common import ( ConstraintLabelIndex, LocIndexer, @@ -159,8 +159,8 @@ def rhs_on_grid(store: CSRStore, rhs: Any) -> DataArray | Reason: grid returns the reason to fall back to the dense path. An expression rhs is moved to the lhs before, see :meth:`LinearExpression.to_constraint`. """ - da = as_dataarray(rhs) grid = store.grid + da = as_dataarray(pivot_table(rhs, grid.indexes)) if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(grid.dims): return "to_constraint(rhs over dimensions outside the grid)" if bool(da.isnull().any()): @@ -1081,6 +1081,7 @@ def soften( >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") >>> slack = budget_constraint.soften(penalty=budget_penalty) """ + penalty = pivot_table(penalty, self.coords) if not bool(np.all(np.asarray(penalty) > 0)): raise ValueError("Penalty is not positive.") diff --git a/linopy/expressions.py b/linopy/expressions.py index 483c6c59..20dda535 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -13,6 +13,7 @@ from abc import ABC, abstractmethod from collections.abc import ( Callable, + Collection, Hashable, ItemsView, Iterable, @@ -134,6 +135,7 @@ pruned_terms, summed, ) +from linopy.table import to_pandas from linopy.types import ( CONSTANT_TYPES, ConstantLike, @@ -247,6 +249,23 @@ def _multikey_value_frame(group: Any, data: Dataset) -> pd.DataFrame | None: return data[names].to_dataframe()[names] +def _relation_series(group: Any, dims: Collection[Hashable]) -> pd.Series | None: + """ + A relation grouper as a Series of groups indexed by member, or None. + + A relation is a two-column frame, member dim then group, whose rows + may repeat a member; a pandas frame indexed by a dim is a multi-key + grouper instead. + """ + if not isinstance(group, pd.DataFrame | pl.DataFrame) or group.shape[1] != 2: + return None + frame = to_pandas(group) + member, target = frame.columns + if member not in dims or frame.index.name in dims: + return None + return frame.set_index(member)[target] + + def _flatten_multidim_group( group: DataArray, data: Dataset ) -> tuple[pd.Series, Dataset]: @@ -433,7 +452,9 @@ class LinearExpressionGroupby: """ obj: xr.Dataset | BaseExpression - group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame + group: ( + Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame | pl.DataFrame + ) model: Any kwargs: Mapping[str, Any] = field(default_factory=dict) @@ -550,13 +571,27 @@ def sum( "`observed=True` is not supported with `use_fallback=True`." ) + relation = _relation_series(self.group, self.obj.dims) + if relation is not None: + 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, relation) + + 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 +684,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: pd.Series) -> 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 = str(relation.index.name) + if relation.name in store.dims: + raise ValueError( + f"relation group column {relation.name!r} must not name an " + f"existing dimension {list(store.dims)}." + ) + positions = store.grid.indexes[member].get_indexer(relation.index) + if (positions == -1).any(): + missing = relation.index[positions == -1].unique().tolist() + raise ValueError( + f"relation members {missing} are not in dimension {member!r}" + ) + rows = pd.RangeIndex(len(relation), 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.to_numpy(), index=rows, name=relation.name) + store_cls = result_store_cls(store, self.model) + res = store_cls.from_grouper( + gathered, self.model, grouper, str(relation.name), 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 +1896,7 @@ def diff(self, dim: str, n: int = 1) -> Self: def groupby( self, - group: DataFrame | Series | DataArray, + group: DataFrame | pl.DataFrame | Series | DataArray, restore_coord_dims: bool | None = None, **kwargs: Any, ) -> LinearExpressionGroupby: @@ -1842,9 +1908,12 @@ def groupby( Parameters ---------- - group : str, DataArray or IndexVariable + group : str, DataArray, IndexVariable, pandas.Series or pandas/polars DataFrame 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 two-column pandas or polars frame (member dimension, group) is a + relation whose rows may repeat a member: ``sum`` adds each row's + member to its group. restore_coord_dims : bool, optional If True, also restore the dimension order of multi-dimensional coordinates. diff --git a/linopy/model.py b/linopy/model.py index 9c1b59a6..053b28de 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -716,10 +716,10 @@ def add_variables( Parameters ---------- - lower : float/array_like, optional + lower : float/array_like/linopy.Table, optional Lower bound of the variable(s). For binary variables it defaults to 0 and, if given, must be 0 or 1. The default is -inf. - upper : TYPE, optional + upper : float/array_like/linopy.Table, optional Upper bound of the variable(s). For binary variables it defaults to 1 and, if given, must be 0 or 1. The default is inf. coords : list/dict/xarray.Coordinates, optional @@ -1205,7 +1205,7 @@ def add_constraints( coordinates. sign : str/array_like Relation between the lhs and rhs, valid values are {'=', '>=', '<='}. - rhs : int/float/array_like + rhs : int/float/array_like/linopy.Table Right hand side of the constraint(s). name : str, optional Reference name of the added constraints. The default None results 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/table.py b/linopy/table.py new file mode 100644 index 00000000..569d6982 --- /dev/null +++ b/linopy/table.py @@ -0,0 +1,103 @@ +""" +Long-form input tables. + +A :class:`Table` holds values in long form, one row per key combination, +and pivots onto the target coordinates of the position it is used in. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Mapping +from dataclasses import dataclass + +import numpy as np +import pandas as pd +import polars as pl +from xarray import DataArray + + +def to_pandas(frame: pd.DataFrame | pl.DataFrame) -> pd.DataFrame: + """``frame`` as a pandas DataFrame, a polars one converted column by column.""" + if isinstance(frame, pl.DataFrame): + return pd.DataFrame(frame.to_dict(as_series=False)) + return frame + + +@dataclass(frozen=True, eq=False) +class Table: + """ + A constant in long form: key columns plus one value column. + + Every column but ``value`` is a key column named after a dimension. Where + a constant is accepted (variable bounds, coefficients, arithmetic + constants, constraint right-hand sides), the table is pivoted onto the + target coordinates: key combinations not listed take ``default``, and + the values broadcast over target dimensions without a key column. + + The pivot materialises the full cube over the key dimensions, which is + what the expression stores hold anyway; a sparsely filled table costs + the memory of its cube. + + Parameters + ---------- + frame : pandas.DataFrame or polars.DataFrame + The long-form data. + value : str + Name of the value column. + default : float + Value of key combinations missing from ``frame``. Defaults to NaN. + + Examples + -------- + >>> import pandas as pd + >>> import linopy + >>> m = linopy.Model() + >>> i = pd.RangeIndex(3, name="i") + >>> t = linopy.Table(pd.DataFrame({"i": [0, 2], "ub": [5.0, 7.0]}), "ub", 0) + >>> m.add_variables(upper=t, coords=[i], name="x").upper.values + array([5., 0., 7.]) + """ + + frame: pd.DataFrame | pl.DataFrame + value: str + default: float = np.nan + + def __neg__(self) -> Table: + frame = to_pandas(self.frame) + negated = frame.assign(**{self.value: -frame[self.value]}) + return Table(negated, self.value, -self.default) + + def to_dataarray(self, coords: Mapping[Hashable, pd.Index]) -> DataArray: + """ + Pivot the table onto ``coords``, one dim per key column in column + order. Raises for key columns without a target dimension, key + values outside the target coordinates and duplicate keys. + """ + frame = to_pandas(self.frame) + if self.value not in frame.columns: + raise ValueError(f"Table has no value column {self.value!r}.") + keys = [k for k in frame.columns if k != self.value] + unknown = [k for k in keys if k not in coords] + if unknown or not keys: + raise ValueError( + f"Table key columns {keys} must name target dimensions " + f"{list(coords)}, got unmatched {unknown}." + ) + index = pd.MultiIndex.from_frame(frame[keys]) + if index.has_duplicates: + duplicates = index[index.duplicated()].unique().tolist() + raise ValueError(f"Table has duplicate keys {duplicates}.") + target = pd.MultiIndex.from_product([coords[k] for k in keys], names=keys) + outside = index.difference(target) + if len(outside): + raise ValueError( + f"Table keys {outside.tolist()} lie outside the target coordinates." + ) + values = pd.Series(frame[self.value].to_numpy(), index=index) + values = values.reindex(target, fill_value=self.default) + shape = [len(coords[k]) for k in keys] + return DataArray( + values.to_numpy().reshape(shape), + coords={k: coords[k] for k in keys}, + dims=keys, + ) diff --git a/linopy/types.py b/linopy/types.py index 6ca69efb..b41fa4c5 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -9,6 +9,8 @@ from xarray import DataArray from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates +from linopy.table import Table + if TYPE_CHECKING: from linopy.constraints import AnonymousScalarConstraint, Constraint from linopy.expressions import ( @@ -33,6 +35,7 @@ | Series | DataFrame | pl.Series + | Table ) CONSTANT_TYPES: tuple[type, ...] = get_args(ConstantLike) diff --git a/linopy/variables.py b/linopy/variables.py index 794de240..8f8704f6 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -31,7 +31,7 @@ from xarray.core.utils import Frozen import linopy.expressions as expressions -from linopy.alignment import as_dataarray, broadcast_to_coords +from linopy.alignment import as_dataarray, broadcast_to_coords, pivot_table from linopy.common import ( LocIndexer, VariableLabelIndex, @@ -382,6 +382,7 @@ def _linexpr( self, store_cls: type[Store], coefficient: ConstantLike = 1 ) -> expressions.LinearExpression: """:meth:`to_linexpr` built on ``store_cls``.""" + coefficient = pivot_table(coefficient, self.coords) if not np.isscalar(coefficient): coeff_da = as_dataarray(coefficient) enforce_aux_conflict([self.labels, coeff_da]) @@ -764,7 +765,7 @@ def dot(self, other: ndarray | Variable) -> QuadraticExpression | LinearExpressi def groupby( self, - group: DataArray, + group: pd.DataFrame | pl.DataFrame | pd.Series | DataArray, 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 or pandas/polars DataFrame 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 two-column pandas or polars frame (member dimension, group) is a + relation whose rows may repeat a member: ``sum`` adds each row's + member to its group. restore_coord_dims : bool, optional If True, also restore the dimension order of multi-dimensional coordinates. diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 306aeaeb..50bafc50 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -50,6 +50,13 @@ 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) +TABLE_AB = linopy.Table( + pd.DataFrame({"a": [0, 2, 1, 0], "b": ["p", "p", "q", "q"], "v": [1.0, 2, 3, 4]}), + "v", + default=5.0, +) +TABLE_A = linopy.Table(pd.DataFrame({"a": [2, 0, 1], "v": [1.0, 2, 3]}), "v") +RELATION = pd.DataFrame({"a": [0, 0, 1, 2, 2], "g": ["u", "v", "v", "u", "w"]}) def var(e: LinearExpression, name: str) -> Variable: @@ -428,6 +435,10 @@ 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), + "mul-table": Case(grid, lambda e: e * TABLE_AB, KEPT), + "add-table-broadcast": Case(grid, lambda e: e + TABLE_A, KEPT), + "to-constraint-table": Case(grid, lambda e: e <= TABLE_AB, KEPT), + "groupby-relation": Case(grid, lambda e: e.groupby(RELATION).sum(), KEPT), } CONST_REJECTED: dict[str, tuple[Any, str]] = { diff --git a/test/test_table.py b/test/test_table.py new file mode 100644 index 00000000..b7d7e602 --- /dev/null +++ b/test/test_table.py @@ -0,0 +1,178 @@ +"""Long-form inputs: ``linopy.Table`` constants and relation group-by.""" + +from __future__ import annotations + +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, TABLE_A, TABLE_AB, A, B, assert_same + +from linopy import LinearExpression, Model, Table, Variable +from linopy.testing import assert_equivalent + +TABLES: dict[str, tuple[Table, xr.DataArray]] = { + "all-keys-default": ( + TABLE_AB, + xr.DataArray([[1.0, 4], [5, 3], [2, 5]], coords=[A, B]), + ), + "missing-key-column": (TABLE_A, xr.DataArray([2.0, 3, 1], coords=[A])), + "polars": ( + Table(pl.DataFrame({"b": ["q", "p"], "a": [1, 1], "v": [3.0, 4]}), "v", 1), + xr.DataArray([[1.0, 4, 1], [1, 3, 1]], coords=[B, A]), + ), +} + + +def x(m: Model) -> Variable: + return m.add_variables(coords=[A, B], name="x") + + +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, + "div": lambda m, c: x(m) / c, + "quad-mul": lambda m, c: c * x(m) ** 2, +} + + +@pytest.mark.parametrize("table", list(TABLES)) +@pytest.mark.parametrize("position", list(POSITIONS)) +def test_table_matches_dataarray(backend: Model, position: str, table: str) -> None: + tab, da = TABLES[table] + want = POSITIONS[position](Model(sparse=backend.sparse), da) + assert_same(POSITIONS[position](backend, tab), want) + + +INVALID: dict[str, tuple[Table, str]] = { + "duplicate-keys": ( + Table(pd.DataFrame({"a": [0, 0], "v": [1.0, 2]}), "v"), + "duplicate keys", + ), + "unknown-key-column": ( + Table(pd.DataFrame({"a": [0], "z": [0], "v": [1.0]}), "v"), + "unmatched \\['z'\\]", + ), + "key-outside-coords": ( + Table(pd.DataFrame({"a": [7], "v": [1.0]}), "v"), + "outside the target coordinates", + ), + "no-value-column": (Table(pd.DataFrame({"a": [0]}), "v"), "no value column"), + "nan-default": (Table(pd.DataFrame({"a": [0], "v": [1.0]}), "v"), "NaN"), +} + + +@pytest.mark.parametrize("position", ["coefficient", "add-constant", "rhs-le"]) +@pytest.mark.parametrize("case", list(INVALID)) +def test_invalid_table_raises(backend: Model, case: str, position: str) -> None: + tab, match = INVALID[case] + with pytest.raises(ValueError, match=match): + POSITIONS[position](backend, tab) + + +def test_table_reflected_comparison(backend: Model) -> None: + assert_same( + TABLE_AB >= 2 * x(backend), 2 * x(Model(sparse=backend.sparse)) <= TABLE_AB + ) + + +def test_table_penalty(backend: Model) -> None: + def softened(m: Model, penalty: Any) -> LinearExpression: + v = x(m) + m.add_objective(v.sum()) + m.add_constraints(v >= 0, penalty=penalty) + return m.objective.expression + + want = softened(Model(sparse=backend.sparse), TABLES["all-keys-default"][1]) + assert_same(softened(backend, TABLE_AB), want) + + +def test_table_without_target_coords_raises() -> None: + with pytest.raises(ValueError, match="needs target coordinates"): + Model().add_variables(lower=TABLE_A) + + +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( + "relation", [RELATION, RELATION[RELATION.a != 1]], ids=["all", "unlisted"] +) +@pytest.mark.parametrize( + "to_frame", [pd.DataFrame, lambda r: pl.from_dict(r.to_dict("list"))] +) +def test_relation_groupby_matches_incidence_matmul( + backend: Model, to_frame: Callable[[pd.DataFrame], Any], relation: pd.DataFrame +) -> None: + expr = 2 * backend.add_variables(coords=[B, A], name="x") + 1 + res = expr.groupby(to_frame(relation)).sum() + assert res.is_sparse == backend.sparse + assert_equivalent(res, expr @ incidence(relation).reindex(a=A, fill_value=0.0)) + + +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", "match"), + [ + (pd.DataFrame({"a": [0, 9], "g": ["u", "v"]}), {}, r"relation members \[9\]"), + (pd.DataFrame({"a": [0, 1], "b": ["p", "q"]}), {}, "must not name"), + (RELATION, {"use_fallback": True}, "not supported"), + (RELATION, {"observed": True}, "not supported"), + ], +) +def test_invalid_relation_groupby_raises( + backend: Model, relation: pd.DataFrame, kwargs: dict[str, bool], match: str +) -> None: + with pytest.raises(ValueError, match=match): + (1 * x(backend)).groupby(relation).sum(**kwargs) + + +LINKS = pd.Index(["l0", "l1", "l2", "l3"], name="link") +SNAPSHOTS = 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) -> pd.DataFrame: + return PORTS.loc[PORTS.port == i, ["link", "bus"]] + + +def test_link_ports_to_bus(backend: Model) -> None: + p = backend.add_variables(coords=[SNAPSHOTS, LINKS], name="p") + eff = Table(EFFICIENCY, "eff") + touching = p.groupby(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) From 5c8bf174f5e724f4dc8d425f72dc6a1dc8937fbd Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 24 Oct 2026 16:16:05 +0200 Subject: [PATCH 2/3] test: annotations for mypy over the test suite with Table in ConstantLike --- test/test_fix_relax.py | 2 +- test/test_table.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_fix_relax.py b/test/test_fix_relax.py index 30ee12f7..bb55828e 100644 --- a/test/test_fix_relax.py +++ b/test/test_fix_relax.py @@ -648,7 +648,7 @@ def variable(self) -> Variable: @pytest.mark.parametrize("value, expected", ALIGNED_VALUES) def test_aligns_to_named_dimension( - self, variable: Variable, value: ConstantLike, expected: ConstantLike + self, variable: Variable, value: ConstantLike, expected: list[float] ) -> None: variable.fix(value) assert variable.lower.dims == ("time",) diff --git a/test/test_table.py b/test/test_table.py index b7d7e602..817ec674 100644 --- a/test/test_table.py +++ b/test/test_table.py @@ -12,7 +12,7 @@ import xarray as xr from test_store_conformance import RELATION, TABLE_A, TABLE_AB, A, B, assert_same -from linopy import LinearExpression, Model, Table, Variable +from linopy import LinearExpression, Model, QuadraticExpression, Table, Variable from linopy.testing import assert_equivalent TABLES: dict[str, tuple[Table, xr.DataArray]] = { @@ -88,7 +88,7 @@ def test_table_reflected_comparison(backend: Model) -> None: def test_table_penalty(backend: Model) -> None: - def softened(m: Model, penalty: Any) -> LinearExpression: + def softened(m: Model, penalty: Any) -> LinearExpression | QuadraticExpression: v = x(m) m.add_objective(v.sum()) m.add_constraints(v >= 0, penalty=penalty) From 6536fbf0c9933b61fc6c4220cbc354d8a5c19e7c Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 24 Oct 2026 21:14:10 +0200 Subject: [PATCH 3/3] refactor: eager from_long with polars-join pivot and explicit Relation replace lazy Table Long tables convert to a DataArray at the call site with strict key checks; nothing leaks into ConstantLike or alignment. groupby takes linopy.Relation instead of detecting relation frames by shape. --- doc/api.rst | 3 +- doc/release_notes.rst | 4 +- linopy/__init__.py | 5 +- linopy/alignment.py | 19 +- linopy/constraints.py | 5 +- linopy/expressions.py | 68 +++--- linopy/long.py | 284 ++++++++++++++++++++++++ linopy/model.py | 6 +- linopy/table.py | 103 --------- linopy/types.py | 3 - linopy/variables.py | 14 +- test/test_fix_relax.py | 2 +- test/test_long.py | 390 +++++++++++++++++++++++++++++++++ test/test_store_conformance.py | 12 +- test/test_table.py | 178 --------------- 15 files changed, 726 insertions(+), 370 deletions(-) create mode 100644 linopy/long.py delete mode 100644 linopy/table.py create mode 100644 test/test_long.py delete mode 100644 test/test_table.py diff --git a/doc/api.rst b/doc/api.rst index 0ac8100b..0f22b881 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -644,7 +644,8 @@ Utilities align options ABSENT - Table + from_long + Relation ``options`` holds these keys: diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 9c771637..88583246 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -64,8 +64,8 @@ Version 1.0.0 The following content was generated by AI. -* ``linopy.Table(frame, value, default=nan)`` passes a constant in long form, a pandas or polars frame of key columns plus one value column, wherever a constant is accepted: ``add_variables(lower=..., upper=...)``, coefficients (``table * x``, ``x.to_linexpr(table)``), arithmetic constants (``expr + table``), right-hand sides (``expr <= table``, ``add_constraints(..., rhs)``) and penalties (``add_constraints(..., penalty=table)``, ``Constraint.soften``). Each key column names a dimension; the table is pivoted onto the target coordinates, key combinations not listed take ``default`` and the values broadcast over dimensions without a key column. Duplicate keys, key columns without a target dimension and key values outside the target coordinates raise. The pivot builds the full cube over the key dimensions, which is what the stores hold. (`#1049 `__) -* ``expr.groupby(relation).sum()`` accepts a relation, a two-column pandas or polars frame of member dimension and group, 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 ``@``. It stays sparse on a sparse expression; a quadratic expression raises ``NotImplementedError``. (`#1049 `__) +* ``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** diff --git a/linopy/__init__.py b/linopy/__init__.py index e3529204..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 ( @@ -36,7 +37,6 @@ from linopy.remote import RemoteHandler from linopy.semantics import ABSENT from linopy.solvers import SolverFeature, available_solvers, licensed_solvers -from linopy.table import Table try: from linopy.remote import OetcCredentials, OetcHandler, OetcSettings # noqa: F401 @@ -61,14 +61,15 @@ "QuadraticExpression", "RemoteHandler", "Slopes", + "Relation", "SolverFeature", - "Table", "Variable", "Variables", "align", "available_solvers", "licensed_solvers", "breakpoints", + "from_long", "merge", "options", "read_netcdf", diff --git a/linopy/alignment.py b/linopy/alignment.py index e1e30d6e..597e4a05 100644 --- a/linopy/alignment.py +++ b/linopy/alignment.py @@ -48,7 +48,6 @@ enforce_no_multiindex, first_mismatched_dim, ) -from linopy.table import Table from linopy.types import UNLABELED_TYPES, CoordsLike, DimsLike @@ -361,9 +360,8 @@ def as_dataarray( Picks the right constructor for each supported input type (pandas, polars, numpy, scalar, DataArray) and labels positional axes with - ``dims`` / ``coords``. A :class:`~linopy.table.Table` is pivoted onto - ``coords`` for its key columns. The result is not reshaped against - ``coords``: dims are neither expanded nor reordered. Use :func:`broadcast_to_coords` when + ``dims`` / ``coords``. The result is not reshaped against ``coords``: + dims are neither expanded nor reordered. Use :func:`broadcast_to_coords` when ``coords`` should govern the result's shape. Parameters @@ -382,12 +380,7 @@ def as_dataarray( DataArray The converted input, dims and entries as ``arr`` provides them. """ - if isinstance(arr, Table): - if coords is None: - raise ValueError("A Table needs target coordinates to pivot onto.") - target = _coords_to_dict(coords, dims=dims) - arr = arr.to_dataarray({d: _as_index(v) for d, v in target.items()}) - elif isinstance(arr, pd.Series | pd.DataFrame): + if isinstance(arr, pd.Series | pd.DataFrame): arr = pandas_to_dataarray(arr, dims=dims, **kwargs) elif isinstance(arr, np.ndarray): arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs) @@ -414,7 +407,6 @@ def as_dataarray( np.ndarray, DataArray, pl.Series, - Table, ] supported_types_str = ", ".join([t.__name__ for t in supported_types]) raise TypeError( @@ -425,11 +417,6 @@ def as_dataarray( return fill_missing_coords(arr) -def pivot_table(arr: Any, coords: CoordsLike) -> Any: - """``arr`` pivoted onto ``coords`` if it is a :class:`~linopy.table.Table`, else ``arr``.""" - return as_dataarray(arr, coords=coords) if isinstance(arr, Table) else arr - - def _pair_axes_by_size( shape: tuple[int, ...], sizes: dict[Hashable, int] ) -> tuple[list[Hashable] | None, str | None]: diff --git a/linopy/constraints.py b/linopy/constraints.py index 5bdcbf7b..9762d468 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -35,7 +35,7 @@ from xarray.core.utils import Frozen from linopy import expressions, variables -from linopy.alignment import as_dataarray, broadcast_to_coords, pivot_table +from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( ConstraintLabelIndex, LocIndexer, @@ -159,8 +159,8 @@ def rhs_on_grid(store: CSRStore, rhs: Any) -> DataArray | Reason: grid returns the reason to fall back to the dense path. An expression rhs is moved to the lhs before, see :meth:`LinearExpression.to_constraint`. """ + da = as_dataarray(rhs) grid = store.grid - da = as_dataarray(pivot_table(rhs, grid.indexes)) if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(grid.dims): return "to_constraint(rhs over dimensions outside the grid)" if bool(da.isnull().any()): @@ -1081,7 +1081,6 @@ def soften( >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") >>> slack = budget_constraint.soften(penalty=budget_penalty) """ - penalty = pivot_table(penalty, self.coords) if not bool(np.all(np.asarray(penalty) > 0)): raise ValueError("Penalty is not positive.") diff --git a/linopy/expressions.py b/linopy/expressions.py index 20dda535..d07e7cdc 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -13,7 +13,6 @@ from abc import ABC, abstractmethod from collections.abc import ( Callable, - Collection, Hashable, ItemsView, Iterable, @@ -107,6 +106,7 @@ TERM_DIM, ) from linopy.csr import CSRStore, Grid, _densify_notice +from linopy.long import Relation from linopy.semantics import ( AbsentType, FillValueLike, @@ -135,7 +135,6 @@ pruned_terms, summed, ) -from linopy.table import to_pandas from linopy.types import ( CONSTANT_TYPES, ConstantLike, @@ -220,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: @@ -249,23 +255,6 @@ def _multikey_value_frame(group: Any, data: Dataset) -> pd.DataFrame | None: return data[names].to_dataframe()[names] -def _relation_series(group: Any, dims: Collection[Hashable]) -> pd.Series | None: - """ - A relation grouper as a Series of groups indexed by member, or None. - - A relation is a two-column frame, member dim then group, whose rows - may repeat a member; a pandas frame indexed by a dim is a multi-key - grouper instead. - """ - if not isinstance(group, pd.DataFrame | pl.DataFrame) or group.shape[1] != 2: - return None - frame = to_pandas(group) - member, target = frame.columns - if member not in dims or frame.index.name in dims: - return None - return frame.set_index(member)[target] - - def _flatten_multidim_group( group: DataArray, data: Dataset ) -> tuple[pd.Series, Dataset]: @@ -452,9 +441,7 @@ class LinearExpressionGroupby: """ obj: xr.Dataset | BaseExpression - group: ( - Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame | pl.DataFrame - ) + group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame | Relation model: Any kwargs: Mapping[str, Any] = field(default_factory=dict) @@ -571,8 +558,7 @@ def sum( "`observed=True` is not supported with `use_fallback=True`." ) - relation = _relation_series(self.group, self.obj.dims) - if relation is not None: + if isinstance(self.group, Relation): if isinstance(self.obj, QuadraticExpression): raise NotImplementedError( "Grouping a quadratic expression by a relation is not supported." @@ -582,7 +568,7 @@ def sum( "`use_fallback` and `observed` are not supported when " "grouping by a relation." ) - return self._relation_sum(self._store, relation) + return self._relation_sum(self._store, self.group) store = self._store @@ -684,7 +670,7 @@ 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: pd.Series) -> LinearExpression: + 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 @@ -692,26 +678,26 @@ def _relation_sum(self, store: Store, relation: pd.Series) -> LinearExpression: positional member dim, then scattered by :meth:`~linopy.store.Store.from_grouper`. """ - member = str(relation.index.name) - if relation.name in store.dims: + member, group = relation.member, relation.group + if group in store.dims: raise ValueError( - f"relation group column {relation.name!r} must not name an " + f"relation group column {group!r} must not name an " f"existing dimension {list(store.dims)}." ) - positions = store.grid.indexes[member].get_indexer(relation.index) - if (positions == -1).any(): - missing = relation.index[positions == -1].unique().tolist() + if member not in store.dims: raise ValueError( - f"relation members {missing} are not in dimension {member!r}" + f"relation member column {member!r} must name a dimension " + f"of {list(store.dims)}." ) - rows = pd.RangeIndex(len(relation), name=member) + 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.to_numpy(), index=rows, name=relation.name) + 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, str(relation.name), True, store.dims + gathered, self.model, grouper, group, True, store.dims ) return LinearExpression._from_store(res, self.model) @@ -1896,7 +1882,7 @@ def diff(self, dim: str, n: int = 1) -> Self: def groupby( self, - group: DataFrame | pl.DataFrame | Series | DataArray, + group: DataFrame | Series | DataArray | Relation, restore_coord_dims: bool | None = None, **kwargs: Any, ) -> LinearExpressionGroupby: @@ -1908,12 +1894,12 @@ def groupby( Parameters ---------- - group : str, DataArray, IndexVariable, pandas.Series or pandas/polars DataFrame + 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 two-column pandas or polars frame (member dimension, group) is a - relation whose rows may repeat a member: ``sum`` adds each row's - member to its group. + 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/model.py b/linopy/model.py index 053b28de..9c1b59a6 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -716,10 +716,10 @@ def add_variables( Parameters ---------- - lower : float/array_like/linopy.Table, optional + lower : float/array_like, optional Lower bound of the variable(s). For binary variables it defaults to 0 and, if given, must be 0 or 1. The default is -inf. - upper : float/array_like/linopy.Table, optional + upper : TYPE, optional Upper bound of the variable(s). For binary variables it defaults to 1 and, if given, must be 0 or 1. The default is inf. coords : list/dict/xarray.Coordinates, optional @@ -1205,7 +1205,7 @@ def add_constraints( coordinates. sign : str/array_like Relation between the lhs and rhs, valid values are {'=', '>=', '<='}. - rhs : int/float/array_like/linopy.Table + rhs : int/float/array_like Right hand side of the constraint(s). name : str, optional Reference name of the added constraints. The default None results diff --git a/linopy/table.py b/linopy/table.py deleted file mode 100644 index 569d6982..00000000 --- a/linopy/table.py +++ /dev/null @@ -1,103 +0,0 @@ -""" -Long-form input tables. - -A :class:`Table` holds values in long form, one row per key combination, -and pivots onto the target coordinates of the position it is used in. -""" - -from __future__ import annotations - -from collections.abc import Hashable, Mapping -from dataclasses import dataclass - -import numpy as np -import pandas as pd -import polars as pl -from xarray import DataArray - - -def to_pandas(frame: pd.DataFrame | pl.DataFrame) -> pd.DataFrame: - """``frame`` as a pandas DataFrame, a polars one converted column by column.""" - if isinstance(frame, pl.DataFrame): - return pd.DataFrame(frame.to_dict(as_series=False)) - return frame - - -@dataclass(frozen=True, eq=False) -class Table: - """ - A constant in long form: key columns plus one value column. - - Every column but ``value`` is a key column named after a dimension. Where - a constant is accepted (variable bounds, coefficients, arithmetic - constants, constraint right-hand sides), the table is pivoted onto the - target coordinates: key combinations not listed take ``default``, and - the values broadcast over target dimensions without a key column. - - The pivot materialises the full cube over the key dimensions, which is - what the expression stores hold anyway; a sparsely filled table costs - the memory of its cube. - - Parameters - ---------- - frame : pandas.DataFrame or polars.DataFrame - The long-form data. - value : str - Name of the value column. - default : float - Value of key combinations missing from ``frame``. Defaults to NaN. - - Examples - -------- - >>> import pandas as pd - >>> import linopy - >>> m = linopy.Model() - >>> i = pd.RangeIndex(3, name="i") - >>> t = linopy.Table(pd.DataFrame({"i": [0, 2], "ub": [5.0, 7.0]}), "ub", 0) - >>> m.add_variables(upper=t, coords=[i], name="x").upper.values - array([5., 0., 7.]) - """ - - frame: pd.DataFrame | pl.DataFrame - value: str - default: float = np.nan - - def __neg__(self) -> Table: - frame = to_pandas(self.frame) - negated = frame.assign(**{self.value: -frame[self.value]}) - return Table(negated, self.value, -self.default) - - def to_dataarray(self, coords: Mapping[Hashable, pd.Index]) -> DataArray: - """ - Pivot the table onto ``coords``, one dim per key column in column - order. Raises for key columns without a target dimension, key - values outside the target coordinates and duplicate keys. - """ - frame = to_pandas(self.frame) - if self.value not in frame.columns: - raise ValueError(f"Table has no value column {self.value!r}.") - keys = [k for k in frame.columns if k != self.value] - unknown = [k for k in keys if k not in coords] - if unknown or not keys: - raise ValueError( - f"Table key columns {keys} must name target dimensions " - f"{list(coords)}, got unmatched {unknown}." - ) - index = pd.MultiIndex.from_frame(frame[keys]) - if index.has_duplicates: - duplicates = index[index.duplicated()].unique().tolist() - raise ValueError(f"Table has duplicate keys {duplicates}.") - target = pd.MultiIndex.from_product([coords[k] for k in keys], names=keys) - outside = index.difference(target) - if len(outside): - raise ValueError( - f"Table keys {outside.tolist()} lie outside the target coordinates." - ) - values = pd.Series(frame[self.value].to_numpy(), index=index) - values = values.reindex(target, fill_value=self.default) - shape = [len(coords[k]) for k in keys] - return DataArray( - values.to_numpy().reshape(shape), - coords={k: coords[k] for k in keys}, - dims=keys, - ) diff --git a/linopy/types.py b/linopy/types.py index b41fa4c5..6ca69efb 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -9,8 +9,6 @@ from xarray import DataArray from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates -from linopy.table import Table - if TYPE_CHECKING: from linopy.constraints import AnonymousScalarConstraint, Constraint from linopy.expressions import ( @@ -35,7 +33,6 @@ | Series | DataFrame | pl.Series - | Table ) CONSTANT_TYPES: tuple[type, ...] = get_args(ConstantLike) diff --git a/linopy/variables.py b/linopy/variables.py index 8f8704f6..0e93ae28 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -31,7 +31,7 @@ from xarray.core.utils import Frozen import linopy.expressions as expressions -from linopy.alignment import as_dataarray, broadcast_to_coords, pivot_table +from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( LocIndexer, VariableLabelIndex, @@ -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, @@ -382,7 +383,6 @@ def _linexpr( self, store_cls: type[Store], coefficient: ConstantLike = 1 ) -> expressions.LinearExpression: """:meth:`to_linexpr` built on ``store_cls``.""" - coefficient = pivot_table(coefficient, self.coords) if not np.isscalar(coefficient): coeff_da = as_dataarray(coefficient) enforce_aux_conflict([self.labels, coeff_da]) @@ -765,7 +765,7 @@ def dot(self, other: ndarray | Variable) -> QuadraticExpression | LinearExpressi def groupby( self, - group: pd.DataFrame | pl.DataFrame | pd.Series | DataArray, + group: DataFrame | pd.Series | DataArray | Relation, restore_coord_dims: bool | None = None, ) -> LinearExpressionGroupby: """ @@ -776,12 +776,12 @@ def groupby( Parameters ---------- - group : str, DataArray, IndexVariable, pandas.Series or pandas/polars DataFrame + 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 two-column pandas or polars frame (member dimension, group) is a - relation whose rows may repeat a member: ``sum`` adds each row's - member to its group. + 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_fix_relax.py b/test/test_fix_relax.py index bb55828e..30ee12f7 100644 --- a/test/test_fix_relax.py +++ b/test/test_fix_relax.py @@ -648,7 +648,7 @@ def variable(self) -> Variable: @pytest.mark.parametrize("value, expected", ALIGNED_VALUES) def test_aligns_to_named_dimension( - self, variable: Variable, value: ConstantLike, expected: list[float] + self, variable: Variable, value: ConstantLike, expected: ConstantLike ) -> None: variable.fix(value) assert variable.lower.dims == ("time",) 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 50bafc50..cfcd226d 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -50,13 +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) -TABLE_AB = linopy.Table( - pd.DataFrame({"a": [0, 2, 1, 0], "b": ["p", "p", "q", "q"], "v": [1.0, 2, 3, 4]}), - "v", - default=5.0, -) -TABLE_A = linopy.Table(pd.DataFrame({"a": [2, 0, 1], "v": [1.0, 2, 3]}), "v") -RELATION = pd.DataFrame({"a": [0, 0, 1, 2, 2], "g": ["u", "v", "v", "u", "w"]}) +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: @@ -435,9 +430,6 @@ 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), - "mul-table": Case(grid, lambda e: e * TABLE_AB, KEPT), - "add-table-broadcast": Case(grid, lambda e: e + TABLE_A, KEPT), - "to-constraint-table": Case(grid, lambda e: e <= TABLE_AB, KEPT), "groupby-relation": Case(grid, lambda e: e.groupby(RELATION).sum(), KEPT), } diff --git a/test/test_table.py b/test/test_table.py deleted file mode 100644 index 817ec674..00000000 --- a/test/test_table.py +++ /dev/null @@ -1,178 +0,0 @@ -"""Long-form inputs: ``linopy.Table`` constants and relation group-by.""" - -from __future__ import annotations - -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, TABLE_A, TABLE_AB, A, B, assert_same - -from linopy import LinearExpression, Model, QuadraticExpression, Table, Variable -from linopy.testing import assert_equivalent - -TABLES: dict[str, tuple[Table, xr.DataArray]] = { - "all-keys-default": ( - TABLE_AB, - xr.DataArray([[1.0, 4], [5, 3], [2, 5]], coords=[A, B]), - ), - "missing-key-column": (TABLE_A, xr.DataArray([2.0, 3, 1], coords=[A])), - "polars": ( - Table(pl.DataFrame({"b": ["q", "p"], "a": [1, 1], "v": [3.0, 4]}), "v", 1), - xr.DataArray([[1.0, 4, 1], [1, 3, 1]], coords=[B, A]), - ), -} - - -def x(m: Model) -> Variable: - return m.add_variables(coords=[A, B], name="x") - - -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, - "div": lambda m, c: x(m) / c, - "quad-mul": lambda m, c: c * x(m) ** 2, -} - - -@pytest.mark.parametrize("table", list(TABLES)) -@pytest.mark.parametrize("position", list(POSITIONS)) -def test_table_matches_dataarray(backend: Model, position: str, table: str) -> None: - tab, da = TABLES[table] - want = POSITIONS[position](Model(sparse=backend.sparse), da) - assert_same(POSITIONS[position](backend, tab), want) - - -INVALID: dict[str, tuple[Table, str]] = { - "duplicate-keys": ( - Table(pd.DataFrame({"a": [0, 0], "v": [1.0, 2]}), "v"), - "duplicate keys", - ), - "unknown-key-column": ( - Table(pd.DataFrame({"a": [0], "z": [0], "v": [1.0]}), "v"), - "unmatched \\['z'\\]", - ), - "key-outside-coords": ( - Table(pd.DataFrame({"a": [7], "v": [1.0]}), "v"), - "outside the target coordinates", - ), - "no-value-column": (Table(pd.DataFrame({"a": [0]}), "v"), "no value column"), - "nan-default": (Table(pd.DataFrame({"a": [0], "v": [1.0]}), "v"), "NaN"), -} - - -@pytest.mark.parametrize("position", ["coefficient", "add-constant", "rhs-le"]) -@pytest.mark.parametrize("case", list(INVALID)) -def test_invalid_table_raises(backend: Model, case: str, position: str) -> None: - tab, match = INVALID[case] - with pytest.raises(ValueError, match=match): - POSITIONS[position](backend, tab) - - -def test_table_reflected_comparison(backend: Model) -> None: - assert_same( - TABLE_AB >= 2 * x(backend), 2 * x(Model(sparse=backend.sparse)) <= TABLE_AB - ) - - -def test_table_penalty(backend: Model) -> None: - 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 - - want = softened(Model(sparse=backend.sparse), TABLES["all-keys-default"][1]) - assert_same(softened(backend, TABLE_AB), want) - - -def test_table_without_target_coords_raises() -> None: - with pytest.raises(ValueError, match="needs target coordinates"): - Model().add_variables(lower=TABLE_A) - - -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( - "relation", [RELATION, RELATION[RELATION.a != 1]], ids=["all", "unlisted"] -) -@pytest.mark.parametrize( - "to_frame", [pd.DataFrame, lambda r: pl.from_dict(r.to_dict("list"))] -) -def test_relation_groupby_matches_incidence_matmul( - backend: Model, to_frame: Callable[[pd.DataFrame], Any], relation: pd.DataFrame -) -> None: - expr = 2 * backend.add_variables(coords=[B, A], name="x") + 1 - res = expr.groupby(to_frame(relation)).sum() - assert res.is_sparse == backend.sparse - assert_equivalent(res, expr @ incidence(relation).reindex(a=A, fill_value=0.0)) - - -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", "match"), - [ - (pd.DataFrame({"a": [0, 9], "g": ["u", "v"]}), {}, r"relation members \[9\]"), - (pd.DataFrame({"a": [0, 1], "b": ["p", "q"]}), {}, "must not name"), - (RELATION, {"use_fallback": True}, "not supported"), - (RELATION, {"observed": True}, "not supported"), - ], -) -def test_invalid_relation_groupby_raises( - backend: Model, relation: pd.DataFrame, kwargs: dict[str, bool], match: str -) -> None: - with pytest.raises(ValueError, match=match): - (1 * x(backend)).groupby(relation).sum(**kwargs) - - -LINKS = pd.Index(["l0", "l1", "l2", "l3"], name="link") -SNAPSHOTS = 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) -> pd.DataFrame: - return PORTS.loc[PORTS.port == i, ["link", "bus"]] - - -def test_link_ports_to_bus(backend: Model) -> None: - p = backend.add_variables(coords=[SNAPSHOTS, LINKS], name="p") - eff = Table(EFFICIENCY, "eff") - touching = p.groupby(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)