import logging
from contextlib import nullcontext
from copy import deepcopy
from functools import partial
from pathlib import Path
from re import escape
from typing import TYPE_CHECKING
from warnings import warn
import genno.config
import yaml
from genno import Key
from genno.compat.pyam import iamc as handle_iamc
from genno.core.key import single_key
from ixmp.util import discard_on_error
from message_ix import Reporter, Scenario
from message_ix_models import Context, ScenarioInfo
from message_ix_models.util._logging import mark_time, silence_log
from .config import Config
if TYPE_CHECKING:
from genno.core.key import KeyLike # TODO Import from genno.types
from .config import Callback
__all__ = [
"NOT_IMPLEMENTED_IAMC",
"NOT_IMPLEMENTED_MEASURE",
"Config",
"defaults",
"prepare_reporter",
"register",
"report",
]
log = logging.getLogger(__name__)
# Ignore a section in global.yaml used to define YAML anchors
try:
# genno ≥ 1.25
genno.config.handles("_iamc formats", False, False)(genno.config.store)
except AttributeError:
# genno < 1.25
# TODO Remove once the minimum supported version in message-ix-models is ≥ 1.25
@genno.config.handles("_iamc formats")
def _(c: Reporter, info):
pass
PE0 = r"Primary Energy\|(Coal|Gas|Hydro|Nuclear|Solar|Wind)"
PE1 = r"Primary Energy\|(Coal|Gas|Solar|Wind)"
E = (
r"Emissions\|CO2\|Energy\|Demand\|Transportation\|Road Rail and Domestic "
"Shipping"
)
#: Measures for which reporting is not implemented. See :data:`NOT_IMPLEMENTED_IAMC`.
NOT_IMPLEMENTED_MEASURE = {
"(Agricultural|Forestry) (Demand|Production)",
"Capacity Additions",
"Capacity",
"Capital Cost",
"Carbon Sequestration",
"Consumption",
r"Cost\|Cost Nodal Net",
"Cumulative Capacity",
"Efficiency",
r"Emissions\|(BC|CF4|CH4|CO|CO2|F-Gases|HFC|Kyoto Gases|N2O|NH3|NOx|OC|SF6|Sulfur)",
r"Emissions\|(VOC)",
"Fertilizer Use",
"Final Energy",
"Food Demand",
"GDP",
"GLOBIOM",
"Investment",
"Land Cover",
"Lifetime",
"OM Cost",
"Population",
"Price",
r"Resource\|Cumulative Extraction",
"Secondary Energy",
"Trade",
"Useful Energy",
"Water Consumption",
"Water Withdrawal",
"Yield",
}
#: Expressions for IAMC variable names not implemented by :func:`prepare_reporter`.
NOT_IMPLEMENTED_IAMC = [
# Other 'variable' codes are missing from `obs`
rf"variable='({'|'.join(NOT_IMPLEMENTED_MEASURE)})(\|.*)?': no right data",
# 'variable' codes with further parts are missing from `obs`
f"variable='{PE0}.*': no right data",
# For `pe1` (NB: not Hydro or Solar) units and most values differ
f"variable='{PE1}.*': units mismatch .*EJ/yr.*'', nan",
r"variable='Primary Energy|Coal': 220 of 240 values with \|diff",
r"variable='Primary Energy|Gas': 234 of 240 values with \|diff",
r"variable='Primary Energy|Solar': 191 of 240 values with \|diff",
r"variable='Primary Energy|Wind': 179 of 240 values with \|diff",
r"variable='Resource\|Extraction.*': \d{2,3} of 240 values with",
# For `e` units and most values differ
f"variable='{E}': units mismatch: .*Mt CO2/yr.*Mt / a",
rf"variable='{E}': 20 missing right entries",
rf"variable='{E}': 220 of 240 values with \|diff",
#
# Missing from data/test/report/snapshot-1.tar.gz
escape("variable='Resource|Extraction|Uranium': no left data"),
]
@genno.config.handles("iamc")
def iamc(c: Reporter, info):
"""Handle one entry from the ``iamc:`` config section.
This version overrides the version from :mod:`genno.config` to:
- Set some defaults for the `rename` argument for :meth:`.convert_pyam`:
- The `n` and `nl` dimensions are mapped to the "region" IAMC column.
- The `y`, `ya`, and `yv` dimensions are mapped to the "year" column.
- Use the MESSAGEix-GLOBIOM custom :func:`.util.collapse` callback to perform
renaming etc. while collapsing dimensions to the IAMC ones. The "var" key from
the entry, if any, is passed to the `var` argument of that function.
- Provide optional partial sums. The "sums" key of the entry can give a list of
strings such as ``["x", "y", "x-y"]``; in this case, the conversion to IAMC format
is also applied to the same "base" key with a partial sum over the dimension "x";
over "y", and over both "x" and "y". The corresponding dimensions are omitted from
"var". All data are concatenated.
"""
# FIXME the upstream key "variable" for the configuration is confusing; choose a
# better name
from message_ix_models.report.util import collapse
# Common
base_key = Key(info["base"])
# First part of the 'Variable' name
name = info.pop("variable", base_key.name)
# Parts (string literals or dimension names) to concatenate into variable name
var_parts = info.pop("var", [name])
# Use message_ix_models custom collapse() method
info.setdefault("collapse", {})
# Add standard renames
info.setdefault("rename", {})
for dim, target in (
("n", "region"),
("nl", "region"),
("y", "year"),
("ya", "year"),
("yv", "year"),
):
info["rename"].setdefault(dim, target)
# Iterate over partial sums
# TODO move some or all of this logic upstream
keys = [] # Resulting keys
for dims in [""] + info.pop("sums", []):
# Dimensions to partial
# TODO allow iterable of str
dims = dims.split("-")
label = f"{name} {'-'.join(dims) or 'full'}"
# Modified copy of `info` for this invocation
_info = info.copy()
# Base key: use the partial sum over any `dims`. Use a distinct variable name.
_info.update(base=base_key.drop(*dims), variable=label)
# Exclude any summed dimensions from the IAMC Variable to be constructed
_info["collapse"].update(
callback=partial(collapse, var=[v for v in var_parts if v not in dims])
)
# Invoke the genno built-in handler
handle_iamc(c, _info)
keys.append(f"{label}::iamc")
# Concatenate together the multiple tables
c.add("concat", f"{name}::iamc", *keys)
[docs]
def register(name_or_callback: "Callback | str") -> str | None:
"""Deprecated alias for :meth:`.report.Config.register`.
This version uses :meth:`Context.get_instance()` to get the 0-th Context, and calls
that method.
"""
warn(
"message_ix_models.report.register(…) function; use the method "
".report.Config.register(…) or Context.report.register(…) instead",
DeprecationWarning,
stacklevel=2,
)
return Context.get_instance(0).report.register(name_or_callback)
def log_before(context, rep, key) -> None:
log.info(f"Prepare to report {'(DRY RUN)' if context.dry_run else ''}")
log.info(key)
log.log(
logging.INFO
if (context.core.dry_run or context.core.verbose)
else logging.DEBUG,
"\n" + rep.describe(key),
)
mark_time()
[docs]
def report(context: Context, *args, **kwargs):
"""Report (post-process) solution data in a :class:`.Scenario`.
This function provides a single, common interface to call both the :mod:`genno`
-based (:mod:`message_ix_models.report`) and ‘legacy’ (
:mod:`message_ix_models.report.legacy`) reporting codes.
Parameters
----------
context : Context
The code responds to:
- :attr:`.dry_run`: if :obj:`True`, reporting is prepared but nothing is done.
- :attr:`~.Config.scenario_info` and :attr:`~.Config.platform_info`: used to
retrieve the Scenario to be reported.
- :py:`context.report`, which is an instance of :class:`.report.Config`; see
there for available configuration settings.
"""
# Handle deprecated usage that appears in:
# - .model.cli.new_baseline()
# - .model.create.solve()
# - .projects.covid.scenario_runner.ScenarioRunner.solve()
if isinstance(context, Scenario):
warn(
"Calling report(scenario, path, legacy=…); pass a Context instead",
category=DeprecationWarning,
)
# Ensure `context` is actually a Context object for the following code
scenario = context
context = Context.get_instance(-1)
# Transfer args, kwargs to context
context.set_scenario(scenario)
context.report.legacy.update(kwargs.pop("legacy", {}))
if len(args) + len(set(kwargs.keys()) & {"path"}) != 1:
raise TypeError(
f"Unknown mix of deprecated positional {args!r} "
f"and keyword arguments {kwargs!r}"
)
elif len(args) == 1:
out_dir = args[0]
else:
out_dir = kwargs.pop("path")
context.report.legacy.setdefault("out_dir", out_dir)
if context.report.legacy["use"]:
return _invoke_legacy_reporting(context)
with (
nullcontext()
if context.core.verbose
else silence_log("genno message_ix_models")
):
rep, key = prepare_reporter(context)
log_before(context, rep, key)
if context.dry_run:
return
with discard_on_error(rep.graph["scenario"]):
result = rep.get(key)
# Display information about the result
log.info(
f"File output(s), if any, written under:\n{rep.graph['config']['output_dir']}"
)
msg = ["Result"]
if context.core.verbose and isinstance(result, genno.Quantity):
msg.append(result.to_string())
else:
msg.append(str(result))
log.info("\n\n".join(msg + [""]))
def _invoke_legacy_reporting(context):
from .legacy import iamc_report_hackathon
log.info("Using .report.legacy.iamc_report_hackathon.report")
# Convert "legacy" config to kwargs for .legacy.iamc_report_hackathon.report()
kwargs = deepcopy(context.report.legacy)
kwargs.pop("use")
# Read a legacy reporting configuration file and update the arguments
config_file_path = kwargs.pop("config_file_path", None)
if isinstance(config_file_path, Path) and config_file_path.exists():
with open(config_file_path, "r") as f:
kwargs.update(yaml.safe_load(f))
# Retrieve the Scenario and Platform
scen = context.get_scenario()
mp = scen.platform
mark_time()
# `context` is passed only for the "dry_run" setting; the function receives all its
# other settings via the `kwargs`
return iamc_report_hackathon.report(mp=mp, scen=scen, context=context, **kwargs)
[docs]
def prepare_reporter(
context: Context,
scenario: Scenario | None = None,
reporter: Reporter | None = None,
) -> tuple[Reporter, "KeyLike | None"]:
"""Return a :class:`.Reporter` and `key` prepared to report a :class:`.Scenario`.
Every function returned by :attr:`.Config.iter_callbacks` is called, in order, to
allow each to populate the `reporter` with additional tasks.
Parameters
----------
context : .Context
The code responds to :py:`context.report`, which is an instance of
:class:`.report.Config`.
scenario : .Scenario, optional
Scenario to report. If not given, :meth:`.Context.get_scenario` is used to
retrieve a Scenario.
reporter : .Reporter, optional
Existing reporter to extend with computations. If not given, it is created
using :meth:`message_ix.Reporter.from_scenario`.
Returns
-------
.Reporter
Reporter prepared with tasks for reporting a MESSAGEix-GLOBIOM scenario; if
`reporter` is given, this is a reference to the same object.
If :attr:`.cli_output` is given, a task with the key "cli-output" is added that
writes the :attr:`.Config.key` to that path.
.Key
Same as :attr:`.Config.key` if any, but in full resolution; else either
"default" or "cli-output" according to the other settings.
"""
log.info("Prepare reporter")
log.debug(f".report.prepare_reporter: {context.regions = }")
if reporter:
# Existing `Reporter` provided
rep = reporter
has_solution = True
if scenario:
log.warning(f"{scenario = } argument ignored")
scenario = rep.graph["scenario"]
else:
# Retrieve the scenario
scenario = scenario or context.get_scenario()
# Create a new Reporter
rep = Reporter.from_scenario(scenario)
has_solution = scenario.has_solution()
# Append the message_data operators
rep.require_compat("message_ix_models.report.operator")
# Force re-installation of the function iamc() in this file as the handler for
# "iamc:" sections in global.yaml. Until message_data.reporting is removed, then
# importing it will cause the iamc() function in *that* file to override the one
# registered above.
# TODO Remove, once message_data.reporting is removed.
genno.config.handles("iamc")(iamc)
if context.report.use_scenario_path:
# Construct ScenarioInfo
si = ScenarioInfo(scenario, empty=True)
# Use the scenario URL to extend the path
assert context.report.output_dir
context.report.set_output_dir(context.report.output_dir.joinpath(si.path))
# Pass values to genno's configuration; deepcopy to protect from destructive
# operations
rep.configure(
**deepcopy(context.report.genno_config),
fail="raise" if has_solution else logging.NOTSET,
)
rep.configure(model=deepcopy(context.model))
# Apply callbacks for other modules which define additional reporting computations
for callback in context.report.iter_callbacks():
callback(rep, context)
key = context.report.key
if key:
# If just a bare name like "ACT" is given, infer the full key
if Key.bare_name(key):
inferred = rep.infer_keys(key)
if inferred != key:
log.info(f"Infer {inferred!r} for {key!r}")
key = inferred
if context.report.cli_output:
# Add a new task that writes `key` to the specified file
key = single_key(
rep.add(
"cli-output", "write_report", key, path=context.report.cli_output
)
)
elif rep.default_key:
key = rep.default_key
log.info(f"No key given; will use default: {key!r}")
else:
log.info("No key given and no default")
# Create the output directory
context.report.mkdir()
log.info("…done")
return rep, key
[docs]
def defaults(rep: Reporter, context: Context) -> None:
"""Prepare default contents and configuration of `rep` for MESSAGEix-GLOBIOM.
This includes:
- Populate :data:`.key.coords` and :data:`.key.groups`.
- Add a :func:`genno.operator.concat` task with no arguments at
:data:`.key.all_iamc`.
- Call :func:`.add_replacements` for members of the :ref:`commodity-yaml` and
:ref:`technology-yaml` code lists
"""
from message_ix_models.model.structure import get_codes
from . import key as k
from .util import add_replacements
# Add tasks to return coordinates for data manpulation, e.g. expand_dims, select
rep.add(k.coords.n_glb, "node_glb", "n")
# Add tasks to return groups of codes for aggregation
rep.add(k.groups.c, "get_commodity_groups")
# Add a placeholder task to concatenate IAMC-structured data
rep.add(k.all_iamc, "concat")
# Add mappings for conversions to IAMC data structures
add_replacements("c", get_codes("commodity"))
add_replacements("t", get_codes("technology"))