diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 3c6a4ab5c..faeb0724c 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -2,6 +2,7 @@ - [#930](https://github.com/IAMconsortium/pyam/pull/930) Fix the `categorize()` method if invalid datapoints exceed the number of scenarios +- [#907](https://github.com/IAMconsortium/pyam/pull/907) Add fast route for operations with non-SI units # Release v3.1.0 diff --git a/pyam/_ops.py b/pyam/_ops.py index ef2f70ff6..6e28056fc 100644 --- a/pyam/_ops.py +++ b/pyam/_ops.py @@ -61,7 +61,7 @@ def _op_data(df, name, method, axis, fillna=None, args=(), ignore_units=False, * cols = [d for d in df.dimensions if d != axis] - # replace args and and kwds with values of `df._data` if applicable + # replace args and kwds with values of `df._data` if applicable # _data_args and _data_kwds track if an argument was replaced by `df._data` values n = len(args) _args, _data_args, _units_args = [None] * n, [False] * n, [None] * n @@ -71,7 +71,7 @@ def _op_data(df, name, method, axis, fillna=None, args=(), ignore_units=False, * ) _data_kwds, _unit_kwds = {}, {} - for i, (key, value) in enumerate(kwds.items()): + for key, value in kwds.items(): kwds[key], _unit_kwds[key], _data_kwds[key] = _get_values( df, axis, value, cols, key ) @@ -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..c17002053 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -56,44 +56,81 @@ def test_add_raises(test_df_year): @pytest.mark.parametrize( - "arg, df_func, fillna, ignore_units", + "variable, df_func, expected_unit", ( - ("Primary Energy|Coal", df_ops_variable, None, False), - ("Primary Energy|Coal", df_ops_fillna_0, 0, "foo"), - ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}, "foo"), - ("Primary Energy|Coal", df_ops_variable_default, 5, "foo"), - (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, None, False), - (2, df_ops_variable_number, None, "foo"), + ("Primary Energy|Coal", df_ops_variable, "EJ/yr"), + # pint changes the unit into its standard format + (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, "EJ / a"), ), ) @pytest.mark.parametrize("append", (False, True)) -def test_add_variable(test_df_year, arg, df_func, fillna, ignore_units, append): - """Verify that in-dataframe addition works on the default `variable` axis""" +def test_add_variable(test_df_year, variable, df_func, expected_unit, append): + """Check that in-dataframe addition works on the default `variable` axis""" - # change one unit to make ignore-units strictly necessary - if ignore_units: - test_df_year.rename( - variable={"Primary Energy": "Primary Energy"}, - unit={"EJ/yr": "custom_unit"}, - inplace=True, - ) + exp = df_func(operator.add, "Sum", unit=expected_unit, meta=test_df_year.meta) + + if append: + obs = test_df_year.copy() + obs.add("Primary Energy", variable, "Sum", append=True) + exp = test_df_year.append(exp) + else: + obs = test_df_year.add("Primary Energy", variable, "Sum") + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize( + "arg, df_func, fillna", + ( + ("Primary Energy|Coal", df_ops_fillna_0, 0), + ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}), + ("Primary Energy|Coal", df_ops_variable_default, 5), + (2, df_ops_variable_number, None), + ), +) +@pytest.mark.parametrize("append", (False, True)) +def test_add_variable_ignore_units(test_df_year, arg, df_func, fillna, append): + """Check that in-dataframe addition works with ignore_units""" + + # change one unit to make ignore_units strictly necessary + test_df_year.rename( + variable={"Primary Energy": "Primary Energy"}, + unit={"EJ/yr": "custom_unit"}, + inplace=True, + ) - unit = "EJ/yr" if ignore_units is False else ignore_units - exp = df_func(operator.add, "Sum", unit=unit, meta=test_df_year.meta) + exp = df_func(operator.add, "Sum", unit="foo", meta=test_df_year.meta) args = ("Primary Energy", arg, "Sum") - kwds = dict(ignore_units=ignore_units, fillna=fillna) if append: obs = test_df_year.copy() - obs.add(*args, **kwds, append=True) - assert_iamframe_equal(test_df_year.append(exp), obs) + obs.add(*args, ignore_units="foo", fillna=fillna, append=True) + exp = test_df_year.append(exp) else: # check that incompatible units raise the expected error - if ignore_units: - with pytest.raises(pint.UndefinedUnitError): - test_df_year.add(*args, fillna=fillna, ignore_units=False) - obs = test_df_year.add(*args, **kwds) - assert_iamframe_equal(exp, obs) + with pytest.raises(pint.UndefinedUnitError): + test_df_year.add(*args, fillna=fillna) + + # using ignore_units works as expected + obs = test_df_year.add(*args, ignore_units="foo", fillna=fillna) + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize("append", (False, True)) +def test_add_variable_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) + + if append: + obs = df.copy() + obs.add("Primary Energy", "Primary Energy|Coal", "Sum", append=True) + exp = df.append(exp) + else: + obs = df.add("Primary Energy", "Primary Energy|Coal", "Sum") + + assert_iamframe_equal(exp, obs) @pytest.mark.parametrize("append", (False, True)) @@ -120,43 +157,80 @@ def test_add_scenario(test_df_year, append): @pytest.mark.parametrize( - "arg, df_func, fillna, ignore_units", + "arg, df_func, expected_unit", ( - ("Primary Energy|Coal", df_ops_variable, None, False), - ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}, "foo"), - ("Primary Energy|Coal", df_ops_variable_default, 5, "foo"), - (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, None, False), - (2, df_ops_variable_number, None, "foo"), + ("Primary Energy|Coal", df_ops_variable, "EJ/yr"), + # pint changes the unit into its standard format + (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, "EJ / a"), ), ) @pytest.mark.parametrize("append", (False, True)) -def test_subtract_variable(test_df_year, arg, df_func, fillna, append, ignore_units): - """Verify that in-dataframe subtraction works on the default `variable` axis""" +def test_subtract_variable(test_df_year, arg, df_func, expected_unit, append): + """Check that in-dataframe subtraction works on the default `variable` axis""" + + exp = df_func(operator.sub, "Diff", unit=expected_unit, meta=test_df_year.meta) + + if append: + obs = test_df_year.copy() + obs.subtract("Primary Energy", arg, "Diff", append=True) + exp = test_df_year.append(exp) + else: + obs = test_df_year.subtract("Primary Energy", arg, "Diff") + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize( + "arg, df_func, fillna", + ( + ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}), + ("Primary Energy|Coal", df_ops_variable_default, 5), + (2, df_ops_variable_number, None), + ), +) +@pytest.mark.parametrize("append", (False, True)) +def test_subtract_variable_ignore_units(test_df_year, arg, df_func, fillna, append): + """Check that in-dataframe subtraction works witgh ignore_units""" # change one unit to make ignore-units strictly necessary - if ignore_units: - test_df_year.rename( - variable={"Primary Energy": "Primary Energy"}, - unit={"EJ/yr": "custom_unit"}, - inplace=True, - ) + test_df_year.rename( + variable={"Primary Energy": "Primary Energy"}, + unit={"EJ/yr": "custom_unit"}, + inplace=True, + ) - unit = "EJ/yr" if ignore_units is False else ignore_units - exp = df_func(operator.sub, "Diff", unit=unit, meta=test_df_year.meta) + exp = df_func(operator.sub, "Diff", unit="foo", meta=test_df_year.meta) args = ("Primary Energy", arg, "Diff") - kwds = dict(ignore_units=ignore_units, fillna=fillna) if append: obs = test_df_year.copy() - obs.subtract(*args, **kwds, append=True) - assert_iamframe_equal(test_df_year.append(exp), obs) + obs.subtract(*args, ignore_units="foo", fillna=fillna, append=True) + exp = test_df_year.append(exp) else: # check that incompatible units raise the expected error - if ignore_units: - with pytest.raises(pint.UndefinedUnitError): - test_df_year.add(*args, fillna=fillna, ignore_units=False) + with pytest.raises(pint.UndefinedUnitError): + test_df_year.add(*args, fillna=fillna) - assert_iamframe_equal(exp, test_df_year.subtract(*args, **kwds)) + # using ignore_units works as expected + obs = test_df_year.subtract(*args, ignore_units="foo", fillna=fillna) + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize("append", (False, True)) +def test_subtract_variable_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) + + if append: + obs = df.copy() + obs.subtract("Primary Energy", "Primary Energy|Coal", "Diff", append=True) + exp = df.append(exp) + else: + obs = df.subtract("Primary Energy", "Primary Energy|Coal", "Diff") + + assert_iamframe_equal(exp, obs) @pytest.mark.parametrize("append", (False, True)) @@ -183,44 +257,61 @@ def test_subtract_scenario(test_df_year, append): @pytest.mark.parametrize( - "arg, df_func, fillna, ignore_units", + "arg, df_func, expected_unit", ( - ("Primary Energy|Coal", df_ops_variable, None, False), - ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}, "foo"), - ("Primary Energy|Coal", df_ops_variable_default, 5, "foo"), - # note that multiplying with pint reformats the unit - (2, df_ops_variable_number, None, False), + ("Primary Energy|Coal", df_ops_variable, "EJ ** 2 / a ** 2"), + (2, df_ops_variable_number, "EJ / a"), ), ) @pytest.mark.parametrize("append", (False, True)) -def test_multiply_variable(test_df_year, arg, df_func, fillna, ignore_units, append): - """Verify that in-dataframe addition works on the default `variable` axis""" - - if ignore_units: - # change one unit to make ignore-units strictly necessary - test_df_year.rename( - variable={"Primary Energy": "Primary Energy"}, - unit={"EJ/yr": "custom_unit"}, - inplace=True, - ) - unit = ignore_units +def test_multiply_variable(test_df_year, arg, df_func, expected_unit, append): + """Check that in-dataframe addition works on the default `variable` axis""" + + exp = df_func(operator.mul, "Prod", unit=expected_unit, meta=test_df_year.meta) + + if append: + obs = test_df_year.copy() + obs.multiply("Primary Energy", arg, "Prod", append=True) + exp = test_df_year.append(exp) else: - unit = "EJ / a" if isinstance(arg, int) else "EJ ** 2 / a ** 2" - exp = df_func(operator.mul, "Prod", unit=unit, meta=test_df_year.meta) + obs = test_df_year.multiply("Primary Energy", arg, "Prod") + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize( + "arg, df_func, fillna", + ( + ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}), + ("Primary Energy|Coal", df_ops_variable_default, 5), + ), +) +@pytest.mark.parametrize("append", (False, True)) +def test_multiply_variable_ignore_units(test_df_year, arg, df_func, fillna, append): + """Check that in-dataframe addition works with ignore_units""" + + # change one unit to make ignore_units strictly necessary + test_df_year.rename( + variable={"Primary Energy": "Primary Energy"}, + unit={"EJ/yr": "custom_unit"}, + inplace=True, + ) + + exp = df_func(operator.mul, "Prod", unit="foo", meta=test_df_year.meta) args = ("Primary Energy", arg, "Prod") - kwds = dict(ignore_units=ignore_units, fillna=fillna) if append: obs = test_df_year.copy() - obs.multiply(*args, **kwds, append=True) - assert_iamframe_equal(test_df_year.append(exp), obs) + obs.multiply(*args, ignore_units="foo", fillna=fillna, append=True) + exp = test_df_year.append(exp) else: # check that incompatible units raise the expected error - if ignore_units: - with pytest.raises(pint.UndefinedUnitError): - test_df_year.add(*args, fillna=fillna, ignore_units=False) + with pytest.raises(pint.UndefinedUnitError): + test_df_year.add(*args, fillna=fillna) + + obs = test_df_year.multiply(*args, ignore_units="foo", fillna=fillna) - assert_iamframe_equal(exp, test_df_year.multiply(*args, **kwds)) + assert_iamframe_equal(exp, obs) @pytest.mark.parametrize("append", (False, True)) @@ -247,45 +338,79 @@ def test_multiply_scenario(test_df_year, append): @pytest.mark.parametrize( - "arg, df_func, fillna, ignore_units", + "arg, df_func, expected_unit", ( - ("Primary Energy|Coal", df_ops_variable, None, False), - ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}, "foo"), - ("Primary Energy|Coal", df_ops_variable_default, 5, "foo"), - (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, None, False), - (2, df_ops_variable_number, None, False), + ("Primary Energy|Coal", df_ops_variable, ""), + (registry.Quantity(2, "EJ/yr"), df_ops_variable_number, ""), + (2, df_ops_variable_number, "EJ / a"), ), ) @pytest.mark.parametrize("append", (False, True)) -def test_divide_variable(test_df_year, arg, df_func, fillna, append, ignore_units): - """Verify that in-dataframe addition works on the default `variable` axis""" - - # note that dividing with pint reformats the unit - if ignore_units: - # change one unit to make ignore-units strictly necessary - test_df_year.rename( - variable={"Primary Energy": "Primary Energy"}, - unit={"EJ/yr": "custom_unit"}, - inplace=True, - ) - unit = ignore_units +def test_divide_variable(test_df_year, arg, df_func, expected_unit, append): + """Check that in-dataframe addition works on the default `variable` axis""" + + exp = df_func(operator.truediv, "Ratio", unit=expected_unit, meta=test_df_year.meta) + + if append: + obs = test_df_year.copy() + obs.divide("Primary Energy", arg, "Ratio", append=True) + exp = test_df_year.append(exp) else: - unit = "EJ / a" if isinstance(arg, int) else "" - exp = df_func(operator.truediv, "Ratio", unit=unit, meta=test_df_year.meta) + obs = test_df_year.divide("Primary Energy", arg, "Ratio") + + assert_iamframe_equal(exp, obs) + + +@pytest.mark.parametrize( + "arg, df_func, fillna", + ( + ("Primary Energy|Coal", df_ops_variable_default, {"c": 7, "b": 5}), + ("Primary Energy|Coal", df_ops_variable_default, 5), + ), +) +@pytest.mark.parametrize("append", (False, True)) +def test_divide_variable_ignore_units(test_df_year, arg, df_func, fillna, append): + """Check that in-dataframe addition works with ignore_units""" + + # change one unit to make ignore_units strictly necessary + test_df_year.rename( + variable={"Primary Energy": "Primary Energy"}, + unit={"EJ/yr": "custom_unit"}, + inplace=True, + ) + + exp = df_func(operator.truediv, "Ratio", unit="foo", meta=test_df_year.meta) args = ("Primary Energy", arg, "Ratio") - kwds = dict(ignore_units=ignore_units, fillna=fillna) if append: obs = test_df_year.copy() - obs.divide(*args, **kwds, append=True) - assert_iamframe_equal(test_df_year.append(exp), obs) + obs.divide(*args, ignore_units="foo", fillna=fillna, append=True) + exp = test_df_year.append(exp) else: # check that incompatible units raise the expected error - if ignore_units: - with pytest.raises(pint.UndefinedUnitError): - test_df_year.add(*args, fillna=fillna, ignore_units=False) + with pytest.raises(pint.UndefinedUnitError): + test_df_year.add(*args, fillna=fillna) + + # using ignore_units works as expected + obs = test_df_year.divide(*args, ignore_units="foo", fillna=fillna) + + assert_iamframe_equal(exp, obs) + + +@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) + + if append: + obs = df.copy() + obs.divide("Primary Energy", "Primary Energy|Coal", "Ratio", append=True) + exp = df.append(exp) + else: + obs = df.divide("Primary Energy", "Primary Energy|Coal", "Ratio") - assert_iamframe_equal(exp, test_df_year.divide(*args, **kwds)) + assert_iamframe_equal(exp, obs) @pytest.mark.parametrize("append", (False, True))