Repository: PyPSA/linopy
Pull request #32 (2022-07-07 13:32:45)

Enable rule based assignment of expression and constraints as done by pyomo
This PR will help reusing functions created for pyomo workflows.


Patch:
diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml
index f7f6ec0..25c3198 100644
--- a/.github/workflows/CI.yaml
+++ b/.github/workflows/CI.yaml
@@ -59,7 +59,7 @@ jobs:
 
     - name: Test with pytest
       run: |
-        pytest --cov=./ --cov-report=xml
+        pytest --cov=./ --cov-report=xml linopy --doctest-modules test
 
     - name: Upload code coverage report
       uses: codecov/codecov-action@v2
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a147221..ef01ae3 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -33,12 +33,8 @@ repos:
     -   id: codespell
         types_or: [python, rst, markdown]
         files: ^(linopy|doc)/
--   repo: local
-    hooks:
-    -   id: jupyter-nb-clear-output
-        name: jupyter-nb-clear-output
-        files: \.ipynb$
-        exclude: solve-on-remote.ipynb
-        stages: [commit]
-        language: system
-        entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
+-   repo: https://github.com/aflc/pre-commit-jupyter
+    rev: v1.2.1
+    hooks:
+    -   id: jupyter-notebook-cleanup
+        exclude: examples/solve-on-remote.ipynb
diff --git a/benchmark/scripts/merge-benchmarks.py b/benchmark/scripts/merge-benchmarks.py
index c70c8e7..7fdcde1 100644
--- a/benchmark/scripts/merge-benchmarks.py
+++ b/benchmark/scripts/merge-benchmarks.py
@@ -19,8 +19,8 @@ for api in df_time.API.unique():
     df.loc[df.API == api, "Time"] = df_time.loc[df.API == api, "Time"].values
 
 
-df["Number of Variables"] = df.N ** 2 * 2
-df["Number of Constraints"] = df.N ** 2 * 2
+df["Number of Variables"] = df.N**2 * 2
+df["Number of Constraints"] = df.N**2 * 2
 
 
 # Make a correction of the memory usage:
diff --git a/benchmark/scripts/plot-benchmarks.py b/benchmark/scripts/plot-benchmarks.py
index bbbb336..7dcd320 100644
--- a/benchmark/scripts/plot-benchmarks.py
+++ b/benchmark/scripts/plot-benchmarks.py
@@ -13,7 +13,7 @@ import seaborn as sns
 df = [pd.read_csv(fn) for fn in snakemake.input]
 df = pd.concat(df, ignore_index=True)
 
-df["# Variables"] = df.N ** 2 * 2
+df["# Variables"] = df.N**2 * 2
 
 fig, ax = plt.subplots()
 sns.lineplot(x="# Variables", y="Time [s]", hue="API", data=df, ax=ax)
diff --git a/doc/logo.py b/doc/logo.py
index 6540bdb..7c3e7ed 100644
--- a/doc/logo.py
+++ b/doc/logo.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Created on Sun Oct 17 12:08:47 2021
+Created on Sun Oct 17 12:08:47 2021.
 
 @author: fabian
 """
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index 02e274f..ed12f5b 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -4,6 +4,11 @@ Release Notes
 Upcoming Release
 ----------------
 
+* Constraints and expressions can now be created using function that iterate over all combinations of given coordinates. This functionality mirrors the behavior of the Pyomo package. For complicated constraints which are hard to create with arrays of variables, it is easier (thus less efficient) to use an iterating function. For more information see the example notebook in the documentation.
+* When getting the value of a variable, the value of the variable is returned as a `ScalarVariable`. This is useful for the above mentioned creation of expressions and constraints with iterating functions. This affect only the direct getter function, all other functions like `.sel` or `.isel` behave as known from Xarray.
+* The docstring examples are now part of the Continuous Integration.
+
+
 Version 0.0.10
 --------------
 
diff --git a/linopy/constraints.py b/linopy/constraints.py
index 8693920..bd8d233 100644
--- a/linopy/constraints.py
+++ b/linopy/constraints.py
@@ -7,17 +7,19 @@ This module contains implementations for the Constraint{s} class.
 
 import re
 from dataclasses import dataclass
+from itertools import product
 from typing import Any, Sequence, Union
 
 import dask
 import numpy as np
 import pandas as pd
 import xarray as xr
+from numpy import array
 from scipy.sparse import coo_matrix
 from xarray import DataArray, Dataset
 
+from linopy import expressions
 from linopy.common import _merge_inplace, replace_by_map
-from linopy.expressions import LinearExpression
 
 
 class Constraint(DataArray):
@@ -146,11 +148,11 @@ class Constraint(DataArray):
         """
         coeffs = self.coeffs.rename({self.name + "_term": "_term"})
         vars = self.vars.rename({self.name + "_term": "_term"})
-        return LinearExpression(Dataset({"coeffs": coeffs, "vars": vars}))
+        return expressions.LinearExpression(Dataset({"coeffs": coeffs, "vars": vars}))
 
     @lhs.setter
     def lhs(self, value):
-        if not isinstance(value, LinearExpression):
+        if not isinstance(value, expressions.LinearExpression):
             raise TypeError("Assigned lhs must be a LinearExpression.")
 
         value = value.rename(_term=self.name + "_term")
@@ -550,3 +552,120 @@ class Constraints:
             shape = (ncons, nvars)  # same as model.nvars/ncons but already there
 
         return coo_matrix((data, (rows, cols)), shape=shape)
+
+
+class AnonymousConstraint:
+    """
+    A constraint container used for storing multiple constraint arrays.
+    """
+
+    __slots__ = ("lhs", "sign", "rhs")
+
+    def __init__(self, lhs, sign, rhs):
+        """
+        Initialize a anonymous constraint.
+        """
+        self.lhs, self.rhs = xr.align(lhs, DataArray(rhs))
+        self.sign = DataArray(sign)
+
+    def __repr__(self):
+        """
+        Get the string representation of the expression.
+        """
+        lhs_string = self.lhs.to_dataset().__repr__()  # .split("\n", 1)[1]
+        lhs_string = lhs_string.split("Data variables:\n", 1)[1]
+        lhs_string = lhs_string.replace("    coeffs", "coeffs")
+        lhs_string = lhs_string.replace("    vars", "vars")
+        if self.sign.size == 1:
+            sign_string = self.sign.item()
+        else:
+            sign_string = self.sign.__repr__().split("\n", 1)[1]
+        if self.rhs.size == 1:
+            rhs_string = self.rhs.item()
+        else:
+            rhs_string = self.rhs.__repr__().split("\n", 1)[1]
+        return (
+            f"Anonymous Constraint:\n"
+            f"---------------------\n"
+            f"\n{lhs_string}"
+            f"\n{sign_string}"
+            f"\n{rhs_string}"
+        )
+
+    def from_rule(model, rule, coords):
+        """
+        Create a constraint from a rule and a set of coordinates.
+
+        This functionality mirrors the assignment of constraints as done by
+        Pyomo.
+
+
+        Parameters
+        ----------
+        model : linopy.Model
+            Passed to function `rule` as a first argument.
+        rule : callable
+            Function to be called for each combinations in `coords`.
+            The first argument of the function is the underlying `linopy.Model`.
+            The following arguments are given by the coordinates for accessing
+            the variables. The function has to return a
+            `AnonymousScalarConstraint`. Therefore use the direct getter when
+            indexing variables in the linear expression.
+        coords : coordinate-like
+            Coordinates to processed by `xarray.DataArray`.
+            For each combination of coordinates, the function given by `rule` is called.
+            The order and size of coords has to be same as the argument list
+            followed by `model` in function `rule`.
+
+
+        Returns
+        -------
+        linopy.AnonymousConstraint
+
+        Examples
+        --------
+        >>> from linopy import Model, LinearExpression
+        >>> m = Model()
+        >>> coords = pd.RangeIndex(10), ["a", "b"]
+        >>> x = m.add_variables(0, 100, coords)
+        >>> def bound(m, i, j):
+        ...     if i % 2:
+        ...         return (i - 1) * x[i - 1, j] >= 0
+        ...     else:
+        ...         return i * x[i, j] >= 0
+        ...
+        >>> con = AnonymousConstraint.from_rule(m, bound, coords)
+        >>> con = m.add_constraints(con)
+        """
+        if not isinstance(coords, xr.core.dataarray.DataArrayCoordinates):
+            coords = DataArray(coords=coords).coords
+        shape = list(map(len, coords.values()))
+
+        # test output type
+        output = rule(model, *[c.values[0] for c in coords.values()])
+        if not isinstance(output, AnonymousScalarConstraint):
+            msg = f"`rule` has to return AnonymousScalarConstraint not {type(output)}."
+            raise TypeError(msg)
+
+        combinations = product(*[c.values for c in coords.values()])
+        cons = [rule(model, *coord) for coord in combinations]
+        exprs = [con.lhs for con in cons]
+
+        lhs = expressions.LinearExpression._from_scalarexpression_list(exprs, coords)
+        sign = DataArray(array([c.sign for c in cons]).reshape(shape), coords)
+        rhs = DataArray(array([c.rhs for c in cons]).reshape(shape), coords)
+        return AnonymousConstraint(lhs, sign, rhs)
+
+
+@dataclass
+class AnonymousScalarConstraint:
+    """
+    Container for anonymous scalar constraint.
+
+    This contains a left-hand-side (lhs), a sign and a right-hand-side
+    (rhs) for exactly one constraint.
+    """
+
+    lhs: expressions.ScalarLinearExpression
+    sign: str
+    rhs: float
diff --git a/linopy/expressions.py b/linopy/expressions.py
index 9a2e353..0984dd4 100644
--- a/linopy/expressions.py
+++ b/linopy/expressions.py
@@ -8,15 +8,19 @@ This module contains definition related to affine expressions.
 
 import functools
 import logging
+from dataclasses import dataclass
+from itertools import product, zip_longest
 from warnings import warn
 
 import numpy as np
 import pandas as pd
 import xarray as xr
+from numpy import array, nan
 from xarray import DataArray, Dataset
+from xarray.core.dataarray import DataArrayCoordinates
 from xarray.core.groupby import _maybe_reorder, peek_at
 
-from linopy import variables
+from linopy import constraints, variables
 from linopy.common import as_dataarray
 
 
@@ -48,6 +52,8 @@ class LinearExpression(Dataset):
 
     Examples
     --------
+    >>> from linopy import Model
+    >>> import pandas as pd
     >>> m = Model()
     >>> x = m.add_variables(pd.Series([0, 0]), 1, name="x")
     >>> y = m.add_variables(4, pd.Series([8, 10]), name="y")
@@ -56,32 +62,21 @@ class LinearExpression(Dataset):
 
     >>> expr = 3 * x
     >>> type(expr)
-    linopy.model.LinearExpression
+    <class 'linopy.expressions.LinearExpression'>
 
     >>> other = 4 * y
     >>> type(expr + other)
-    linopy.model.LinearExpression
+    <class 'linopy.expressions.LinearExpression'>
 
     Multiplying:
 
     >>> type(3 * expr)
-    linopy.model.LinearExpression
+    <class 'linopy.expressions.LinearExpression'>
 
     Summation over dimensions
 
-    >>> expr.sum(dim="dim_0")
-
-    ::
-
-        Linear Expression with 2 term(s):
-        ----------------------------------
-
-        Dimensions:  (_term: 2)
-        Coordinates:
-            * _term    (_term) int64 0 1
-        Data:
-            coeffs   (_term) int64 3 3
-            vars     (_term) int64 1 2
+    >>> type(expr.sum(dims="dim_0"))
+    <class 'linopy.expressions.LinearExpression'>
     """
 
     __slots__ = ("_cache", "_coords", "_indexes", "_name", "_variable")
@@ -187,13 +182,13 @@ class LinearExpression(Dataset):
         return self.__mul__(other)
 
     def __le__(self, rhs):
-        return AnonymousConstraint(self, "<=", rhs)
+        return constraints.AnonymousConstraint(self, "<=", rhs)
 
     def __ge__(self, rhs):
-        return AnonymousConstraint(self, ">=", rhs)
+        return constraints.AnonymousConstraint(self, ">=", rhs)
 
     def __eq__(self, rhs):
-        return AnonymousConstraint(self, "=", rhs)
+        return constraints.AnonymousConstraint(self, "=", rhs)
 
     def to_dataset(self):
         """
@@ -265,9 +260,11 @@ class LinearExpression(Dataset):
 
         Examples
         --------
+        >>> from linopy import Model
+        >>> import pandas as pd
         >>> m = Model()
         >>> x = m.add_variables(pd.Series([0, 0]), 1)
-        >>> m.add_variables(4, pd.Series([8, 10]))
+        >>> y = m.add_variables(4, pd.Series([8, 10]))
         >>> expr = LinearExpression.from_tuples((10, x), (1, y))
 
         This is the same as calling ``10*x + y`` but a bit more performant.
@@ -295,6 +292,84 @@ class LinearExpression(Dataset):
             ds = ds_list[0].expand_dims("_term")
         return LinearExpression(ds)
 
+    def from_rule(model, rule, coords):
+        """
+        Create a linear expression from a rule and a set of coordinates.
+
+        This functionality mirrors the assignment of linear expression as done by
+        Pyomo.
+
+
+        Parameters
+        ----------
+        model : linopy.Model
+            Passed to function `rule` as a first argument.
+        rule : callable
+            Function to be called for each combinations in `coords`.
+            The first argument of the function is the underlying `linopy.Model`.
+            The following arguments are given by the coordinates for accessing
+            the variables. The function has to return a
+            `ScalarLinearExpression`. Therefore use the `.at` accessor when
+            indexing variables.
+        coords : coordinate-like
+            Coordinates to processed by `xarray.DataArray`.
+            For each combination of coordinates, the function
+            given by `rule` is called. The order and size of coords has
+            to be same as the argument list followed by `model` in
+            function `rule`.
+
+
+        Returns
+        -------
+        linopy.LinearExpression
+
+        Examples
+        --------
+        >>> from linopy import Model, LinearExpression
+        >>> m = Model()
+        >>> coords = pd.RangeIndex(10), ["a", "b"]
+        >>> x = m.add_variables(0, 100, coords)
+        >>> def bound(m, i, j):
+        ...     if i % 2:
+        ...         return (i - 1) * x[i - 1, j]
+        ...     else:
+        ...         return i * x[i, j]
+        ...
+        >>> expr = LinearExpression.from_rule(m, bound, coords)
+        >>> con = m.add_constraints(expr <= 10)
+        """
+        if not isinstance(coords, xr.core.dataarray.DataArrayCoordinates):
+            coords = DataArray(coords=coords).coords
+
+        # test output type
+        output = rule(model, *[c.values[0] for c in coords.values()])
+        if not isinstance(output, ScalarLinearExpression):
+            msg = f"`rule` has to return ScalarLinearExpression not {type(output)}."
+            raise TypeError(msg)
+
+        combinations = product(*[c.values for c in coords.values()])
+        exprs = [rule(model, *coord) for coord in combinations]
+        return LinearExpression._from_scalarexpression_list(exprs, coords)
+
+    def _from_scalarexpression_list(exprs, coords: DataArrayCoordinates):
+        """
+        Create a LinearExpression from a list of lists with different lengths.
+        """
+        shape = list(map(len, coords.values()))
+
+        coeffs = array(tuple(zip_longest(*(e.coeffs for e in exprs), fillvalue=nan)))
+        vars = array(tuple(zip_longest(*(e.vars for e in exprs), fillvalue=-1)))
+
+        nterm = vars.shape[0]
+        coeffs = coeffs.reshape((nterm, *shape))
+        vars = vars.reshape((nterm, *shape))
+
+        coeffs = DataArray(coeffs, coords, dims=("_term", *coords))
+        vars = DataArray(vars, coords, dims=("_term", *coords))
+        ds = Dataset({"coeffs": coeffs, "vars": vars}).transpose(..., "_term")
+
+        return LinearExpression(ds)
+
     def where(self, cond, **kwargs):
         """
         Filter variables based on a condition.
@@ -379,7 +454,9 @@ class LinearExpression(Dataset):
         )
 
         vars = xr.DataArray.rolling(self.vars, **kwargs).construct(
-            "_rolling_term", fill_value=self.fill_value["vars"], keep_attrs=True
+            "_rolling_term",
+            fill_value=self.fill_value["vars"],
+            keep_attrs=True,
         )
 
         ds = xr.Dataset({"coeffs": coeffs, "vars": vars})
@@ -558,36 +635,79 @@ def merge(*exprs, dim="_term"):
     return res
 
 
-class AnonymousConstraint:
-    """
-    A constraint container used for storing multiple constraint arrays.
-    """
+@dataclass
+class ScalarLinearExpression:
+    coeffs: tuple
+    vars: tuple
+    coords: dict = None
 
-    __slots__ = ("lhs", "sign", "rhs")
+    def __add__(self, other):
+        if isinstance(other, variables.ScalarVariable):
+            coeffs = self.coeffs + (1,)
+            vars = self.vars + (other.label,)
+            return ScalarLinearExpression(coeffs, vars)
+        elif not isinstance(other, ScalarLinearExpression):
+            raise TypeError(
+                "unsupported operand type(s) for +: " f"{type(self)} and {type(other)}"
+            )
 
-    def __init__(self, lhs, sign, rhs):
-        """
-        Initialize a anonymous constraint.
-        """
-        self.lhs, self.rhs = xr.align(lhs, DataArray(rhs))
-        self.sign = DataArray(sign)
+        coeffs = self.coeffs + other.coeffs
+        vars = self.vars + other.vars
+        return ScalarLinearExpression(coeffs, vars)
 
-    def __repr__(self):
-        """
-        Get the string representation of the expression.
-        """
-        lhs_string = self.lhs.to_dataset().__repr__()  # .split("\n", 1)[1]
-        lhs_string = lhs_string.split("Data variables:\n", 1)[1]
-        lhs_string = lhs_string.replace("    coeffs", "coeffs")
-        lhs_string = lhs_string.replace("    vars", "vars")
-        if self.rhs.size == 1:
-            rhs_string = self.rhs.item()
-        else:
-            rhs_string = self.rhs.__repr__().split("\n", 1)[1]
-        return (
-            f"Anonymous Constraint:\n"
-            f"---------------------\n"
-            f"\n{lhs_string}"
-            f"\n{self.sign.item()}"
-            f"\n{rhs_string}"
+    @property
+    def nterm(self):
+        return len(self.vars)
+
+    def __sub__(self, other):
+        if isinstance(other, variables.ScalarVariable):
+            other = other.to_linexpr(1)
+        elif not isinstance(other, ScalarLinearExpression):
+            raise TypeError(
+                "unsupported operand type(s) for -: " f"{type(self)} and {type(other)}"
+            )
+
+        return ScalarLinearExpression(
+            self.coeffs + tuple(-c for c in other.coeffs), self.vars + other.vars
         )
+
+    def __neg__(self):
+        return ScalarLinearExpression(tuple(-c for c in self.coeffs), self.vars)
+
+    def __mul__(self, other):
+        if not isinstance(other, (float, int)):
+            raise TypeError(
+                "unsupported operand type(s) for *: " f"{type(self)} and {type(other)}"
+            )
+
+        return ScalarLinearExpression(tuple(other * c for c in self.coeffs), self.vars)
+
+    def __rmul__(self, other):
+        return self.__mul__(other)
+
+    def __div__(self, other):
+        return self.__mul__(1 / other)
+
+    def __le__(self, other):
+        if not isinstance(other, (int, float)):
+            raise TypeError(
+                "unsupported operand type(s) for >=: " f"{type(self)} and {type(other)}"
+            )
+
+        return constraints.AnonymousScalarConstraint(self, "<=", other)
+
+    def __ge__(self, other):
+        if not isinstance(other, (int, float)):
+            raise TypeError(
+                "unsupported operand type(s) for >=: " f"{type(self)} and {type(other)}"
+            )
+
+        return constraints.AnonymousScalarConstraint(self, ">=", other)
+
+    def __eq__(self, other):
+        if not isinstance(other, (int, float)):
+            raise TypeError(
+                "unsupported operand type(s) for ==: " f"{type(self)} and {type(other)}"
+            )
+
+        return constraints.AnonymousScalarConstraint(self, "==", other)
diff --git a/linopy/model.py b/linopy/model.py
index 0e2ac36..82caaa7 100644
--- a/linopy/model.py
+++ b/linopy/model.py
@@ -19,9 +19,9 @@ from xarray import DataArray, Dataset
 
 from linopy import solvers
 from linopy.common import best_int, replace_by_map
-from linopy.constraints import Constraints
+from linopy.constraints import AnonymousConstraint, Constraints
 from linopy.eval import Expr
-from linopy.expressions import AnonymousConstraint, LinearExpression
+from linopy.expressions import LinearExpression
 from linopy.io import to_block_files, to_file, to_netcdf
 from linopy.solvers import available_solvers
 from linopy.variables import Variable, Variables
@@ -380,22 +380,20 @@ class Model:
 
         Examples
         --------
-
-        >>> m = linopy.Model()
+        >>> from linopy import Model
+        >>> import pandas as pd
+        >>> m = Model()
         >>> time = pd.RangeIndex(10, name="Time")
         >>> m.add_variables(lower=0, coords=[time], name="x")
-
-        ::
-
-            Variable container:
-            -------------------
-
-            Variables:
-            array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])
-            Coordinates:
-                * Time     (Time) int64 0 1 2 3 4 5 6 7 8 9
-            Attributes:
-                name:     x
+        Variable 'x':
+        -------------
+        <BLANKLINE>
+        Variable labels:
+        array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        Coordinates:
+          * Time     (Time) int64 0 1 2 3 4 5 6 7 8 9
+        Attributes:
+            binary:   False
         """
         if name is None:
             name = "var" + str(self._varnameCounter)
@@ -445,7 +443,9 @@ class Model:
 
         return self.variables[name]
 
-    def add_constraints(self, lhs, sign=None, rhs=None, name=None, mask=None):
+    def add_constraints(
+        self, lhs, sign=None, rhs=None, name=None, coords=None, mask=None
+    ):
         """
         Assign a new, possibly multi-dimensional array of constraints to the
         model.
@@ -457,10 +457,14 @@ class Model:
 
         Parameters
         ----------
-        lhs : linopy.LinearExpression/linopy.AnonynousConstraint
+        lhs : linopy.LinearExpression/linopy.AnonymousConstraint/callable
             Left hand side of the constraint(s) or optionally full constraint.
             In case a linear expression is passed, `sign` and `rhs` must not be
             None.
+            If a function is passed, it is called for every combination of
+            coordinates given in `coords`. It's first argument has to be the
+            model, followed by scalar argument for each coordinate given in
+            coordinates.
         sign : str/array_like
             Relation between the lhs and rhs, valid values are {'=', '>=', '<='}.
         rhs : int/float/array_like
@@ -468,6 +472,9 @@ class Model:
         name : str, optional
             Reference name of the added constraints. The default None results
             results a name like "con1", "con2" etc.
+        coords : list/xarray.Coordinates, optional
+            The coords of the constraint array. This is only used when lhs is
+            a function. The default is None.
         mask : array_like, optional
             Boolean mask with False values for variables which are skipped.
             The shape of the mask has to match the shape the added variables.
@@ -486,6 +493,11 @@ class Model:
         if name in self.constraints:
             raise ValueError(f"Constraint '{name}' already assigned to model")
 
+        if callable(lhs):
+            assert coords is not None, "`coords` must be given when lhs is a function"
+            rule = lhs
+            lhs = AnonymousConstraint.from_rule(self, rule, coords)
+
         if isinstance(lhs, AnonymousConstraint):
             if sign is not None or rhs is not None:
                 raise ValueError(
@@ -684,18 +696,18 @@ class Model:
         blocks = replace_by_map(self.objective.vars, block_map)
         self.objective = self.objective.assign(blocks=blocks)
 
-    def linexpr(self, *tuples):
+    def linexpr(self, *args):
         """
-        Create a linopy.LinearExpression by using variable names.
-
-        Calls the function LinearExpression.from_tuples but loads variables from
-        the model if a variable name is used.
+        Create a linopy.LinearExpression from argument list.
 
         Parameters
         ----------
-        tuples : tuples of (coefficients, variables)
-            Each tuple represents on term in the linear expression, which can
-            span over multiple dimensions:
+        args : tuples of (coefficients, variables) or tuples of
+               coordinates and a function
+            If *args is a collection of coefficients-variables-tuples, the resulting
+            linear expression is built with the function LinearExpression.from_tuples.
+            In this case, each tuple represents on term in the linear expression,
+            which can span over multiple dimensions:
 
             * coefficients : int/float/array_like
                 The coefficient(s) in the term, if the coefficients array
@@ -705,6 +717,24 @@ class Model:
                 The variable(s) going into the term. These may be referenced
                 by name.
 
+            If *args is a collection of coordinates with an appended function at the
+            end, the function LinearExpression.from_rule is used to build the linear
+            expression. Then, the argument are expected to contain:
+
+            * rule : callable
+                Function to be called for each combinations in `coords`.
+                The first argument of the function is the underlying `linopy.Model`.
+                The following arguments are given by the coordinates for accessing
+                the variables. The function has to return a
+                `ScalarLinearExpression`. Therefore use the direct getter when
+                indexing variables.
+            * coords : coordinate-like
+                Coordinates to be processed by `xarray.DataArray`. For each
+                combination of coordinates, the function `rule` is called.
+                The order and size of coords has to be same as the argument list
+                followed by `model` in function `rule`.
+
+
         Returns
         -------
         linopy.LinearExpression
@@ -712,16 +742,42 @@ class Model:
         Examples
         --------
 
+        For creating an expression from tuples:
+        >>> from linopy import Model
+        >>> import pandas as pd
         >>> m = Model()
-        >>> m.add_variables(pd.Series([0, 0]), 1, name="x")
-        >>> m.add_variables(4, pd.Series([8, 10]), name="y")
+        >>> x = m.add_variables(pd.Series([0, 0]), 1, name="x")
+        >>> y = m.add_variables(4, pd.Series([8, 10]), name="y")
         >>> expr = m.linexpr((10, "x"), (1, "y"))
+
+        For creating an expression from a rule:
+        >>> m = Model()
+        >>> coords = pd.RangeIndex(10), ["a", "b"]
+        >>> a = m.add_variables(coords=coords)
+        >>> def rule(m, i, j):
+        ...     return a[i, j] + a[(i + 1) % 10, j]
+        ...
+        >>> expr = m.linexpr(rule, coords)
+
+        See also
+        --------
+        LinearExpression.from_tuples, LinearExpression.from_rule
         """
-        tuples = [
-            (c, self.variables[v]) if isinstance(v, str) else (c, v)
-            for (c, v) in tuples
-        ]
-        return LinearExpression.from_tuples(*tuples, chunk=self.chunk)
+        if callable(args[0]):
+            assert len(args) == 2, (
+                "When first argument is a function, only one second argument "
+                "containing a tuple or a single set of coords must be given."
+            )
+            rule, coords = args
+            return LinearExpression.from_rule(self, rule, coords)
+        if isinstance(args, tuple):
+            args = [
+                (c, self.variables[v]) if isinstance(v, str) else (c, v)
+                for (c, v) in args
+            ]
+            return LinearExpression.from_tuples(*args, chunk=self.chunk)
+        else:
+            raise TypeError(f"Not supported type {args}.")
 
     def _eval(self, expr: str, **kwargs):
         from pandas.core.computation.eval import eval as pd_eval
@@ -765,11 +821,11 @@ class Model:
         >>> m = linopy.Model()
         >>> lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
         >>> upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
-        >>> m.vareval("@lower <= x <= @upper")
+        >>> x = m.vareval("@lower <= x <= @upper")  # doctest:+SKIP
 
         This is the same as
 
-        >>> m.add_variables(lower, upper, name="x")
+        >>> x = m.add_variables(lower, upper, name="x")
         """
         if eval_kw is None:
             eval_kw = {}
@@ -812,17 +868,17 @@ class Model:
         >>> m = linopy.Model()
         >>> lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
         >>> upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
-        >>> m.add_variables(lower, upper, name="x")
-        >>> m.add_variables(lower, upper, name="y")
+        >>> x = m.add_variables(lower, upper, name="x")
+        >>> y = m.add_variables(lower, upper, name="y")
         >>> c = xr.DataArray(np.random.rand(10, 10), coords=[range(10), range(10)])
 
         Now create the linear expression
 
-        >>> m.lineval("@c * x - y")
+        >>> con = m.lineval("@c * x - y")
 
         This is the same as
 
-        >>> m.linexpr((c, "x"), (-1, "y"))
+        >>> con = m.linexpr((c, "x"), (-1, "y"))
         """
         if eval_kw is None:
             eval_kw = {}
@@ -866,18 +922,22 @@ class Model:
         >>> m = linopy.Model()
         >>> lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
         >>> upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
-        >>> m.add_variables(lower, upper, name="x")
-        >>> m.add_variables(lower, upper, name="y")
+        >>> x = m.add_variables(lower, upper, name="x")
+        >>> y = m.add_variables(lower, upper, name="y")
         >>> c = xr.DataArray(np.random.rand(10, 10), coords=[range(10), range(10)])
 
         Now create the constraint:
 
-        >>> m.coneval("@c * x - y <= 5 ")
+        >>> con = m.coneval("@c * x - y <= 5 ")
 
         This is the same as
 
         >>> lhs = m.linexpr((c, "x"), (-1, "y"))
-        >>> m.add_constraints(lhs, "<=", 5)
+        >>> con = m.add_constraints(lhs, "<=", 5)
+
+        or
+
+        >>> con = m.add_constraints(c * x - y <= 5)
         """
         if eval_kw is None:
             eval_kw = {}
diff --git a/linopy/remote.py b/linopy/remote.py
index bfd383f..0b5f4b0 100644
--- a/linopy/remote.py
+++ b/linopy/remote.py
@@ -98,18 +98,18 @@ class RemoteHandler:
     >>> coords = [arange(N), arange(N)]
     >>> x = m.add_variables(coords=coords)
     >>> y = m.add_variables(coords=coords)
-    >>> m.add_constraints(x - y >= arange(N))
-    >>> m.add_constraints(x + y >= 0)
-    >>> m.add_objective((2 * x + y).sum())
+    >>> con1 = m.add_constraints(x - y >= arange(N))
+    >>> con2 = m.add_constraints(x + y >= 0)
+    >>> obj = m.add_objective((2 * x + y).sum())
     >>>
-    >>> host = "gridlock.fias.uni-frankfurt.de"
-    >>> username = "hofmann"
-    >>> handler = linopy.remote.RemoteHandler(host, username=username)
+    >>> host = "my-remote-machine.com"
+    >>> username = "my-username"
+    >>> handler = linopy.remote.RemoteHandler(host, username=username)  # doctest: +SKIP
     >>>
     >>> # optionally activate a conda environment
-    >>> handler.execute("conda activate my-linopy-env")
+    >>> handler.execute("conda activate my-linopy-env")  # doctest: +SKIP
     >>>
-    >>> m = handler.solve_on_remote(m)
+    >>> m = handler.solve_on_remote(m)  # doctest: +SKIP
     """
 
     hostname: str = None
diff --git a/linopy/variables.py b/linopy/variables.py
index 228dac4..62c0f4e 100644
--- a/linopy/variables.py
+++ b/linopy/variables.py
@@ -48,44 +48,39 @@ class Variable(DataArray):
 
     Examples
     --------
+    >>> from linopy import Model
+    >>> import pandas as pd
     >>> m = Model()
     >>> x = m.add_variables(pd.Series([0, 0]), 1, name="x")
     >>> y = m.add_variables(4, pd.Series([8, 10]), name="y")
 
     Add variable together:
 
-    >>> x + y
-
-    ::
-
-        Linear Expression with 2 term(s):
-        ----------------------------------
-
-        Dimensions:  (dim_0: 2, _term: 2)
-        Coordinates:
-            * dim_0    (dim_0) int64 0 1
-            * _term    (_term) int64 0 1
-        Data:
-            coeffs   (dim_0, _term) int64 1 1 1 1
-            vars     (dim_0, _term) int64 1 3 2 4
-
+    >>> x + y  # doctest: +SKIP
+    Linear Expression with 2 term(s):
+    ----------------------------------
+    <BLANKLINE>
+    Dimensions:  (dim_0: 2, _term: 2)
+    Coordinates:
+      * dim_0    (dim_0) int64 0 1
+    Dimensions without coordinates: _term
+    Data:
+        coeffs   (dim_0, _term) int64 1 1 1 1
+        vars     (dim_0, _term) int64 0 2 1 3
 
     Multiply them with a coefficient:
 
-    >>> 3 * x
-
-    ::
-
-        Linear Expression with 1 term(s):
-        ----------------------------------
-
-        Dimensions:  (dim_0: 2, _term: 1)
-        Coordinates:
-            * _term    (_term) int64 0
-            * dim_0    (dim_0) int64 0 1
-        Data:
-            coeffs   (dim_0, _term) int64 3 3
-            vars     (dim_0, _term) int64 1 2
+    >>> 3 * x  # doctest: +SKIP
+    Linear Expression with 1 term(s):
+    ----------------------------------
+    <BLANKLINE>
+    Dimensions:  (dim_0: 2, _term: 1)
+    Coordinates:
+      * dim_0    (dim_0) int64 0 1
+    Dimensions without coordinates: _term
+    Data:
+        coeffs   (dim_0, _term) int64 3 3
+        vars     (dim_0, _term) int64 0 1
 
 
     Further operations like taking the negative and subtracting are supported.
@@ -114,6 +109,18 @@ class Variable(DataArray):
     __array_ufunc__ = None
     __array_priority__ = 10000
 
+    def __getitem__(self, keys) -> "ScalarVariable":
+        keys = (keys,) if not isinstance(keys, tuple) else keys
+        assert all(map(np.isscalar, keys)), (
+            "The get function of Variable is different as of xarray.DataArray. "
+            "Set single values for each dimension in order to obtain a "
+            "ScalarVariable. For all other purposes, use `.sel` and `.isel`."
+        )
+        assert self.ndim == len(keys), f"expected {self.ndim} keys, got {len(keys)}."
+        key = dict(zip(self.dims, keys))
+        selector = [self.get_index(k).get_loc(v) for k, v in key.items()]
+        return ScalarVariable(self.data[tuple(selector)])
+
     def to_array(self):
         """
         Convert the variable array to a xarray.DataArray.
@@ -606,3 +613,32 @@ class Variables:
             res[np.ravel(labels)] = np.ravel(block_map[name])
         res[-1] = -1
         return res
+
+
+@dataclass
+class ScalarVariable:
+    label: int
+    coords: dict = None
+
+    def to_linexpr(self, coeff=1):
+        if not isinstance(coeff, (int, float)):
+            raise TypeError(f"Coefficient must be a numeric value, got {type(coeff)}.")
+        return expressions.ScalarLinearExpression((coeff,), (self.label,))
+
+    def __neg__(self):
+        return self.to_linexpr(-1)
+
+    def __add__(self, other):
+        return self.to_linexpr(1) + other
+
+    def __sub__(self, other):
+        return self.to_linexpr(1) - other
+
+    def __mul__(self, coeff):
+        return self.to_linexpr(coeff)
+
+    def __rmul__(self, coeff):
+        return self.to_linexpr(coeff)
+
+    def __div__(self, coeff):
+        return self.to_linexpr(1 / coeff)
diff --git a/setup.py b/setup.py
index 5b2a6b0..0d071e6 100755
--- a/setup.py
+++ b/setup.py
@@ -49,7 +49,12 @@ setup(
             "ipykernel",
             "matplotlib",
         ],
-        "dev": ["pytest", "pytest-cov", "pre-commit"],
+        "dev": [
+            "pytest",
+            "pytest-cov",
+            "pre-commit",
+            "paramiko",
+        ],
     },
     classifiers=[
         "Development Status :: 3 - Alpha",


Test patch:
diff --git a/test/test_eval.py b/test/test_eval.py
index 43d4278..4020e67 100644
--- a/test/test_eval.py
+++ b/test/test_eval.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Created on Wed Oct 27 22:47:04 2021
+Created on Wed Oct 27 22:47:04 2021.
 
 @author: fabian
 """
diff --git a/test/test_io.py b/test/test_io.py
index 27f5b91..d12d9ab 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -41,7 +41,7 @@ def test_str_arrays_with_nans():
     x = m.add_variables(4, pd.Series([8, 10]), name="x")
     # now expand the second dimension, expended values of x will be nan
     y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]), name="y")
-    assert m["x"][-1].item() == -1
+    assert m["x"].data[-1].item() == -1
 
     da = int_to_str(m["x"].values)
     assert da.dtype == object
diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py
index 565d917..21501fe 100644
--- a/test/test_linear_expression.py
+++ b/test/test_linear_expression.py
@@ -13,6 +13,7 @@ import xarray as xr
 from xarray.testing import assert_equal
 
 from linopy import LinearExpression, Model, merge
+from linopy.expressions import ScalarLinearExpression
 
 m = Model()
 
@@ -282,3 +283,21 @@ def test_shift():
     assert shifted.nterm == 1
     assert shifted.coeffs.loc[:1].isnull().all()
     assert (shifted.vars.loc[:1] == -1).all()
+
+
+def test_scalarexpression_creation():
+    expr = 10 * x[0] + y[1] + z[1, 1]
+    assert isinstance(expr, ScalarLinearExpression)
+
+
+def test_scalarexpression_operations():
+    expr = 10 * x[0]
+    assert isinstance(expr, ScalarLinearExpression)
+
+    expr2 = 2 * expr
+    assert isinstance(expr2, ScalarLinearExpression)
+    assert expr2.coeffs == (20,)
+
+    expr3 = -expr
+    assert isinstance(expr3, ScalarLinearExpression)
+    assert expr3.coeffs == (-10,)
diff --git a/test/test_model_creation.py b/test/test_model_creation.py
index 8103e98..5ca2acf 100755
--- a/test/test_model_creation.py
+++ b/test/test_model_creation.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
-Created on Wed Mar 10 11:23:13 2021
+Created on Wed Mar 10 11:23:13 2021.
 
 @author: fabulous
 """
@@ -184,8 +184,10 @@ def test_masked_variables():
 
 def test_variable_merging():
     """
-    Test the merger of a variables with same dimension name but with
-    different lengths. Missing values should be filled up with -1.
+    Test the merger of a variables with same dimension name but with different
+    lengths.
+
+    Missing values should be filled up with -1.
     """
     m = Model()
 
diff --git a/test/test_variable.py b/test/test_variable.py
index 536f701..5b94606 100644
--- a/test/test_variable.py
+++ b/test/test_variable.py
@@ -15,6 +15,22 @@ import linopy
 from linopy import LinearExpression, Model
 
 
+def test_variable_getter():
+    m = Model()
+    x = m.add_variables(coords=[range(10)])
+
+    assert isinstance(x[0], linopy.variables.ScalarVariable)
+
+    with pytest.raises(AssertionError):
+        x[0, 0]
+
+    with pytest.raises(AssertionError):
+        x[0:5]
+
+    with pytest.raises(AssertionError):
+        x[[1, 2, 3]]
+
+
 def test_variable_repr():
     m = Model()
     m.variables.__repr__()