def test_reporting_key():
k1 = Key('foo', ['a', 'b', 'c'])
# Representation
assert repr(k1) == 'foo:a-b-c'
# Key hashes the same as its string representation
assert hash(k1) == hash('foo:a-b-c')
# Key compares equal to its string representation
assert k1 == 'foo:a-b-c'
# Number of partial sums for a 3-dimensional quantity
assert sum(1 for a in k1.iter_sums()) == 7
def test_missing_keys():
"""Adding computations that refer to missing keys raises KeyError."""
r = Reporter()
r.add('a', 3)
r.add('d', 4)
def gen(other):
"""A generator for apply()."""
return (lambda a, b: a * b, 'a', other)
# One missing key
with pytest.raises(KeyError, match=r"\['b'\]"):
r.add_product('ab', 'a', 'b')
# Two missing keys
with pytest.raises(KeyError, match=r"\['c', 'b'\]"):
r.add_product('abc', 'c', 'a', 'b')
# Using apply() targeted at non-existent keys also raises an Exception
with pytest.raises(KeyError, match=r"\['e', 'f'\]"):
r.apply(gen, 'd', 'e', 'f')
# add(..., strict=True) checks str or Key arguments
g = Key('g', 'hi')
with pytest.raises(KeyError, match=r"\['b', g:h-i\]"):
r.add('foo', (computations.product, 'a', 'b', g), strict=True)
# aggregate() and disaggregate() call add(), which raises the exception
with pytest.raises(KeyError, match=r"\[g:h-i\]"):
r.aggregate(g, 'tag', 'i')
with pytest.raises(KeyError, match=r"\[g:h-i\]"):
r.disaggregate(g, 'j')
def convert_pyam(self,
quantities,
year_time_dim,
tag='iamc',
drop={},
collapse=None):
"""Add conversion of one or more **quantities** to IAMC format.
Parameters
----------
quantities : str or Key or list of (str, Key)
Quantities to transform to :mod:`pyam`/IAMC format.
year_time_dim : str
Label of the dimension use for the ‘Year’ or ‘Time’ column of the
resulting :class:`pyam.IamDataFrame`. The column is labelled ‘Time’
if ``year_time_dim=='h'``, otherwise ‘Year’.
tag : str, optional
Tag to append to new Keys.
drop : iterable of str, optional
Label of additional dimensions to drop from the resulting data
frame. Dimensions ``h``, ``y``, ``ya``, ``yr``, and ``yv``—
except for the one named by `year_time_dim`—are automatically
dropped.
collapse : callable, optional
Callback to handle additional dimensions of the quantity. A
:class:`pandas.DataFrame` is passed as the sole argument to
`collapse`, which must return a modified dataframe.
Returns
-------
list of Key
Each key converts a :class:`Quantity
<ixmp.reporting.utils.Quantity>` into a :class:`pyam.IamDataFrame`.
See also
--------
message_ix.reporting.computations.as_pyam
"""
if isinstance(quantities, (str, Key)):
quantities = [quantities]
quantities = self.check_keys(*quantities)
keys = []
for qty in quantities:
# Dimensions to drop automatically
qty = Key.from_str_or_key(qty)
to_drop = set(drop) | set(
qty.dims) & ({'h', 'y', 'ya', 'yr', 'yv'} - {year_time_dim})
new_key = ':'.join([qty.name, tag])
comp = partial(computations.as_pyam,
year_time_dim=year_time_dim,
drop=to_drop,
collapse=collapse)
self.add(new_key, (comp, 'scenario', qty))
keys.append(new_key)
return keys
def test_reporter_add():
"""Adding computations that refer to missing keys raises KeyError."""
r = Reporter()
r.add('a', 3)
r.add('d', 4)
# Adding an existing key with strict=True
with pytest.raises(KeyExistsError, match=r"key 'a' already exists"):
r.add('a', 5, strict=True)
def gen(other):
"""A generator for apply()."""
return (lambda a, b: a * b, 'a', other)
def msg(*keys):
"""Return a regex for str(MissingKeyError(*keys))."""
return 'required keys {!r} not defined'.format(tuple(keys)) \
.replace('(', '\\(') \
.replace(')', '\\)')
# One missing key
with pytest.raises(MissingKeyError, match=msg('b')):
r.add_product('ab', 'a', 'b')
# Two missing keys
with pytest.raises(MissingKeyError, match=msg('c', 'b')):
r.add_product('abc', 'c', 'a', 'b')
# Using apply() targeted at non-existent keys also raises an Exception
with pytest.raises(MissingKeyError, match=msg('e', 'f')):
r.apply(gen, 'd', 'e', 'f')
# add(..., strict=True) checks str or Key arguments
g = Key('g', 'hi')
with pytest.raises(MissingKeyError, match=msg('b', g)):
r.add('foo', (computations.product, 'a', 'b', g), strict=True)
# aggregate() and disaggregate() call add(), which raises the exception
with pytest.raises(MissingKeyError, match=msg(g)):
r.aggregate(g, 'tag', 'i')
with pytest.raises(MissingKeyError, match=msg(g)):
r.disaggregate(g, 'j')
# add(..., sums=True) also adds partial sums
r.add('foo:a-b-c', [], sums=True)
assert 'foo:b' in r
# add(name, ...) where name is the name of a computation
r.add('select', 'bar', 'a', indexers={'dim': ['d0', 'd1', 'd2']})
# add(name, ...) with keyword arguments not recognized by the computation
# raises an exception
msg = "unexpected keyword argument 'bad_kwarg'"
with pytest.raises(TypeError, match=msg):
r.add('select', 'bar', 'a', bad_kwarg='foo', index=True)
def test_reporter_disaggregate():
r = Reporter()
foo = Key('foo', ['a', 'b', 'c'])
r.add(foo, '<foo data>')
r.add('d_shares', '<share data>')
# Disaggregation works
r.disaggregate(foo, 'd', args=['d_shares'])
assert 'foo:a-b-c-d' in r.graph
assert r.graph['foo:a-b-c-d'] == (computations.disaggregate_shares,
'foo:a-b-c', 'd_shares')
# Invalid method
with pytest.raises(ValueError):
r.disaggregate(foo, 'd', method='baz')
def test_prepare_plots(dantzig_reporter):
# Function runs without error
prepare_plots(dantzig_reporter)
# Plot keys are added; contains a task with 2 elements
(func, key) = dantzig_reporter.graph["plot new capacity"]
# First element is a callable partial object with certain keywords
assert callable(func)
assert func.keywords == dict(
dims=("nl", "t", "yv"),
units="GWa",
title="Energy System New Capacity",
cf=1.0,
stacked=True,
)
# Second element is a key
assert Key("CAP_NEW", ["nl", "t", "yv"]) == key
def test_reporting_key():
k1 = Key('foo', ['a', 'b', 'c'])
k2 = Key('bar', ['d', 'c', 'b'])
# String
assert str(k1) == 'foo:a-b-c'
# Representation
assert repr(k1) == '<foo:a-b-c>'
# Key hashes the same as its string representation
assert hash(k1) == hash('foo:a-b-c')
# Key compares equal to its string representation
assert k1 == 'foo:a-b-c'
# product:
assert Key.product('baz', k1, k2) == Key('baz', ['a', 'b', 'c', 'd'])
# iter_sums: Number of partial sums for a 3-dimensional quantity
assert sum(1 for a in k1.iter_sums()) == 7
# Key with name and tag but no dimensions
assert Key('foo', tag='baz') == 'foo::baz'
def convert_pyam(
self,
quantities,
year_time_dim,
tag="iamc",
drop={},
collapse=None,
unit=None,
replace_vars=None,
):
"""Add conversion of one or more **quantities** to IAMC format.
Parameters
----------
quantities : str or Key or list of (str, Key)
Quantities to transform to :mod:`pyam`/IAMC format.
year_time_dim : str
Label of the dimension use for the ‘Year’ or ‘Time’ column of the
resulting :class:`pyam.IamDataFrame`. The column is labelled ‘Time’
if ``year_time_dim=='h'``, otherwise ‘Year’.
tag : str, optional
Tag to append to new Keys.
drop : iterable of str, optional
Label of additional dimensions to drop from the resulting data
frame. Dimensions ``h``, ``y``, ``ya``, ``yr``, and ``yv``—
except for the one named by `year_time_dim`—are automatically
dropped.
collapse : callable, optional
Callback to handle additional dimensions of the quantity. A
:class:`pandas.DataFrame` is passed as the sole argument to
`collapse`, which must return a modified dataframe.
unit : str or pint.Unit, optional
Convert values to these units.
replace_vars : str or Key
Other reporting key containing a :class:`dict` mapping variable
names to replace.
Returns
-------
list of Key
Each key converts a :class:`Quantity
<ixmp.reporting.utils.Quantity>` into a :class:`pyam.IamDataFrame`.
See also
--------
message_ix.reporting.computations.as_pyam
"""
if isinstance(quantities, (str, Key)):
quantities = [quantities]
quantities = self.check_keys(*quantities)
keys = []
for qty in quantities:
# Key for the new quantity
qty = Key.from_str_or_key(qty)
new_key = ":".join([qty.name, tag])
# Dimensions to drop automatically
to_drop = set(drop) | set(
qty.dims) & ({"h", "y", "ya", "yr", "yv"} - {year_time_dim})
# Prepare the computation
comp = [
partial(
computations.as_pyam,
year_time_dim=year_time_dim,
drop=to_drop,
collapse=collapse,
unit=unit,
),
"scenario",
qty,
]
if replace_vars:
comp.append(replace_vars)
# Add and store
self.add(new_key, tuple(comp))
keys.append(new_key)
return keys
