add offset method to compute offsets relative to base-year values (#659)

* add offset method to compute offsets relative to base-year values * blacking * apply faster set_index Co-authored-by: Jonas Hörsch <[email protected]> * provide a fill_value * add to release notes * stickler Co-authored-by: Jonas Hörsch <[email protected]>
IAMconsortium · May 24, 2022 · de6b47b · de6b47b
1 parent f8e50aa
commit de6b47b
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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")