Skip to content

Commit

Permalink
Add .project.ssp.data with a ExoDataSource for SSP data
Browse files Browse the repository at this point in the history
  • Loading branch information
khaeru committed Sep 13, 2023
1 parent 87d86b5 commit 2446271
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 0 deletions.
47 changes: 47 additions & 0 deletions message_ix_models/project/ssp/data.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
from message_ix_models.tools.exo_data import (
ExoDataSource,
iamc_like_data_for_query,
register_source,
)
from message_ix_models.util import private_data_path


@register_source
class SSPUpdate(ExoDataSource):
"""Provider of exogenous data from the SSP Update database."""

id = "SSP update"

def __init__(self, source, source_kw):
s = "ICONICS:SSP(2024)."
if not source.startswith(s):
raise ValueError(source)

*parts, self.ssp_number = source.partition(s)

# Map the `measure` keyword to a string appearing in the data
self.measure = {
"GDP": "GDP|PPP",
"POP": "Population",
}[source_kw.pop("measure")]

# Store the model ID, if any
self.model = source_kw.pop("model", None)

if len(source_kw):
raise ValueError(source_kw)

def __call__(self):
# Assemble a query string
query = " and ".join(
[
f"Scenario == 'SSP{self.ssp_number} - Review Phase 1'",
f"Variable == '{self.measure}'",
f"Model == '{self.model}'" if self.model else "True",
]
)

path = private_data_path("ssp", "SSP-Review-Phase-1.csv.gz")
assert path.exists(), "TODO handle the case where message_data is not insalled"

return iamc_like_data_for_query(path, query)
37 changes: 37 additions & 0 deletions message_ix_models/tests/project/test_ssp.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import pytest
from genno import Computer

from message_ix_models.project.ssp import (
SSP,
Expand All @@ -8,6 +9,8 @@
parse,
ssp_field,
)
from message_ix_models.project.ssp.data import SSPUpdate # noqa: F401
from message_ix_models.tools.exo_data import prepare_computer


def test_generate(tmp_path, test_context):
Expand Down Expand Up @@ -81,3 +84,37 @@ class Foo:

def test_cli(mix_models_cli):
mix_models_cli.assert_exit_0(["ssp", "gen-structures", "--dry-run"])


class TestSSPUpdate:
@pytest.mark.parametrize(
"source",
(
"ICONICS:SSP(2024).1",
"ICONICS:SSP(2024).2",
"ICONICS:SSP(2024).3",
"ICONICS:SSP(2024).4",
"ICONICS:SSP(2024).5",
),
)
@pytest.mark.parametrize(
"source_kw",
(
dict(measure="POP"),
dict(measure="GDP", model="IIASA GDP 2023"),
dict(measure="GDP", model="OECD ENV-Growth 2023"),
),
)
def test_prepare_computer(self, test_context, source, source_kw):
c = Computer()
keys = prepare_computer(test_context, c, source, source_kw)

# Preparation of data runs successfully
result = c.get(keys[0])

# Data has the expected dimensions
assert ("n", "y") == result.dims

# Data is complete
assert 14 == len(result.coords["n"])
assert 14 == len(result.coords["y"])

0 comments on commit 2446271

Please sign in to comment.