Merge pull request #297 from PyPSA/methods-addition

xarray methods additions

doc/release_notes.rst

@@ -19,6 +19,10 @@ Version 0.3.10
 * A new memory-efficient and super fast LP file writing method was added which uses the `Polars package <https://github.com/pola-rs/polars>`_. It is still in experimental mode but seems to be very promising. Activate it with the `io_api="lp-polars"` argument in the `solve` function.
 
 
+* The Constraint class now supports the methods `assign`, `assign_attrs`, `assign_coords`, `broadcast_like`, `chunk`, `drop_sel`, `drop_isel`, `expand_dims`, `sel`, `isel`, `shift`, `swap_dims`, `set_index`, `reindex`, `reindex_like`, `rename`, `rename_dims`, `roll`, `stack`. These methods allow to manipulation of a (anonymous) constraint more flexibly.
+
+* The Variable, expressions and Constraint classes now have new methods `swap_dims` and `set_index`. The `swap_dims` method allows to swap the dimensions of the object. The `set_index` method allows to set a new index for the object. Both methods are useful for reshaping the object more flexibly.
+
 Version 0.3.9
 -------------
 

linopy/constraints.py

@@ -591,10 +591,49 @@ def to_polars(self):
         df = df.filter(is_non_null & ~prev_non_is_null | ~is_non_null)
         return df[["labels", "coeffs", "vars", "sign", "rhs"]]
 
+    # Wrapped function which would convert variable to dataarray
+    assign = conwrap(Dataset.assign)
+
+    assign_attrs = conwrap(Dataset.assign_attrs)
+
+    assign_coords = conwrap(Dataset.assign_coords)
+
+    # bfill = conwrap(Dataset.bfill)
+
+    broadcast_like = conwrap(Dataset.broadcast_like)
+
+    chunk = conwrap(Dataset.chunk)
+
+    drop_sel = conwrap(Dataset.drop_sel)
+
+    drop_isel = conwrap(Dataset.drop_isel)
+
+    expand_dims = conwrap(Dataset.expand_dims)
+
+    # ffill = conwrap(Dataset.ffill)
+
     sel = conwrap(Dataset.sel)
 
     isel = conwrap(Dataset.isel)
 
+    shift = conwrap(Dataset.shift)
+
+    swap_dims = conwrap(Dataset.swap_dims)
+
+    set_index = conwrap(Dataset.set_index)
+
+    reindex = conwrap(Dataset.reindex, fill_value=_fill_value)
+
+    reindex_like = conwrap(Dataset.reindex_like, fill_value=_fill_value)
+
+    rename = conwrap(Dataset.rename)
+
+    rename_dims = conwrap(Dataset.rename_dims)
+
+    roll = conwrap(Dataset.roll)
+
+    stack = conwrap(Dataset.stack)
+
 
 @dataclass(repr=False)
 class Constraints:

linopy/expressions.py

@@ -1317,6 +1317,10 @@ def to_polars(self) -> pl.DataFrame:
 
     shift = exprwrap(Dataset.shift)
 
+    swap_dims = exprwrap(Dataset.swap_dims)
+
+    set_index = exprwrap(Dataset.set_index)
+
     reindex = exprwrap(Dataset.reindex, fill_value=_fill_value)
 
     reindex_like = exprwrap(Dataset.reindex_like, fill_value=_fill_value)

linopy/variables.py

@@ -1006,6 +1006,10 @@ def equals(self, other):
 
     shift = varwrap(Dataset.shift, fill_value=_fill_value)
 
+    swap_dims = varwrap(Dataset.swap_dims)
+
+    set_index = varwrap(Dataset.set_index)
+
     rename = varwrap(Dataset.rename)
 
     roll = varwrap(Dataset.roll)

test/test_constraint.py

@@ -176,12 +176,59 @@ def test_constraint_inherited_properties(x, y):
     assert con.values is None
 
 
+def test_constraint_wrapped_methods(x, y):
+    con = 10 * x + y <= 10
+
+    # Test wrapped methods
+    con.assign({"new_var": xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])})
+    con.assign_attrs({"new_attr": "value"})
+    con.assign_coords(
+        {"new_coord": xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])}
+    )
+    # con.bfill(dim="first")
+    con.broadcast_like(con.data)
+    con.chunk()
+    con.drop_sel({"first": 0})
+    con.drop_isel({"first": 0})
+    con.expand_dims("new_dim")
+    # con.ffill(dim="first")
+    con.shift({"first": 1})
+    con.reindex({"first": [0, 1]})
+    con.reindex_like(con.data)
+    con.rename({"first": "new_labels"})
+    con.rename_dims({"first": "new_labels"})
+    con.roll({"first": 1})
+    con.stack(new_dim=("first", "second"))
+
+
 def test_anonymous_constraint_sel(x, y):
     expr = 10 * x + y
     con = expr <= 10
     assert isinstance(con.sel(first=[1, 2]), Constraint)
 
 
+def test_anonymous_constraint_swap_dims(x, y):
+    expr = 10 * x + y
+    con = expr <= 10
+    con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)})
+    con = con.swap_dims({"second": "third"})
+    assert isinstance(con, Constraint)
+    assert con.coord_dims == ("first", "third")
+
+
+def test_anonymous_constraint_set_index(x, y):
+    expr = 10 * x + y
+    con = expr <= 10
+    con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)})
+    con = con.set_index({"multi": ["second", "third"]})
+    assert isinstance(con, Constraint)
+    assert con.coord_dims == (
+        "first",
+        "multi",
+    )
+    assert isinstance(con.indexes["multi"], pd.MultiIndex)
+
+
 def test_anonymous_constraint_loc(x, y):
     expr = 10 * x + y
     con = expr <= 10

test/test_linear_expression.py

@@ -608,6 +608,23 @@ def test_linear_expression_shift(v):
     assert (shifted.vars.loc[:1] == -1).all()
 
 
+def test_linear_expression_swap_dims(v):
+    expr = v.to_linexpr()
+    expr = expr.assign_coords({"second": ("dim_2", expr.indexes["dim_2"] + 100)})
+    expr = expr.swap_dims({"dim_2": "second"})
+    assert isinstance(expr, LinearExpression)
+    assert expr.coord_dims == ("second",)
+
+
+def test_linear_expression_set_index(v):
+    expr = v.to_linexpr()
+    expr = expr.assign_coords({"second": ("dim_2", expr.indexes["dim_2"] + 100)})
+    expr = expr.set_index({"multi": ["dim_2", "second"]})
+    assert isinstance(expr, LinearExpression)
+    assert expr.coord_dims == ("multi",)
+    assert isinstance(expr.indexes["multi"], pd.MultiIndex)
+
+
 def test_linear_expression_fillna(v):
     expr = np.arange(20) * v + 10
     assert expr.const.sum() == 200

test/test_variable.py

@@ -204,6 +204,21 @@ def test_variable_shift(x):
     assert x.labels[0] == -1
 
 
+def test_variable_swap_dims(x):
+    x = x.assign_coords({"second": ("first", x.indexes["first"] + 100)})
+    x = x.swap_dims({"first": "second"})
+    assert isinstance(x, linopy.variables.Variable)
+    assert x.dims == ("second",)
+
+
+def test_variable_set_index(x):
+    x = x.assign_coords({"second": ("first", x.indexes["first"] + 100)})
+    x = x.set_index({"multi": ["first", "second"]})
+    assert isinstance(x, linopy.variables.Variable)
+    assert x.dims == ("multi",)
+    assert isinstance(x.indexes["multi"], pd.MultiIndex)
+
+
 def test_isnull(x):
     x = x.where([True] * 4 + [False] * 6)
     assert isinstance(x.isnull(), xr.DataArray)