diff --git a/linopy/expressions.py b/linopy/expressions.py index a15e2233..8e111aec 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -11,7 +11,7 @@ import warnings from dataclasses import dataclass, field from itertools import product, zip_longest -from typing import Any, Mapping, Union +from typing import Any, Mapping, Optional, Union import numpy as np import pandas as pd @@ -23,6 +23,7 @@ from scipy.sparse import csc_matrix from xarray import DataArray, Dataset from xarray.core.dataarray import DataArrayCoordinates +from xarray.core.types import Dims from linopy import constraints, expressions, variables from linopy.common import ( @@ -586,6 +587,54 @@ def sum(self, dims=None, drop_zeros=False) -> "LinearExpression": return res + def cumsum( + self, + dim: Dims = None, + *, + skipna: Optional[bool] = None, + keep_attrs: Optional[bool] = None, + **kwargs: Any, + ) -> "LinearExpression": + """ + Cumulated sum along a given axis. + + Docstring and arguments are borrowed from `xarray.Dataset.cumsum` + + Parameters + ---------- + dim : str, Iterable of Hashable, "..." or None, default: None + Name of dimension[s] along which to apply ``cumsum``. For e.g. ``dim="x"`` + or ``dim=["x", "y"]``. If "..." or None, will reduce over all dimensions. + skipna : bool or None, optional + If True, skip missing values (as marked by NaN). By default, only + skips missing values for float dtypes; other dtypes either do not + have a sentinel missing value (int) or ``skipna=True`` has not been + implemented (object, datetime64 or timedelta64). + keep_attrs : bool or None, optional + If True, ``attrs`` will be copied from the original + object to the new one. If False, the new object will be + returned without attributes. + **kwargs : Any + Additional keyword arguments passed on to the appropriate array + function for calculating ``cumsum`` on this object's data. + These could include dask-specific kwargs like ``split_every``. + + Returns + ------- + linopy.expression.LinearExpression + """ + # Along every dimensions, we want to perform cumsum along, get the size of the + # dimension to pass that to self.rolling. + if not dim: + # If user did not specify a dimension to sum over, use all relevant + # dimensions + dim = self.coord_dims + if isinstance(dim, str): + # Make sure, single mentioned dimensions is handled correctly. + dim = [dim] + dim_dict = {dim_name: self.data.dims[dim_name] for dim_name in dim} + return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) + @classmethod def from_tuples(cls, *tuples, model=None, chunk=None): """ diff --git a/linopy/variables.py b/linopy/variables.py index c6e4e1b1..b10934be 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -9,7 +9,7 @@ import logging from collections.abc import Iterable from dataclasses import dataclass, field -from typing import Any, Dict, Mapping, Sequence, Union +from typing import Any, Dict, Mapping, Optional, Sequence, Union from warnings import warn import dask @@ -18,6 +18,7 @@ from deprecation import deprecated from numpy import floating, inf, issubdtype from xarray import DataArray, Dataset, align, broadcast, zeros_like +from xarray.core.types import Dims import linopy.expressions as expressions from linopy.common import ( @@ -415,6 +416,46 @@ def rolling( dim=dim, min_periods=min_periods, center=center, **window_kwargs ) + def cumsum( + self, + dim: Dims = None, + *, + skipna: Optional[bool] = None, + keep_attrs: Optional[bool] = None, + **kwargs: Any, + ) -> "expressions.LinearExpression": + """ + Cumulated sum along a given axis. + + Docstring and arguments are borrowed from `xarray.Dataset.cumsum` + + Parameters + ---------- + dim : str, Iterable of Hashable, "..." or None, default: None + Name of dimension[s] along which to apply ``cumsum``. For e.g. ``dim="x"`` + or ``dim=["x", "y"]``. If "..." or None, will reduce over all dimensions. + skipna : bool or None, optional + If True, skip missing values (as marked by NaN). By default, only + skips missing values for float dtypes; other dtypes either do not + have a sentinel missing value (int) or ``skipna=True`` has not been + implemented (object, datetime64 or timedelta64). + keep_attrs : bool or None, optional + If True, ``attrs`` will be copied from the original + object to the new one. If False, the new object will be + returned without attributes. + **kwargs : Any + Additional keyword arguments passed on to the appropriate array + function for calculating ``cumsum`` on this object's data. + These could include dask-specific kwargs like ``split_every``. + + Returns + ------- + linopy.expression.LinearExpression + """ + return self.to_linexpr().cumsum( + dim=dim, skipna=skipna, keep_attrs=keep_attrs, **kwargs + ) + @property def name(self): """ diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index eb3f6bc2..4684dc63 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -603,3 +603,16 @@ def test_rename(x, y, z): renamed = expr.rename({"dim_0": "dim_1", "dim_1": "dim_2"}) assert set(renamed.dims) == {"dim_1", "dim_2", TERM_DIM} assert renamed.nterm == 3 + + +@pytest.mark.parametrize("multiple", [1.0, 0.5, 2.0, 0.0]) +def test_cumsum(m, multiple): + # Test cumsum on variable x + var = m.variables["x"] + cumsum = (multiple * var).cumsum() + cumsum.nterm == 2 + + # Test cumsum on sum of variables + var = m.variables["x"] + m.variables["y"] + cumsum = (multiple * var).cumsum() + cumsum.nterm == 2