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)