From e9b70876d6bb9bd6e1cd6824941048450a4c60f0 Mon Sep 17 00:00:00 2001 From: Daniel Huppmann Date: Thu, 13 Feb 2025 09:01:36 +0100 Subject: [PATCH] Add fast route for operations with non-SI units --- pyam/_ops.py | 2 +- tests/test_ops.py | 46 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/pyam/_ops.py b/pyam/_ops.py index ef2f70ff6..a02f7a8c0 100644 --- a/pyam/_ops.py +++ b/pyam/_ops.py @@ -83,7 +83,7 @@ def _op_data(df, name, method, axis, fillna=None, args=(), ignore_units=False, * and fillna is None and len(_unit_kwds["a"]) == 1 and len(_unit_kwds["b"]) == 1 - and registry.Unit(_unit_kwds["a"][0]) == registry.Unit(_unit_kwds["b"][0]) + and _unit_kwds["a"][0] == _unit_kwds["b"][0] ): # activate ignore-units feature ignore_units = _unit_kwds["a"][0] if method in [add, subtract] else "" diff --git a/tests/test_ops.py b/tests/test_ops.py index 9ddf2dc5b..f97c9f665 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -96,6 +96,21 @@ def test_add_variable(test_df_year, arg, df_func, fillna, ignore_units, append): assert_iamframe_equal(exp, obs) +@pytest.mark.parametrize("append", (False, True)) +def test_add_scenario_non_si_unit(test_df_year, append): + df = test_df_year.rename(unit={"EJ/yr": "foo"}) + exp = df_ops_variable(operator.add, "Sum", unit="foo", meta=test_df_year.meta) + args = ("Primary Energy", "Primary Energy|Coal", "Sum") + + if append: + obs = df.copy() + obs.add(*args, append=True) + assert_iamframe_equal(df.append(exp), obs) + else: + obs = df.add(*args) + assert_iamframe_equal(exp, obs) + + @pytest.mark.parametrize("append", (False, True)) def test_add_scenario(test_df_year, append): """Verify that in-dataframe addition works on a custom axis (`scenario`)""" @@ -159,6 +174,22 @@ def test_subtract_variable(test_df_year, arg, df_func, fillna, append, ignore_un assert_iamframe_equal(exp, test_df_year.subtract(*args, **kwds)) +@pytest.mark.parametrize("append", (False, True)) +def test_subtract_scenario_non_si_unit_unit(test_df_year, append): + df = test_df_year.rename(unit={"EJ/yr": "foo"}) + exp = df_ops_variable(operator.sub, "Diff", unit="foo", meta=test_df_year.meta) + args = ("Primary Energy", "Primary Energy|Coal", "Diff") + + if append: + obs = df.copy() + obs.subtract(*args, append=True) + assert_iamframe_equal(df.append(exp), obs) + else: + obs = df.subtract(*args) + assert_iamframe_equal(exp, obs) + + + @pytest.mark.parametrize("append", (False, True)) def test_subtract_scenario(test_df_year, append): """Verify that in-dataframe subtraction works on a custom axis (`scenario`)""" @@ -288,6 +319,21 @@ def test_divide_variable(test_df_year, arg, df_func, fillna, append, ignore_unit assert_iamframe_equal(exp, test_df_year.divide(*args, **kwds)) +@pytest.mark.parametrize("append", (False, True)) +def test_divide_variable_non_si_unit_unit(test_df_year, append): + df = test_df_year.rename(unit={"EJ/yr": "foo"}) + exp = df_ops_variable(operator.truediv, "Ratio", unit="", meta=test_df_year.meta) + args = ("Primary Energy", "Primary Energy|Coal", "Ratio") + + if append: + obs = df.copy() + obs.divide(*args, append=True) + assert_iamframe_equal(df.append(exp), obs) + else: + obs = df.divide(*args) + assert_iamframe_equal(exp, obs) + + @pytest.mark.parametrize("append", (False, True)) def test_divide_scenario(test_df_year, append): """Verify that in-dataframe addition works on a custom axis (`scenario`)"""