def test_reporter_apply():
# Reporter with two scalar values
r = Reporter()
r.add('foo', (lambda x: x, 42))
r.add('bar', (lambda x: x, 11))
N = len(r.keys())
# A computation
def _product(a, b):
return a * b
# A generator function that yields keys and computations
def baz_qux(key):
yield key + ':baz', (_product, key, 0.5)
yield key + ':qux', (_product, key, 1.1)
# Apply the generator to two targets
r.apply(baz_qux, 'foo')
r.apply(baz_qux, 'bar')
# Four computations were added to the reporter
N += 4
assert len(r.keys()) == N
assert r.get('foo:baz') == 42 * 0.5
assert r.get('foo:qux') == 42 * 1.1
assert r.get('bar:baz') == 11 * 0.5
assert r.get('bar:qux') == 11 * 1.1
# A generator that takes two arguments
def twoarg(key1, key2):
yield key1 + '__' + key2, (_product, key1, key2)
r.apply(twoarg, 'foo:baz', 'bar:qux')
# One computation added to the reporter
N += 1
assert len(r.keys()) == N
assert r.get('foo:baz__bar:qux') == 42 * 0.5 * 11 * 1.1
# A useless generator that does nothing
def useless():
return
r.apply(useless)
# Nothing added to the reporter
assert len(r.keys()) == N
# Adding with a generator that takes Reporter as the first argument
def add_many(rep: Reporter, max=5):
[rep.add(f'foo{x}', _product, 'foo', x) for x in range(max)]
r.apply(add_many, max=10)
# Function was called, adding keys
assert len(r.keys()) == N + 10
# Keys work
assert r.get('foo9') == 42 * 9
def test_reporter_add_queue():
r = Reporter()
r.add('foo-0', (lambda x: x, 42))
# A computation
def _product(a, b):
return a * b
# A queue of computations to add. Only foo-1 succeeds on the first pass;
# only foo-2 on the second pass, etc.
strict = dict(strict=True)
queue = [
(('foo-4', _product, 'foo-3', 10), strict),
(('foo-3', _product, 'foo-2', 10), strict),
(('foo-2', _product, 'foo-1', 10), strict),
(('foo-1', _product, 'foo-0', 10), {}),
]
# Maximum 3 attempts → foo-4 fails on the start of the 3rd pass
with pytest.raises(MissingKeyError, match='foo-3'):
r.add(queue, max_tries=3, fail='raise')
# But foo-2 was successfully added on the second pass, and gives the
# correct result
assert r.get('foo-2') == 42 * 10 * 10
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 test_reporter_apply():
# Reporter with two scalar values
r = Reporter()
r.add('foo', (lambda x: x, 42))
r.add('bar', (lambda x: x, 11))
N = len(r.keys())
# A computation
def _product(a, b):
return a * b
# A generator method that yields keys and computations
def baz_qux(key):
yield key + ':baz', (_product, key, 0.5)
yield key + ':qux', (_product, key, 1.1)
# Apply the generator to two targets
r.apply(baz_qux, 'foo')
r.apply(baz_qux, 'bar')
# Four computations were added to the reporter
N += 4
assert len(r.keys()) == N
assert r.get('foo:baz') == 42 * 0.5
assert r.get('foo:qux') == 42 * 1.1
assert r.get('bar:baz') == 11 * 0.5
assert r.get('bar:qux') == 11 * 1.1
# A generator that takes two arguments
def twoarg(key1, key2):
yield key1 + '__' + key2, (_product, key1, key2)
r.apply(twoarg, 'foo:baz', 'bar:qux')
# One computation added to the reporter
N += 1
assert len(r.keys()) == N
assert r.get('foo:baz__bar:qux') == 42 * 0.5 * 11 * 1.1
# A useless generator that does nothing
def useless(key):
return
r.apply(useless, 'foo:baz__bar:qux')
# Nothing added to the reporter
assert len(r.keys()) == N
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_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
def test_reporter_full_key():
r = Reporter()
# Without index, the full key cannot be retrieved
r.add('a:i-j-k', [])
with pytest.raises(KeyError, match='a'):
r.full_key('a')
# Using index=True adds the full key to the index
r.add('a:i-j-k', [], index=True)
assert r.full_key('a') == 'a:i-j-k'
# The full key can be retrieved by giving only some of the indices
assert r.full_key('a:j') == 'a:i-j-k'
# Same with a tag
r.add('a:i-j-k:foo', [], index=True)
# Original and tagged key can both be retrieved
assert r.full_key('a') == 'a:i-j-k'
assert r.full_key('a::foo') == 'a:i-j-k:foo'
def test_reporting_units():
"""Test handling of units within Reporter computations."""
r = Reporter()
# Create some dummy data
dims = dict(coords=['a b c'.split()], dims=['x'])
r.add('energy:x',
Quantity(xr.DataArray([1., 3, 8], attrs={'_unit': 'MJ'}, **dims)))
r.add('time',
Quantity(xr.DataArray([5., 6, 8], attrs={'_unit': 'hour'}, **dims)))
r.add('efficiency', Quantity(xr.DataArray([0.9, 0.8, 0.95], **dims)))
# Aggregation preserves units
r.add('energy', (computations.sum, 'energy:x', None, ['x']))
assert r.get('energy').attrs['_unit'] == ureg.parse_units('MJ')
# Units are derived for a ratio of two quantities
r.add('power', (computations.ratio, 'energy:x', 'time'))
assert r.get('power').attrs['_unit'] == ureg.parse_units('MJ/hour')
# Product of dimensioned and dimensionless quantities keeps the former
r.add('energy2', (computations.product, 'energy:x', 'efficiency'))
assert r.get('energy2').attrs['_unit'] == ureg.parse_units('MJ')
def add_many(rep: Reporter, max=5):
[rep.add(f'foo{x}', _product, 'foo', x) for x in range(max)]
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 (f'required keys {repr(tuple(keys))} not defined'.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)
