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