Refactor LinearExpression and Constraint using composition

doc/release_notes.rst

@@ -1,8 +1,11 @@
 Release Notes
 =============
 
-.. Upcoming Release
-.. ----------------
+Upcoming Release
+----------------
+
+* The internal data structure of linopy classes were updated to a safer design. Instead of being defined as inherited xarray classes, the class `Variable`, `LinearExpression` and `Constraint` are now dataclasses with containing the xarray objects in the data field. This allows the package to have more flexible function design and a reduced set of wrapped functions that are sensible to use in the optimization context.
+* The class `Variable` and `LinearExpression` have new functions `groupby` and `rolling` imitating the corresponding xarray functions but with safe type inheritance and application of appended operations.
 
 
 Version 0.0.15

linopy/__init__.py

@@ -6,6 +6,9 @@
 @author: fabulous
 """
 
+# Note: For intercepting multiplications between xarray dataarrays, Variables and Expressions
+# we need to extend their __mul__ functions with a quick special case
+import linopy.monkey_patch_xarray
 from linopy import model, remote
 from linopy.expressions import merge
 from linopy.io import read_netcdf

linopy/common.py

@@ -6,7 +6,7 @@
 This module contains commonly used functions.
 """
 
-from functools import wraps
+from functools import partialmethod, update_wrapper, wraps
 
 import numpy as np
 from xarray import DataArray, apply_ufunc, merge
@@ -98,3 +98,20 @@ def wrapper(self, arg):
         return func(self, arg)
 
     return wrapper
+
+
+def forward_as_properties(**routes):
+    def add_accessor(cls, item, attr):
+        @property
+        def get(self):
+            return getattr(getattr(self, item), attr)
+
+        setattr(cls, attr, get)
+
+    def deco(cls):
+        for item, attrs in routes.items():
+            for attr in attrs:
+                add_accessor(cls, item, attr)
+        return cls
+
+    return deco

linopy/constraints.py

@@ -14,58 +14,37 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
+from deprecation import deprecated
 from numpy import arange, array
 from scipy.sparse import coo_matrix
 from xarray import DataArray, Dataset
 
 from linopy import expressions, variables
 from linopy.common import (
     _merge_inplace,
-    has_assigned_model,
     has_optimized_model,
     is_constant,
     replace_by_map,
 )
 
 
-class Constraint(DataArray):
+@dataclass(repr=False)
+class Constraint:
     """
-    Constraint container for storing constraint labels.
+    Projection to a single constraint in a model.
 
     The Constraint class is a subclass of xr.DataArray hence most xarray
     functions can be applied to it.
     """
 
-    __slots__ = ("_cache", "_coords", "_indexes", "_name", "_variable", "model")
-
-    def __init__(self, *args, **kwargs):
-
-        # workaround until https://github.com/pydata/xarray/pull/5984 is merged
-        if isinstance(args[0], DataArray):
-            da = args[0]
-            args = (da.data, da.coords)
-            kwargs.update({"attrs": da.attrs, "name": da.name})
-
-        self.model = kwargs.pop("model", None)
-        super().__init__(*args, **kwargs)
-        assert self.name is not None, "Constraint data does not have a name."
-
-    # We have to set the _reduce_method to None, in order to overwrite basic
-    # reduction functions as `sum`. There might be a better solution (?).
-    _reduce_method = None
-
-    # Disable array function, only function defined below are supported
-    # and set priority higher than pandas/xarray/numpy
-    __array_ufunc__ = None
-    __array_priority__ = 10000
+    name: str
+    model: Any
 
     def __repr__(self):
         """
         Get the string representation of the constraints.
         """
-        data_string = (
-            "Constraint labels:\n" + self.to_array().__repr__().split("\n", 1)[1]
-        )
+        data_string = "Constraint labels:\n" + self.labels.__repr__().split("\n", 1)[1]
         extend_line = "-" * len(self.name)
         return (
             f"Constraint '{self.name}':\n"
@@ -78,18 +57,26 @@ def _repr_html_(self):
         Get the html representation of the variables.
         """
         # return self.__repr__()
-        data_string = self.to_array()._repr_html_()
+        data_string = self.labels._repr_html_()
         data_string = data_string.replace("xarray.DataArray", "linopy.Constraint")
         return data_string
 
+    @deprecated(details="Use the `labels` property instead of `to_array`")
     def to_array(self):
         """
         Convert the variable array to a xarray.DataArray.
         """
-        return DataArray(self)
+        return self.labels
+
+    @property
+    def labels(self):
+        return self.model.constraints.labels[self.name]
+
+    @labels.setter
+    def labels(self, value):
+        raise RuntimeError("Labels are read-only")
 
     @property
-    @has_assigned_model
     def coeffs(self):
         """
         Get the left-hand-side coefficients of the constraint.
@@ -100,7 +87,6 @@ def coeffs(self):
         return self.model.constraints.coeffs[self.name]
 
     @coeffs.setter
-    @has_assigned_model
     def coeffs(self, value):
         term_dim = self.name + "_term"
         value = DataArray(value).broadcast_like(self.vars, exclude=[term_dim])
@@ -115,7 +101,6 @@ def coeffs(self, value):
             self.model.constraints.coeffs = coeffs
 
     @property
-    @has_assigned_model
     def vars(self):
         """
         Get the left-hand-side variables of the constraint.
@@ -126,7 +111,6 @@ def vars(self):
         return self.model.constraints.vars[self.name]
 
     @vars.setter
-    @has_assigned_model
     def vars(self, value):
         term_dim = self.name + "_term"
         value = DataArray(value).broadcast_like(self.coeffs, exclude=[term_dim])
@@ -141,7 +125,6 @@ def vars(self, value):
             self.model.constraints.vars = vars
 
     @property
-    @has_assigned_model
     def lhs(self):
         """
         Get the left-hand-side linear expression of the constraint.
@@ -155,7 +138,6 @@ def lhs(self):
         return expressions.LinearExpression(Dataset({"coeffs": coeffs, "vars": vars}))
 
     @lhs.setter
-    @has_assigned_model
     def lhs(self, value):
         if not isinstance(value, expressions.LinearExpression):
             raise TypeError("Assigned lhs must be a LinearExpression.")
@@ -165,7 +147,6 @@ def lhs(self, value):
         self.vars = value.vars
 
     @property
-    @has_assigned_model
     def sign(self):
         """
         Get the signs of the constraint.
@@ -176,16 +157,14 @@ def sign(self):
         return self.model.constraints.sign[self.name]
 
     @sign.setter
-    @has_assigned_model
     @is_constant
     def sign(self, value):
-        value = DataArray(value).broadcast_like(self)
+        value = DataArray(value).broadcast_like(self.sign)
         if (value == "==").any():
             raise ValueError('Sign "==" not supported, use "=" instead.')
         self.model.constraints.sign[self.name] = value
 
     @property
-    @has_assigned_model
     def rhs(self):
         """
         Get the right hand side constants of the constraint.
@@ -196,12 +175,11 @@ def rhs(self):
         return self.model.constraints.rhs[self.name]
 
     @rhs.setter
-    @has_assigned_model
     @is_constant
     def rhs(self, value):
         if isinstance(value, (variables.Variable, expressions.LinearExpression)):
             raise TypeError(f"Assigned rhs must be a constant, got {type(value)}).")
-        value = DataArray(value).broadcast_like(self)
+        value = DataArray(value).broadcast_like(self.rhs)
         self.model.constraints.rhs[self.name] = value
 
     @property
@@ -219,6 +197,10 @@ def dual(self):
             )
         return self.model.dual[self.name]
 
+    @property
+    def shape(self):
+        return self.labels.shape
+
 
 @dataclass(repr=False)
 class Constraints:
@@ -268,7 +250,7 @@ def __getitem__(
         self, names: Union[str, Sequence[str]]
     ) -> Union[Constraint, "Constraints"]:
         if isinstance(names, str):
-            return Constraint(self.labels[names], model=self.model)
+            return Constraint(names, model=self.model)
 
         return self.__class__(
             self.labels[names],
@@ -581,7 +563,8 @@ def __init__(self, lhs, sign, rhs):
         """
         if isinstance(rhs, (variables.Variable, expressions.LinearExpression)):
             raise TypeError(f"Assigned rhs must be a constant, got {type(rhs)}).")
-        self._lhs, self._rhs = xr.align(lhs, DataArray(rhs))
+        lhs, self._rhs = xr.align(lhs.data, DataArray(rhs))
+        self._lhs = expressions.LinearExpression(lhs)
         self._sign = DataArray(sign)
 
     @property
@@ -696,7 +679,7 @@ class AnonymousScalarConstraint:
     (rhs) for exactly one constraint.
     """
 
-    lhs: expressions.ScalarLinearExpression
+    lhs: "expressions.ScalarLinearExpression"
     sign: str
     rhs: float