From b9f425e15974eb4ec8e273975358ec16e093c77e Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Wed, 12 Jul 2023 08:03:49 +0200 Subject: [PATCH 01/10] Implement Variable.cumsum --- linopy/variables.py | 43 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/linopy/variables.py b/linopy/variables.py index c6e4e1b1..24f04769 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -17,7 +17,8 @@ import pandas as pd from deprecation import deprecated from numpy import floating, inf, issubdtype -from xarray import DataArray, Dataset, align, broadcast, zeros_like +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: bool | None = None, + keep_attrs: bool | None = 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): """ From 88f9c1301adafd9e11beb8d6c7886c580a59804c Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Wed, 12 Jul 2023 08:04:32 +0200 Subject: [PATCH 02/10] Implement LinearExpression.cumsum --- linopy/expressions.py | 53 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/linopy/expressions.py b/linopy/expressions.py index a15e2233..fda54d08 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -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,58 @@ def sum(self, dims=None, drop_zeros=False) -> "LinearExpression": return res + + def cumsum( + self, + dim: Dims = None, + *, + skipna: bool | None = None, + keep_attrs: bool | None = 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 = [d for d in self.data.dims.keys() if d != "_term"] + 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): """ From 25565eeaa79a352fc8f4b6c44a683c6573ddc316 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 12 Jul 2023 06:10:37 +0000 Subject: [PATCH 03/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- linopy/expressions.py | 6 +----- linopy/variables.py | 2 +- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index fda54d08..a6170344 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -587,7 +587,6 @@ def sum(self, dims=None, drop_zeros=False) -> "LinearExpression": return res - def cumsum( self, dim: Dims = None, @@ -633,10 +632,7 @@ def cumsum( 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 - } + 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 diff --git a/linopy/variables.py b/linopy/variables.py index 24f04769..3aad483f 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -17,7 +17,7 @@ import pandas as pd from deprecation import deprecated from numpy import floating, inf, issubdtype -from xarray import DataArray, Dataset, align, broadcast, zeros_like +from xarray import DataArray, Dataset, align, broadcast, zeros_like from xarray.core.types import Dims import linopy.expressions as expressions From 0cc1f38dba21087a51502fccb6715fee3f14a424 Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Wed, 12 Jul 2023 08:19:56 +0200 Subject: [PATCH 04/10] Fix type annotations in variables.py --- linopy/variables.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/linopy/variables.py b/linopy/variables.py index 3aad483f..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 @@ -420,8 +420,8 @@ def cumsum( self, dim: Dims = None, *, - skipna: bool | None = None, - keep_attrs: bool | None = None, + skipna: Optional[bool] = None, + keep_attrs: Optional[bool] = None, **kwargs: Any, ) -> "expressions.LinearExpression": """ From e9239a18104697977d6d8b9594503bac55f01a30 Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Wed, 12 Jul 2023 08:20:46 +0200 Subject: [PATCH 05/10] Fix type annotations in expressions.py --- linopy/expressions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index a6170344..ab6b18e8 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 @@ -591,8 +591,8 @@ def cumsum( self, dim: Dims = None, *, - skipna: bool | None = None, - keep_attrs: bool | None = None, + skipna: Optional[bool] = None, + keep_attrs: Optional[bool] = None, **kwargs: Any, ) -> "LinearExpression": """ From 9bd968dda30ef6b9f85409fd949c0dc832fd5a8e Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Thu, 13 Jul 2023 18:51:31 +0200 Subject: [PATCH 06/10] directly access coord dimensions Co-authored-by: Fabian Hofmann --- linopy/expressions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index ab6b18e8..8e111aec 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -628,7 +628,7 @@ def cumsum( if not dim: # If user did not specify a dimension to sum over, use all relevant # dimensions - dim = [d for d in self.data.dims.keys() if d != "_term"] + dim = self.coord_dims if isinstance(dim, str): # Make sure, single mentioned dimensions is handled correctly. dim = [dim] From be3aa2e8d07ada09adf60b634acbe8d29e7ed419 Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Fri, 14 Jul 2023 16:36:51 +0200 Subject: [PATCH 07/10] Add test for LinearExpression.cumsum --- test/test_linear_expression.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index eb3f6bc2..d171297b 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -603,3 +603,19 @@ 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() + assert_linequal(cumsum.loc[0], multiple * var.loc[0]) + assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) + + # Test cumsum on sum of variables + var = m.variables["x"] + var = m.variables["y"] + cumsum = (multiple * var).cumsum() + assert_linequal(cumsum.loc[0], multiple * var.loc[0]) + assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) + From 52e60da3f82e99128b7aec87c37035e73356b9b7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Jul 2023 14:37:00 +0000 Subject: [PATCH 08/10] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- test/test_linear_expression.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index d171297b..7a54e111 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -618,4 +618,3 @@ def test_cumsum(m, multiple): cumsum = (multiple * var).cumsum() assert_linequal(cumsum.loc[0], multiple * var.loc[0]) assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) - From 1a2ecabea4b861b7bbf321c871fe3d469ef6b4ca Mon Sep 17 00:00:00 2001 From: Henning Blunck Date: Fri, 14 Jul 2023 16:38:35 +0200 Subject: [PATCH 09/10] Bugfix in test --- 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 7a54e111..4382bc47 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -614,7 +614,7 @@ def test_cumsum(m, multiple): assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) # Test cumsum on sum of variables - var = m.variables["x"] + var = m.variables["y"] + var = m.variables["x"] + m.variables["y"] cumsum = (multiple * var).cumsum() assert_linequal(cumsum.loc[0], multiple * var.loc[0]) assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) From 3a56e6e6e435e5e7a615276079ee1b7e5cb4c868 Mon Sep 17 00:00:00 2001 From: Fabian Hofmann Date: Tue, 18 Jul 2023 19:00:47 +0200 Subject: [PATCH 10/10] Apply suggestions from code review --- test/test_linear_expression.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 4382bc47..4684dc63 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -610,11 +610,9 @@ def test_cumsum(m, multiple): # Test cumsum on variable x var = m.variables["x"] cumsum = (multiple * var).cumsum() - assert_linequal(cumsum.loc[0], multiple * var.loc[0]) - assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) + cumsum.nterm == 2 # Test cumsum on sum of variables var = m.variables["x"] + m.variables["y"] cumsum = (multiple * var).cumsum() - assert_linequal(cumsum.loc[0], multiple * var.loc[0]) - assert_linequal(cumsum.loc[0], multiple * (var.loc[0] + var.loc[1])) + cumsum.nterm == 2