Merge pull request #65 from khaeru/enh/2022-W29

Miscellaneous enhancements for 2022-W29

.gitignore

@@ -21,3 +21,6 @@ myfunc2-*.pkl
 doc/*.svg
 doc/.ipynb_checkpoints
 doc/_build
+
+# Editor/IDE-specific files
+/.vscode

doc/api.rst

@@ -299,6 +299,7 @@ Computations
       disaggregate_shares
       div
       group_sum
+      index_to
       interpolate
       mul
       pow

doc/config.rst

@@ -305,6 +305,9 @@ Each item contains:
   Refers to the name of a computation that is available in the namespace of :mod:`genno.computations`, or custom computations registered by compatibility modules or third-party packages.
   See :meth:`Computer.add` and :meth:`Computer.get_comp`.
   E.g. if "product", then :meth:`.Computer.add_product` is called, which also automatically infers the correct dimensions for each input.
+
+  If omitted, :data:`None`, or YAML ``null``, no specific callable is used, but instead ``key:`` is configured to retrieve a simple :class:`list`` of the ``inputs:``.
+  In this case, ``args:`` are ignored.
 ``key:``
    The key for the computed quantity.
 

doc/whatsnew.rst

@@ -9,7 +9,16 @@ What's new
 Next release
 ============
 
-- :meth:`.require_compat` can handle arbitrary module names and module objects (:pull:`63`).
+- New computation :func:`.index_to` (:pull:`65`).
+- :ref:`config-general` configuration items are more flexible (:pull:`65`).
+
+  - ``comp: null`` or omitted allows to specify a simple collection of other computations.
+  - A bare string ``key:`` is left as-is; only keys with (a) dimension(s) and/or tag are parsed to :class:`.Key`.
+
+- :func:`.repr` of :class:`.Quantity` displays its units (:pull:`65`).
+- Bug fix: :meth:`.Computer.convert_pyam` handles its `tag` argument correctly, generating keys like ``foo:x-y-z:iamc`` or ``bar::iamc`` when applied to existing keys like ``foo:x-y-z`` or ``bar`` (:pull:`65`).
+  Previously the generated keys would be e.g. ``bar:iamc``, which incorrectly treats "iamc" as a (sole) dimension rather than a tag.
+- :meth:`.Computer.require_compat` can handle arbitrary module names as strings, as well as module objects (:pull:`63`).
 
 v1.11.0 (2022-04-20)
 ====================

genno/computations.py

@@ -1,10 +1,10 @@
 """Elementary computations for genno."""
 # Notes:
-# - To avoid ambiguity, computations should not have default arguments. Define
-#   default values for the corresponding methods on the Computer class.
+# - To avoid ambiguity, computations should not have default arguments. Define default
+#   values for the corresponding methods on the Computer class.
 import logging
 from pathlib import Path
-from typing import Any, Hashable, Mapping, Union, cast
+from typing import Any, Hashable, Mapping, Optional, Union, cast
 
 import pandas as pd
 import pint
@@ -25,6 +25,7 @@
     "disaggregate_shares",
     "div",
     "group_sum",
+    "index_to",
     "interpolate",
     "load_file",
     "mul",
@@ -90,7 +91,7 @@ def add(*quantities, fill_value=0.0):
     return result
 
 
-def aggregate(quantity, groups, keep):
+def aggregate(quantity, groups: Mapping[Hashable, Mapping], keep: bool):
     """Aggregate *quantity* by *groups*.
 
     Parameters
@@ -129,7 +130,9 @@ def aggregate(quantity, groups, keep):
             values.append(agg)
 
         # Reassemble to a single dataarray
-        quantity = concat(*values, dim=dim)
+        quantity = concat(
+            *values, **({} if isinstance(quantity, AttrSeries) else {"dim": dim})
+        )
 
     # Preserve attrs
     quantity.attrs = attrs
@@ -402,6 +405,48 @@ def load_file(path, dims={}, units=None, name=None):
         return open(path).read()
 
 
+def index_to(
+    qty: Quantity,
+    dim_or_selector: Union[str, Mapping],
+    label: Optional[Hashable] = None,
+) -> Quantity:
+    """Compute an index of `qty` against certain of its values.
+
+    If the label is not provided, :func:`index_to` uses the label in the first position
+    along the identified dimension.
+
+    Parameters
+    ----------
+    qty : :class:`~genno.Quantity`
+    dim_or_selector : str or mapping
+        If a string, the ID of the dimension to index along.
+        If a mapping, it must have only one element, mapping a dimension ID to a label.
+    label : Hashable
+        Label to select along the dimension, required if `dim_or_selector` is a string.
+    Raises
+    ------
+    TypeError
+        if `dim_or_selector` is a mapping with length != 1.
+    """
+    if isinstance(dim_or_selector, Mapping):
+        if len(dim_or_selector) != 1:
+            raise TypeError(
+                f"Got {dim_or_selector}; expected a mapping from 1 key to 1 value"
+            )
+        dim, label = dict(dim_or_selector).popitem()
+    else:
+        # Unwrap dask.core.literals
+        dim = getattr(dim_or_selector, "data", dim_or_selector)
+        label = getattr(label, "data", label)
+
+    if label is None:
+        # Choose a label on which to normalize
+        label = qty.coords[dim][0].item()
+        log.info(f"Normalize quantity {qty.name} on {dim}={label}")
+
+    return div(qty, qty.sel({dim: label}))
+
+
 def pow(a, b):
     """Compute `a` raised to the power of `b`.
 

genno/config.py

@@ -2,7 +2,7 @@
 from copy import copy
 from functools import partial
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union
 
 import pint
 import yaml
@@ -236,25 +236,31 @@ def general(c: Computer, info):
         log.info(f"Add {repr(key)} using .add_product()")
     else:
         # The resulting key
-        key = Key.from_str_or_key(info["key"])
+        key = info["key"]
+        key = key if Key.bare_name(key) else Key.from_str_or_key(key)
 
         # Infer the dimensions of the resulting key if ":*:" is given for the dims
-        if set(key.dims) == {"*"}:
+        if set(getattr(key, "dims", {})) == {"*"}:
             key = Key.product(key.name, *inputs, tag=key.tag)
             # log.debug(f"Inferred dimensions ({', '.join(key.dims)}) for '*'")
 
-        # Retrieve the function for the computation
-        f = c.get_comp(info["comp"])
+        # If info["comp"] is None, the task is a list that collects other keys
+        _seq: Type = list
+        task = []
 
-        if f is None:
-            raise ValueError(info["comp"])
+        if info["comp"] is not None:
+            _seq = tuple  # Task is a computation
+            # Retrieve the function for the computation
+            f = c.get_comp(info["comp"])
+            if f is None:
+                raise ValueError(info["comp"])
+            task = [partial(f, **info.get("args", {}))]
 
-        log.info(f"Add {repr(key)} using {f.__module__}{f.__name__}(...)")
+            log.info(f"Add {repr(key)} using {f.__module__}.{f.__name__}(…)")
 
-        kwargs = info.get("args", {})
-        task = tuple([partial(f, **kwargs)] + list(inputs))
+        task.extend(inputs)
 
-        added = c.add(key, task, strict=True, sums=info.get("sums", False))
+        added = c.add(key, _seq(task), strict=True, sums=info.get("sums", False))
 
         if isinstance(added, tuple):
             log.info(f"    + {len(added)-1} partial sums")
@@ -265,7 +271,7 @@ def report(c: Computer, info):
     """Handle one entry from the ``report:`` config section."""
     log.info(f"Add report {info['key']} with {len(info['members'])} table(s)")
 
-    # Concatenate pyam data structures
+    # Concatenatqe pyam data structures
     c.add(info["key"], tuple([c.get_comp("concat")] + info["members"]), strict=True)
 
 

genno/core/attrseries.py

@@ -82,6 +82,9 @@ def __init__(self, data=None, *args, name=None, attrs=None, **kwargs):
         # Update the attrs after initialization
         self.attrs.update(attrs)
 
+    def __repr__(self):
+        return super().__repr__() + f", units: {self.units}"
+
     @classmethod
     def from_series(cls, series, sparse=None):
         """Like :meth:`xarray.DataArray.from_series`."""

genno/core/computer.py

@@ -142,7 +142,7 @@ def require_compat(self, pkg: Union[str, ModuleType]):
               :mod:`genno.compat`, e.g. :mod:`genno.compat.plotnine`.
               ``genno.compat.{pkg}.computations`` is added.
             - the name of an arbitary module, e.g. "foo.bar"
-            - a previously imported module.
+            - a previously imported module object.
 
         Raises
         ------
@@ -856,11 +856,11 @@ def convert_pyam(self, quantities, tag="iamc", **kwargs):
 
         keys = []
         for qty in quantities:
-            # Key for the input quantity
+            # Key for the input quantity, e.g. foo:x-y-z
             key = Key.from_str_or_key(qty)
 
-            # Key for the task
-            keys.append(":".join([key.name, tag]))
+            # Key for the task/output, e.g. foo::iamc
+            keys.append(Key(key.name, tag=tag))
 
             # Add the task and store the key
             self.add_single(keys[-1], (comp, "scenario", key))

genno/core/quantity.py

@@ -37,6 +37,10 @@ def from_series(cls, series, sparse=True):
         assert sparse
         return cls._get_class().from_series(series, sparse)
 
+    @property
+    def name(self) -> Hashable:
+        ...  # pragma: no cover
+
     @property
     def units(self):
         """Retrieve or set the units of the Quantity.

genno/core/sparsedataarray.py

@@ -168,7 +168,11 @@ def _item(self, *args):
         if len(args):  # pragma: no cover
             super().item(*args)
         elif len(self.data.shape) == 0:
-            return self.data.data[0]
+            return (
+                self.data.data[0]
+                if isinstance(self.data, sparse.COO)
+                else self.data.item()
+            )
         else:
             raise ValueError("can only convert an array of size 1 to a Python scalar")
 

genno/tests/compat/test_pyam.py

@@ -132,7 +132,7 @@ def add_tm(df, name="Activity"):
     key2 = c.convert_pyam(ACT, rename=rename, collapse=add_tm)
 
     # Keys of added node(s) are returned
-    assert ACT.name + ":iamc" == key2
+    assert ACT.name + "::iamc" == key2
 
     caplog.clear()
 

genno/tests/test_computations.py

@@ -210,6 +210,37 @@ def test_group_sum(ureg):
     assert 2 == len(result)
 
 
+def test_index_to():
+    q = random_qty(dict(x=3, y=5))
+    q.name = "Foo"
+
+    exp = q / q.sel(x="x0")
+    exp.units = ""
+
+    # Called with a mapping
+    result = computations.index_to(q, dict(x="x0"))
+    assert_qty_equal(exp, result)
+    assert exp.name == result.name
+
+    # Called with two positional arguments
+    result = computations.index_to(q, "x", "x0")
+    assert_qty_equal(exp, result)
+
+    # Default first index selected if 'None' is given
+    result = computations.index_to(q, "x")
+    assert_qty_equal(exp, result)
+
+    result = computations.index_to(q, dict(x=None))
+    assert_qty_equal(exp, result)
+
+    # Invalid calls
+    with pytest.raises(TypeError, match="expected a mapping from 1 key to 1 value"):
+        computations.index_to(q, dict(x="x0", y="y0"))  # Length != 1
+
+    with pytest.raises(KeyError):
+        computations.index_to(q, dict(x="x99"))  # Mapping to something invalid
+
+
 @pytest.mark.parametrize(
     "shape",
     [