From 9d89c97bf671dbafa75ed70d8f118a5e83a5d2a1 Mon Sep 17 00:00:00 2001 From: Fridolin Glatter Date: Wed, 29 Nov 2023 14:28:12 +0100 Subject: [PATCH] Rework a sample test for the water module * Use test_context as in other tests * Type hint Context in the code * Fix other small test-related issues --- message_ix_models/model/water/__init__.py | 3 +- message_ix_models/model/water/build.py | 11 +- message_ix_models/model/water/cli.py | 14 +- .../model/water/data/__init__.py | 6 +- message_ix_models/model/water/data/demands.py | 12 +- .../model/water/data/infrastructure.py | 14 +- .../model/water/data/irrigation.py | 6 +- .../model/water/data/water_for_ppl.py | 17 ++- .../model/water/data/water_supply.py | 15 ++- message_ix_models/model/water/reporting.py | 5 +- message_ix_models/model/water/utils.py | 7 +- .../tests/model/water/test_cooling.py | 121 +++++++++--------- .../tests/model/water/test_utils.py | 17 +-- 13 files changed, 141 insertions(+), 107 deletions(-) diff --git a/message_ix_models/model/water/__init__.py b/message_ix_models/model/water/__init__.py index c68e1257fd..46583ecf73 100644 --- a/message_ix_models/model/water/__init__.py +++ b/message_ix_models/model/water/__init__.py @@ -1,3 +1,4 @@ +from .data import demands, water_supply from .utils import read_config -__all__ = ["read_config"] +__all__ = ["demands", "read_config", "water_supply"] diff --git a/message_ix_models/model/water/build.py b/message_ix_models/model/water/build.py index 465e1bb7fb..b32f836106 100644 --- a/message_ix_models/model/water/build.py +++ b/message_ix_models/model/water/build.py @@ -1,6 +1,6 @@ import logging from functools import lru_cache, partial -from typing import Mapping +from typing import TYPE_CHECKING, Mapping import pandas as pd from sdmx.model.v21 import Code @@ -12,10 +12,13 @@ from .utils import read_config +if TYPE_CHECKING: + from message_ix_models import Context + log = logging.getLogger(__name__) -def get_spec(context) -> Mapping[str, ScenarioInfo]: +def get_spec(context: Context) -> Mapping[str, ScenarioInfo]: """Return the specification for nexus implementation Parameters @@ -190,7 +193,7 @@ def generate_set_elements(set_name, match=None): return results -def map_basin(context) -> Mapping[str, ScenarioInfo]: +def map_basin(context: Context) -> Mapping[str, ScenarioInfo]: """Return specification for mapping basins to regions The basins are spatially consolidated from HydroSHEDS basins delineation @@ -242,7 +245,7 @@ def map_basin(context) -> Mapping[str, ScenarioInfo]: return dict(require=require, remove=remove, add=add) -def main(context, scenario, **options): +def main(context: Context, scenario, **options): """Set up MESSAGEix-Nexus on `scenario`. See also diff --git a/message_ix_models/model/water/cli.py b/message_ix_models/model/water/cli.py index a58b3cca6f..63fbe8b68a 100644 --- a/message_ix_models/model/water/cli.py +++ b/message_ix_models/model/water/cli.py @@ -1,10 +1,14 @@ import logging +from typing import TYPE_CHECKING import click from message_ix_models.model.structure import get_codes from message_ix_models.util.click import common_params +if TYPE_CHECKING: + from message_ix_models import Context + log = logging.getLogger(__name__) @@ -13,12 +17,12 @@ @common_params("regions") @click.option("--time", help="Manually defined time") @click.pass_obj -def cli(context, regions, time): +def cli(context: "Context", regions, time): """MESSAGEix-Water and Nexus variant.""" water_ini(context, regions, time) -def water_ini(context, regions, time): +def water_ini(context: "Context", regions, time): """Add components of the MESSAGEix-Nexus module This function modifies model name & scenario name @@ -61,7 +65,7 @@ def water_ini(context, regions, time): context.regions = regions # create a mapping ISO code : - # region name, for other scripts + # a region name, for other scripts # only needed for 1-country models nodes = get_codes(f"node/{context.regions}") nodes = list(map(str, nodes[nodes.index("World")].child)) @@ -106,7 +110,7 @@ def water_ini(context, regions, time): help="Defines whether the model solves with macro", ) @common_params("regions") -def nexus_cli(context, regions, rcps, sdgs, rels, macro=False): +def nexus_cli(context: "Context", regions, rcps, sdgs, rels, macro=False): """ Add basin structure connected to the energy sector and water balance linking different water demands to supply. @@ -115,7 +119,7 @@ def nexus_cli(context, regions, rcps, sdgs, rels, macro=False): nexus(context, regions, rcps, sdgs, rels, macro) -def nexus(context, regions, rcps, sdgs, rels, macro=False): +def nexus(context: "Context", regions, rcps, sdgs, rels, macro=False): """Add basin structure connected to the energy sector and water balance linking different water demands to supply. diff --git a/message_ix_models/model/water/data/__init__.py b/message_ix_models/model/water/data/__init__.py index 931af079e1..0ccb9da269 100644 --- a/message_ix_models/model/water/data/__init__.py +++ b/message_ix_models/model/water/data/__init__.py @@ -1,6 +1,7 @@ """Generate input data.""" import logging +from typing import TYPE_CHECKING from message_ix_models import ScenarioInfo from message_ix_models.util import add_par_data @@ -11,6 +12,9 @@ from .water_for_ppl import cool_tech, non_cooling_tec from .water_supply import add_e_flow, add_water_supply +if TYPE_CHECKING: + from message_ix_models import Context + log = logging.getLogger(__name__) DATA_FUNCTIONS = [ @@ -40,7 +44,7 @@ ] -def add_data(scenario, context, dry_run=False): +def add_data(scenario, context: "Context", dry_run=False): """Populate `scenario` with MESSAGEix-Nexus data.""" info = ScenarioInfo(scenario) diff --git a/message_ix_models/model/water/data/demands.py b/message_ix_models/model/water/data/demands.py index 95df9878f0..821d678b77 100644 --- a/message_ix_models/model/water/data/demands.py +++ b/message_ix_models/model/water/data/demands.py @@ -1,6 +1,7 @@ """Prepare data for adding demands""" import os +from typing import TYPE_CHECKING import numpy as np import pandas as pd @@ -9,6 +10,9 @@ from message_ix_models.util import broadcast, package_data_path +if TYPE_CHECKING: + from message_ix_models import Context + def get_basin_sizes(basin, node): """Returns the sizes of developing and developed basins for a given node""" @@ -141,7 +145,7 @@ def target_rate_trt(df, basin): return df -def add_sectoral_demands(context): +def add_sectoral_demands(context: "Context"): """ Adds water sectoral demands Parameters @@ -699,7 +703,7 @@ def add_sectoral_demands(context): return results -def read_water_availability(context): +def read_water_availability(context: "Context"): """ Reads water availability data and bias correct it for the historical years and no climate @@ -824,7 +828,7 @@ def read_water_availability(context): return df_sw, df_gw -def add_water_availability(context): +def add_water_availability(context: "Context"): """ Adds water supply constraints Parameters @@ -892,7 +896,7 @@ def add_water_availability(context): return results -def add_irrigation_demand(context): +def add_irrigation_demand(context: "Context"): """ Adds endogenous irrigation water demands from GLOBIOM emulator Parameters diff --git a/message_ix_models/model/water/data/infrastructure.py b/message_ix_models/model/water/data/infrastructure.py index 92d1c801f0..68e87752a7 100644 --- a/message_ix_models/model/water/data/infrastructure.py +++ b/message_ix_models/model/water/data/infrastructure.py @@ -2,6 +2,7 @@ treatment in urban & rural""" from collections import defaultdict +from typing import TYPE_CHECKING import pandas as pd from message_ix import make_df @@ -15,8 +16,11 @@ same_time, ) +if TYPE_CHECKING: + from message_ix_models import Context -def add_infrastructure_techs(context): # noqa: C901 + +def add_infrastructure_techs(context: "Context"): """Process water distribution data for a scenario instance. Parameters ---------- @@ -39,8 +43,7 @@ def add_infrastructure_techs(context): # noqa: C901 # load the scenario from context scen = context.get_scenario() - year_wat = [2010, 2015] - year_wat.extend(info.Y) + year_wat = (2010, 2015, *info.Y) # first activity year for all water technologies is 2020 first_year = scen.firstmodelyear @@ -547,7 +550,7 @@ def add_infrastructure_techs(context): # noqa: C901 return results -def add_desalination(context): +def add_desalination(context: "Context"): """Add desalination infrastructure Two types of desalination are considered; 1. Membrane @@ -572,8 +575,7 @@ def add_desalination(context): # load the scenario from context scen = context.get_scenario() - year_wat = [2010, 2015] - year_wat.extend(info.Y) + year_wat = (2010, 2015, *info.Y) # first activity year for all water technologies is 2020 first_year = scen.firstmodelyear diff --git a/message_ix_models/model/water/data/irrigation.py b/message_ix_models/model/water/data/irrigation.py index 8a8e50e778..d8cf67259b 100644 --- a/message_ix_models/model/water/data/irrigation.py +++ b/message_ix_models/model/water/data/irrigation.py @@ -1,13 +1,17 @@ """Prepare data for water use for cooling & energy technologies.""" +from typing import TYPE_CHECKING import pandas as pd from message_ix import make_df from message_ix_models.util import broadcast, package_data_path +if TYPE_CHECKING: + from message_ix_models import Context + # water & electricity for irrigation -def add_irr_structure(context): +def add_irr_structure(context: "Context"): """Add irrigation withdrawal infrastructure The irrigation demands are added in Parameters diff --git a/message_ix_models/model/water/data/water_for_ppl.py b/message_ix_models/model/water/data/water_for_ppl.py index 2bbfec746c..e1350f2f7b 100644 --- a/message_ix_models/model/water/data/water_for_ppl.py +++ b/message_ix_models/model/water/data/water_for_ppl.py @@ -1,5 +1,7 @@ """Prepare data for water use for cooling & energy technologies.""" +from typing import TYPE_CHECKING, Any + import numpy as np import pandas as pd from message_ix import make_df @@ -12,9 +14,12 @@ same_node, ) +if TYPE_CHECKING: + from message_ix_models import Context + # water & electricity for cooling technologies -def cool_tech(context): # noqa: C901 +def cool_tech(context: "Context"): """Process cooling technology data for a scenario instance. The input values of parent technologies are read in from a scenario instance and then cooling fractions are calculated by using the data from @@ -363,7 +368,7 @@ def cooling_fr(x): ].drop_duplicates() search_cols_cooling_fraction = [col for col in search_cols if col != "technology"] - def shares(x, context): + def shares(x, context: "Context"): """Process share and cooling fraction. Returns ------- @@ -379,7 +384,7 @@ def shares(x, context): ]["cooling_fraction"] x[col] = x[col] * cooling_fraction - results = [] + results: list[Any] = [] for i in x: if isinstance(i, str): results.append(i) @@ -397,7 +402,7 @@ def shares(x, context): hold_cost = cost[search_cols].apply(shares, axis=1, context=context) hold_cost = hold_cost[hold_cost["technology"] != "delme"] - def hist_act(x, context): + def hist_act(x, context: "Context"): """Calculate historical activity of cooling technology. The data for shares is read from ``cooltech_cost_and_shares_ssp_msg.csv`` Returns @@ -447,7 +452,7 @@ def hist_act(x, context): # dataframe for historical activities of cooling techs act_value_df = pd.DataFrame(changed_value_series_flat, columns=columns) - def hist_cap(x, context): + def hist_cap(x, context: "Context"): """Calculate historical capacity of cooling technology. The data for shares is read from ``cooltech_cost_and_shares_ssp_msg.csv`` Returns @@ -741,7 +746,7 @@ def hist_cap(x, context): # Water use & electricity for non-cooling technologies -def non_cooling_tec(context): +def non_cooling_tec(context: "Context"): """Process data for water usage of power plants (non-cooling technology related). Water withdrawal values for power plants are read in from ``tech_water_performance_ssp_msg.csv`` diff --git a/message_ix_models/model/water/data/water_supply.py b/message_ix_models/model/water/data/water_supply.py index 5605ea778d..d4a2a25e8b 100644 --- a/message_ix_models/model/water/data/water_supply.py +++ b/message_ix_models/model/water/data/water_supply.py @@ -1,5 +1,7 @@ """Prepare data for water use for cooling & energy technologies.""" +from typing import TYPE_CHECKING + import numpy as np import pandas as pd from message_ix import make_df @@ -8,8 +10,11 @@ from message_ix_models.model.water.utils import map_yv_ya_lt from message_ix_models.util import broadcast, package_data_path, same_node, same_time +if TYPE_CHECKING: + from message_ix_models import Context + -def map_basin_region_wat(context): +def map_basin_region_wat(context: "Context"): """ Calculate share of water avaialbility of basins per each parent region. @@ -112,7 +117,7 @@ def map_basin_region_wat(context): return df_sw -def add_water_supply(context): +def add_water_supply(context: "Context"): """Add Water supply infrastructure This function links the water supply based on different settings and options. It defines the supply linkages for freshwater, groundwater and salinewater. @@ -135,9 +140,9 @@ def add_water_supply(context): # load the scenario from context scen = context.get_scenario() - year_wat = [2010, 2015] + # year_wat = (2010, 2015) fut_year = info.Y - year_wat.extend(info.Y) + year_wat = (2010, 2015, *info.Y) sub_time = context.time # first activity year for all water technologies is 2020 @@ -688,7 +693,7 @@ def add_water_supply(context): return results -def add_e_flow(context): +def add_e_flow(context: "Context"): """Add environmental flows This function bounds the available water and allocates the environmental flows.Environmental flow bounds are calculated using Variable Monthly Flow diff --git a/message_ix_models/model/water/reporting.py b/message_ix_models/model/water/reporting.py index 2e73d89966..ce67f35975 100644 --- a/message_ix_models/model/water/reporting.py +++ b/message_ix_models/model/water/reporting.py @@ -260,7 +260,6 @@ def multiply_electricity_output_of_hydro(elec_hydro_var, report_iam): # TODO -# flake8: noqa: C901 def report(sc=False, reg="", sdgs=False): """Report nexus module results""" @@ -1372,8 +1371,8 @@ def report(sc=False, reg="", sdgs=False): and country_n in group["region"].values ): report_pd.drop(group.index, inplace=True) - # Step 4: Rename "world" to "country" and remove rows with - # region = "country" + # Step 4: Rename "world" to "country" and remove rows + # with region = "country" group = group[group["region"] == "World"] group.loc[group["region"] == "World", "region"] = country_n # Step 5: Update the original dataframe with the modified group diff --git a/message_ix_models/model/water/utils.py b/message_ix_models/model/water/utils.py index 5c2547ce0c..72043498ec 100644 --- a/message_ix_models/model/water/utils.py +++ b/message_ix_models/model/water/utils.py @@ -2,7 +2,7 @@ from collections import defaultdict from functools import lru_cache from itertools import product -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple import numpy as np import pandas as pd @@ -13,6 +13,9 @@ from message_ix_models.model.structure import get_codes from message_ix_models.util import load_package_data +if TYPE_CHECKING: + from message_ix_models import Context + log = logging.getLogger(__name__) # Configuration files @@ -24,7 +27,7 @@ ] -def read_config(context=None): +def read_config(context: Context | None = None): """Read the water model configuration / metadata from file. Numerical values are converted to computation-ready data structures. diff --git a/message_ix_models/tests/model/water/test_cooling.py b/message_ix_models/tests/model/water/test_cooling.py index 425fabb5e2..0e192435a8 100644 --- a/message_ix_models/tests/model/water/test_cooling.py +++ b/message_ix_models/tests/model/water/test_cooling.py @@ -1,68 +1,71 @@ -from unittest.mock import MagicMock, patch - -import pandas as pd +from message_ix import Scenario from message_ix_models.model.water.data.water_for_ppl import non_cooling_tec -def test_non_cooling_tec(): - # Mock the context - context = { - "water build info": {"Y": [2020, 2030, 2040]}, - "type_reg": "country", - "regions": "test_region", - "map_ISO_c": {"test_region": "test_ISO"}, - "get_scenario": MagicMock( - return_value=MagicMock( - par=MagicMock( - return_value=pd.DataFrame( - { - "technology": ["tech1", "tech2"], - "node_loc": ["loc1", "loc2"], - "node_dest": ["dest1", "dest2"], - "year_vtg": ["2020", "2020"], - "year_act": ["2020", "2020"], - } - ) - ) - ) - ), +def test_non_cooling_tec(test_context): + context = test_context + mp = context.get_platform() + scenario_info = { + "mp": mp, + "model": "test water model", + "scenario": "test water scenario", + "version": "new", } + s = Scenario(**scenario_info) + s.add_horizon(year=[2020, 2030, 2040]) + s.add_set("technology", ["tech1", "tech2"]) + + # TODO: this is where you would add + # "node_loc": ["loc1", "loc2"], + # "node_dest": ["dest1", "dest2"], + # "year_vtg": ["2020", "2020"], + # "year_act": ["2020", "2020"], etc + # to the scenario as per usual. However, I don't know if that's necesarry as the + # test is passing without it, too. + + s.commit(comment="basic water test model") + + # set_scenario() updates Context.scenario_info + context.set_scenario(s) + print(context.get_scenario()) + context["water build info"] = {"Y": [2020, 2030, 2040]} + context.type_reg = "country" + context.regions = "test_region" + context.map_ISO_c = {"test_region": "test_ISO"} + + # TODO: only leaving this in so you can see which data you might want to assert to + # be in the result. Please remove after adapting the assertions below: # Mock the DataFrame read from CSV - df = pd.DataFrame( - { - "technology_group": ["cooling", "non-cooling"], - "technology_name": ["cooling_tech1", "non_cooling_tech1"], - "water_supply_type": ["freshwater_supply", "freshwater_supply"], - "water_withdrawal_mid_m3_per_output": [1, 2], - } - ) + # df = pd.DataFrame( + # { + # "technology_group": ["cooling", "non-cooling"], + # "technology_name": ["cooling_tech1", "non_cooling_tech1"], + # "water_supply_type": ["freshwater_supply", "freshwater_supply"], + # "water_withdrawal_mid_m3_per_output": [1, 2], + # } + # ) - # Mock the function 'private_data_path' to return the mocked DataFrame - with patch( - "message_ix_models.util.private_data_path", return_value="path/to/file" - ), patch("pandas.read_csv", return_value=df): - # Call the function to be tested - result = non_cooling_tec(context) + result = non_cooling_tec(context) - # Assert the results - assert isinstance(result, dict) - assert "input" in result - assert all( - col in result["input"].columns - for col in [ - "technology", - "value", - "unit", - "level", - "commodity", - "mode", - "time", - "time_origin", - "node_origin", - "node_loc", - "year_vtg", - "year_act", - ] - ) + # Assert the results + assert isinstance(result, dict) + assert "input" in result + assert all( + col in result["input"].columns + for col in [ + "technology", + "value", + "unit", + "level", + "commodity", + "mode", + "time", + "time_origin", + "node_origin", + "node_loc", + "year_vtg", + "year_act", + ] + ) diff --git a/message_ix_models/tests/model/water/test_utils.py b/message_ix_models/tests/model/water/test_utils.py index d12fa59880..226431564b 100644 --- a/message_ix_models/tests/model/water/test_utils.py +++ b/message_ix_models/tests/model/water/test_utils.py @@ -2,7 +2,7 @@ import pandas as pd import xarray as xr -from sdmx.model.v21 import Code +from sdmx.model.common import Code from message_ix_models import Context from message_ix_models.model.water.utils import ( @@ -13,9 +13,9 @@ ) -def test_read_config(): +def test_read_config(test_context): # Mock the context - context = Context(0) + context = test_context # Mock the data returned by load_private_data mock_data = {"test_key": "test_value"} @@ -33,9 +33,6 @@ def test_read_config(): def test_map_add_on(): - # Mock the context - Context(0) - # Mock the data returned by read_config mock_data = { "water set": { @@ -78,9 +75,9 @@ def test_add_commodity_and_level(): data=[ Code( id="tech1", - anno={"input": {"commodity": "com1", "level": "lev1"}}, + annotations=["input", "commodity", "com1", "level", "lev1"], ), - Code(id="tech2", anno={"input": {"commodity": "com2"}}), + Code(id="tech2", annotations=["input", "commodity", "com2"]), ], name="tech", ) @@ -89,8 +86,8 @@ def test_add_commodity_and_level(): } mock_codes_data = pd.Series( data=[ - Code(id="com1", anno={"level": "lev1"}), - Code(id="com2", anno={"level": "lev2"}), + Code(id="com1", annotations=["level", "lev1"]), + Code(id="com2", annotations=["level", "lev2"]), ], name="com", )