diff --git a/message_ix_models/report/computations.py b/message_ix_models/report/computations.py index 20b5acfb2f..a00c34e044 100644 --- a/message_ix_models/report/computations.py +++ b/message_ix_models/report/computations.py @@ -172,18 +172,22 @@ def model_periods(y: List[int], cat_year: pd.DataFrame) -> List[int]: def remove_ts( scenario: ixmp.Scenario, - config: dict, + config: Optional[dict] = None, after: Optional[int] = None, dump: bool = False, ) -> None: """Remove all time series data from `scenario`. - .. todo:: Improve to provide the option to remove only those periods in the model - horizon. + Note that data stored with :meth:`.add_timeseries` using :py:`meta=True` as a + keyword argument cannot be removed using :meth:`.TimeSeries.remove_timeseries`, and + thus also not with this operator. - .. todo:: Move upstream, e.g. to :mod:`ixmp` alongside :func:`.store_ts`. + .. todo:: Move upstream, to :mod:`ixmp` alongside :func:`.store_ts`. """ - data = scenario.timeseries() + if dump: + raise NotImplementedError + + data = scenario.timeseries().drop("value", axis=1) N = len(data) count = f"{N}" @@ -203,9 +207,6 @@ def remove_ts( else: scenario.commit(f"Remove time series data ({__name__}.remove_all_ts)") - if dump: - raise NotImplementedError - # Non-weak references to objects to keep them alive _FROM_URL_REF: Set[Any] = set() diff --git a/message_ix_models/tests/report/test_computations.py b/message_ix_models/tests/report/test_computations.py index 2c7b4544e7..0f8cd68c70 100644 --- a/message_ix_models/tests/report/test_computations.py +++ b/message_ix_models/tests/report/test_computations.py @@ -3,9 +3,11 @@ import ixmp import message_ix import pandas as pd +import pandas.testing as pdt import pytest import xarray as xr from genno import Computer, Quantity +from ixmp.testing import assert_logs from message_ix.testing import make_dantzig from message_ix_models import ScenarioInfo @@ -90,8 +92,38 @@ def test_from_url(scenario): assert scenario.url == result.url -def test_get_ts(): - get_ts() +def test_get_remove_ts(caplog, scenario): + # get_ts() runs + result0 = get_ts(scenario) + pdt.assert_frame_equal(scenario.timeseries(), result0) + + # Can be used through a Computer + + c = Computer() + c.require_compat("message_ix_models.report.computations") + c.add("scenario", scenario) + + key = c.add("test1", "get_ts", "scenario", filters=dict(variable="GDP")) + result1 = c.get(key) + assert 3 == len(result1) + + # remove_ts() can be used through Computer + key = c.add("test2", "remove_ts", "scenario", "config", after=1964) + + # Task runs, logs + # NB this log message is incorrect, because ixmp's JDBCBackend is unable to delete + # data stored with "meta=True". Only 1 row is removed + with assert_logs(caplog, "Remove 2 of 6 (1964 <= year) rows of time series data"): + c.get(key) + + # See comment above; only one row is removed + assert 6 - 1 == len(scenario.timeseries()) + + # remove_ts() can be used directly + remove_ts(scenario) + + # All non-'meta' data were removed + assert 3 == len(scenario.timeseries()) def test_gwp_factors(): @@ -125,11 +157,6 @@ def test_model_periods(): assert 2020 == min(result) -@pytest.mark.xfail(reason="Incomplete") -def test_remove_ts(): - remove_ts() - - @pytest.mark.xfail(reason="Incomplete") def test_share_curtailment(): share_curtailment()