Test & improve perfomance in .tools.add_year

message_ix/tests/tools/test_add_year.py

@@ -1,8 +1,8 @@
 import pytest
 from ixmp import Platform
 
-from message_ix import Scenario
-from message_ix.tools.add_year import add_year
+from message_ix import Scenario, make_df
+from message_ix.tools.add_year import add_year, i1d_genno, interpolate_1d
 
 
 @pytest.fixture
@@ -143,3 +143,70 @@ def test_add_year_cli(message_ix_cli, base_scen_mp):
     # Bad usage: not giving the base scenario info
     r = message_ix_cli(*cmd[6:], "--dry-run")
     assert r.exit_code == 2
+
+
+@pytest.mark.parametrize("func", [interpolate_1d, i1d_genno])
+def test_interpolate_1d(func):
+    # Input data
+    years_base = [2020, 2030, 2040]
+    df = make_df(
+        "technical_lifetime",
+        node_loc="n",
+        technology="t",
+        value=[10.0, 20, 30],
+        unit="year",
+        year_vtg=years_base,
+    )
+
+    years_new = [2015, 2025, 2035, 2045]
+
+    # With default extrapolate=False, extrapol_neg=None
+    args = dict(
+        df=df,
+        yrs_new=years_new,
+        horizon=years_base,
+        year_col="year_vtg",
+        extrapolate=False,
+        extrapol_neg=None,
+        bound_extend=True,
+    )
+    result = func(**args)
+    print("1", result)
+
+    # Result has all the expected years
+    assert set(years_base + years_new) - {2015} == set(result["year_vtg"])
+
+    # With extrapolate=True
+    args["extrapolate"] = True
+    result = func(**args)
+    print("2", result)
+
+    # Result has all the expected years
+    assert set(years_base + years_new) == set(result["year_vtg"])
+
+    # With extrapolation that produces negative values before `horizon`, extrapol_neg
+    # has no effect
+    years_new = [2000, 2015, 2025, 2035, 2045, 2050]
+    args["yrs_new"] = years_new
+    args["extrapol_neg"] = 1.1
+
+    result = func(**args)
+    print("3", result)
+
+    # Result has all the expected years
+    assert set(years_base + years_new) == set(result["year_vtg"])
+
+    # Extrapolation that produces negative values after `horizon`
+    df.loc[:, "value"] = [30.0, 20.0, 1.5]
+    args["extrapol_neg"] = 1.1
+
+    result = func(**args)
+    print("4", result)
+
+    # Result has all the expected years
+    assert set(years_base + years_new) == set(result["year_vtg"])
+
+    # With bound_extend=False
+    args["bound_extend"] = False
+    result = func(**args)
+    print("5", result)

message_ix/tools/add_year/README.rst

@@ -4,7 +4,7 @@ Add model years to an existing Scenario
 Description
 -----------
 
-This tool adds new modeling years to an existing :class:`message_ix.Scenario` (hereafter "reference scenario"). For instance, in a scenario define with::
+This tool adds new modeling years to an existing :class:`message_ix.Scenario` (hereafter "reference scenario"). For instance, in a scenario defined with::
 
     history = [690]
     model_horizon = [700, 710, 720]
@@ -13,7 +13,7 @@ This tool adds new modeling years to an existing :class:`message_ix.Scenario` (h
         firstmodelyear=model_horizon[0]
     )
 
-…additional years can be added after importing the add_year function::
+…additional years can be added after importing the ``add_year function``::
 
     from message_ix.tools.add_year import add_year
     sc_new = message_ix.Scenario(mp, sc_ref.model, sc_ref.scenario,
@@ -26,7 +26,7 @@ At this point, ``sc_new`` will have the years [700, 705, 710, 712, 718, 720, 725
 The tool operates by creating a new empty Scenario (hereafter "new scenario") and:
 
 - Copying all **sets** from the reference scenario, adding new time steps to relevant sets (e.g., adding 2025 between 2020 and 2030 in the set ``year``)
-- Copying all **parameters** from the reference scenario, adding new years to relevant parameters, and calculating missing values for the added years.
+- Copying all **parameters** from the reference scenario, adding new years to relevant parameters, and calculating missing values for the added years
 
 Features
 ~~~~~~~~