From 90b16c265e39584b60f3658c961ae81d5e89db9a Mon Sep 17 00:00:00 2001 From: Fabian Date: Tue, 4 Feb 2025 23:18:05 +0100 Subject: [PATCH 1/2] add unstack method --- doc/release_notes.rst | 1 + linopy/constraints.py | 2 ++ linopy/expressions.py | 2 ++ linopy/variables.py | 2 ++ test/test_constraint.py | 2 +- test/test_linear_expression.py | 7 +++++++ test/test_variable.py | 6 ++++++ 7 files changed, 21 insertions(+), 1 deletion(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 0a42d6b4..51063d13 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -6,6 +6,7 @@ Upcoming Version * The internal handling of `Solution` objects was improved for more consistency. Solution objects created from solver calls now preserve the exact index names from the input file. * Multiplication of a linear expression by a constant value may now introduce new dimensions. +* Added method `unstack` to `LinearExpression`, `Variable` and `Constraint` to unstack a dimension. Version 0.4.4 -------------- diff --git a/linopy/constraints.py b/linopy/constraints.py index e34a49a8..64864a1b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -675,6 +675,8 @@ def to_polars(self): stack = conwrap(Dataset.stack) + unstack = conwrap(Dataset.unstack) + iterate_slices = iterate_slices diff --git a/linopy/expressions.py b/linopy/expressions.py index cf6443be..a3bb8a75 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1472,6 +1472,8 @@ def to_polars(self) -> pl.DataFrame: stack = exprwrap(Dataset.stack) + unstack = exprwrap(Dataset.unstack) + iterate_slices = iterate_slices diff --git a/linopy/variables.py b/linopy/variables.py index 43f5f7d8..aca54ec5 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -1095,6 +1095,8 @@ def equals(self, other: Variable) -> bool: stack = varwrap(Dataset.stack) + unstack = varwrap(Dataset.unstack) + iterate_slices = iterate_slices diff --git a/test/test_constraint.py b/test/test_constraint.py index 0d0080fa..a5db519d 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -197,7 +197,7 @@ def test_constraint_wrapped_methods(x, y): con.rename({"first": "new_labels"}) con.rename_dims({"first": "new_labels"}) con.roll({"first": 1}) - con.stack(new_dim=("first", "second")) + con.stack(new_dim=("first", "second")).unstack("new_dim") def test_anonymous_constraint_sel(x, y): diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 017d165f..7ac75c15 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -669,6 +669,13 @@ def test_variable_stack(v): assert result.coord_dims == ("new",) +def test_variable_unstack(v): + result = v.to_linexpr().expand_dims("new_dim").stack(new=("new_dim", "dim_2")) + result = result.unstack("new") + assert isinstance(result, LinearExpression) + assert result.coord_dims == ("new_dim", "dim_2") + + def test_linear_expression_diff(v): diff = v.to_linexpr().diff("dim_2") assert diff.nterm == 2 diff --git a/test/test_variable.py b/test/test_variable.py index 321bc3d5..57315219 100644 --- a/test/test_variable.py +++ b/test/test_variable.py @@ -267,6 +267,12 @@ def test_variable_stack(x): assert result.dims == ("new",) +def test_variable_unstack(x): + result = x.expand_dims("new_dim").stack(new=("new_dim", "first")).unstack("new") + assert isinstance(result, linopy.variables.Variable) + assert result.dims == ("new_dim", "first") + + def test_variable_flat(x): result = x.flat assert isinstance(result, pd.DataFrame) From 972cd0695665420c6375acaacaf3ab21bf71f50e Mon Sep 17 00:00:00 2001 From: Fabian Hofmann Date: Wed, 5 Feb 2025 08:30:04 +0100 Subject: [PATCH 2/2] Update test/test_linear_expression.py --- test/test_linear_expression.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 7ac75c15..452ce9b2 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -669,7 +669,7 @@ def test_variable_stack(v): assert result.coord_dims == ("new",) -def test_variable_unstack(v): +def test_linear_expression_unstack(v): result = v.to_linexpr().expand_dims("new_dim").stack(new=("new_dim", "dim_2")) result = result.unstack("new") assert isinstance(result, LinearExpression)