From b3c35ad49785a5580c6834278ca19a87b600bf5c Mon Sep 17 00:00:00 2001
From: Paul Natsuo Kishimoto <mail@paul.kishimoto.name>
Date: Mon, 12 Aug 2024 10:51:15 +0200
Subject: [PATCH] Add .models.shift_period; tests

---
 message_ix/message.py           | 64 +++++++++++++++++++++++++++++++++
 message_ix/tests/test_models.py | 25 +++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/message_ix/message.py b/message_ix/message.py
index 135e7ff05..614eb6b1a 100644
--- a/message_ix/message.py
+++ b/message_ix/message.py
@@ -1,6 +1,7 @@
 import logging
 from collections.abc import MutableMapping
 from functools import partial
+from typing import TYPE_CHECKING
 from warnings import warn
 
 import ixmp.model.gams
@@ -18,6 +19,10 @@
     _item_shorthand,
 )
 
+if TYPE_CHECKING:
+    from message_ix.core import Scenario
+
+
 log = logging.getLogger(__name__)
 
 
@@ -865,3 +870,62 @@ def run(self, scenario: "ixmp.Scenario") -> None:
 )
 equ("TOTAL_CAPACITY_BOUND_LO", "n inv_tec y", "Lower bound on total installed capacity")
 equ("TOTAL_CAPACITY_BOUND_UP", "n inv_tec y", "Upper bound on total installed capacity")
+
+
+def shift_period(scenario: "Scenario", y0: int) -> None:
+    """Shift the first period of the model horizon."""
+    from ixmp.backend.jdbc import JDBCBackend
+
+    # Retrieve existing cat_year information, including the current 'firstmodelyear'
+    cat_year = scenario.set("cat_year")
+    y0_pre = cat_year.query("type_year == 'firstmodelyear'")["year"].item()
+
+    if y0 == y0_pre:
+        log.info(f"First model period is already {y0!r}")
+        return
+    elif y0 < y0_pre:
+        raise NotImplementedError(
+            f"Shift first model period *earlier*, from {y0_pre!r} -> {y0}"
+        )
+
+    # Periods to be shifted from within to before the model horizon
+    periods = list(
+        filter(lambda y: y0_pre <= y < y0, map(int, sorted(cat_year["year"].unique())))
+    )
+    log.info(f"Shift data for period(s): {periods}")
+
+    # Handle historical_* parameters for which the dimensions are a subset of the
+    # corresponding variable's dimensions
+    data = {}
+    for var_name, par_name, filter_dim in (
+        ("ACT", "historical_activity", "year_act"),
+        ("CAP_NEW", "historical_new_capacity", "year_vtg"),
+        ("EXT", "historical_extraction", "year"),
+        ("GDP", "historical_gdp", "year"),
+        ("LAND", "historical_land", "year"),
+    ):
+        # - Filter data for `var_name` along the `filter_dim`, keeping only the periods
+        #   to be shifted.
+        # - Drop the marginal column; rename the level column to "value".
+        # - Group according to the dimensions of the target `par_name`.
+        # - Sum within groups.
+        # - Restore index columns.
+        data[par_name] = (
+            scenario.var(var_name, filters={filter_dim: periods})
+            .drop("mrg", axis=1)
+            .rename(columns={"lvl": "value"})
+            .groupby(list(MESSAGE.items[par_name].dims))
+            .sum()["value"]
+            .reset_index()
+        )
+
+    # TODO Handle "EMISS:n-e-type_tec-y" →
+    # "historical_emission:n-type_emission-type_tec-type_year", in which dimension names
+    # are changed
+
+    # TODO Adjust cat_year
+
+    if isinstance(scenario.platform._backend, JDBCBackend):
+        raise NotImplementedError("Cannot set variable values with JDBCBackend")
+
+    # TODO Store new data
diff --git a/message_ix/tests/test_models.py b/message_ix/tests/test_models.py
index 7f79be583..2d261867e 100644
--- a/message_ix/tests/test_models.py
+++ b/message_ix/tests/test_models.py
@@ -1,7 +1,14 @@
 from contextlib import nullcontext
+from typing import TYPE_CHECKING
 
 import pytest
 
+from message_ix.message import shift_period
+from message_ix.testing import make_dantzig
+
+if TYPE_CHECKING:
+    from ixmp import Platform
+
 
 @pytest.mark.parametrize(
     "name",
@@ -27,3 +34,21 @@ def test_deprecated_import(name: str) -> None:
 
     with ctx:
         getattr(message_ix.models, name)
+
+
+@pytest.mark.parametrize(
+    "y0",
+    (
+        # Not implemented: shifting to an earlier period
+        pytest.param(1962, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        # Does nothing
+        1963,
+        # Not implemented with ixmp.JDBCBackend
+        pytest.param(1964, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        pytest.param(1965, marks=pytest.mark.xfail(raises=NotImplementedError)),
+    ),
+)
+def test_shift_period(test_mp: "Platform", y0: int) -> None:
+    s = make_dantzig(test_mp, solve=True, multi_year=True)
+
+    shift_period(s, y0)