From de6b47bb1ea7e937ad32dc8a6c284beddca98685 Mon Sep 17 00:00:00 2001
From: Matthew Gidden <matthew.gidden@gmail.com>
Date: Tue, 24 May 2022 18:39:37 +0200
Subject: [PATCH] add offset method to compute offsets relative to base-year
 values (#659)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* add offset method to compute offsets relative to base-year values

* blacking

* apply faster set_index

Co-authored-by: Jonas Hörsch <coroa@posteo.de>

* provide a fill_value

* add to release notes

* stickler

Co-authored-by: Jonas Hörsch <coroa@posteo.de>
---
 RELEASE_NOTES.md   |  1 +
 pyam/core.py       | 39 +++++++++++++++++++++++++++++++++++++++
 tests/test_core.py | 21 +++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
index 6654c69cc..b3b370dd6 100644
--- a/RELEASE_NOTES.md
+++ b/RELEASE_NOTES.md
@@ -1,5 +1,6 @@
 # Next release
 
+- [#659](https://github.com/IAMconsortium/pyam/pull/659) Add an `offset` method
 - [#657](https://github.com/IAMconsortium/pyam/pull/657) Add an `IamSlice` class
 
 # Release v1.4.0
diff --git a/pyam/core.py b/pyam/core.py
index cb958c9f1..47baca09c 100755
--- a/pyam/core.py
+++ b/pyam/core.py
@@ -1269,6 +1269,45 @@ def normalize(self, inplace=False, **kwargs):
         if not inplace:
             return ret
 
+    def offset(self, padding=0, fill_value=None, inplace=False, **kwargs):
+        """Compute new data which is offset from a specific data point
+
+        For example, offsetting from `year=2005` will provide data
+        *relative* to `year=2005` such that the value in 2005 is 0 and
+        all other values `value[year] - value[2005]`.
+
+        Conceptually this operation performs as:
+        ```
+        df - df.filter(**kwargs) + padding
+        ```
+
+        Note: Currently only supports normalizing to a specific time.
+
+        Parameters
+        ----------
+        padding : float, optional
+            an additional offset padding
+        fill_value : float or None, optional
+            Applied on subtraction. Fills exisiting missing (NaN) values. See
+            https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.subtract.html
+        inplace : bool, optional
+            if :obj:`True`, do operation inplace and return None
+        kwargs
+            the column and value on which to offset (e.g., `year=2005`)
+        """
+        if len(kwargs) > 1 or self.time_col not in kwargs:
+            raise ValueError("Only time(year)-based normalization supported")
+        ret = self.copy() if not inplace else self
+        data = ret._data
+        value = kwargs[self.time_col]
+        base_value = data.loc[data.index.isin([value], level=self.time_col)].droplevel(
+            self.time_col
+        )
+        ret._data = data.subtract(base_value, fill_value=fill_value) + padding
+
+        if not inplace:
+            return ret
+
     def aggregate(
         self,
         variable,
diff --git a/tests/test_core.py b/tests/test_core.py
index f87d9c173..f13973b8c 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -789,3 +789,24 @@ def test_normalize(test_df):
 def test_normalize_not_time(test_df):
     pytest.raises(ValueError, test_df.normalize, variable="foo")
     pytest.raises(ValueError, test_df.normalize, year=2015, variable="foo")
+
+
+@pytest.mark.parametrize("padding", [0, 2])
+def test_offset(test_df, padding):
+    exp = test_df.data.copy().reset_index(drop=True)
+    exp.loc[1::2, "value"] -= exp["value"][::2].values - padding
+    exp.loc[::2, "value"] -= exp["value"][::2].values - padding
+    # only call with kwarg if padding != 0 (the default)
+    kwargs = {"padding": padding} if padding else {}
+    if "year" in test_df.data:
+        obs = test_df.offset(year=2005, **kwargs).data.reset_index(drop=True)
+    else:
+        obs = test_df.offset(
+            time=datetime.datetime(2005, 6, 17), **kwargs
+        ).data.reset_index(drop=True)
+    pd.testing.assert_frame_equal(obs, exp)
+
+
+def test_offset_not_time(test_df):
+    pytest.raises(ValueError, test_df.offset, variable="foo")
+    pytest.raises(ValueError, test_df.offset, year=2015, variable="foo")