Move unit conversion into IAMCConversion.add_tasks

Author: khaeru

message_ix_models/report/extraction.py

@@ -15,7 +15,7 @@
 
 #: Describe the conversion to IAMC structure.
 CONV = IAMCConversion(
-    base=K.all[2],
+    base=K.all[1],
     var_parts=["Resource|Extraction", "c"],
     sums=["c"],
     unit="EJ/yr",
@@ -25,8 +25,8 @@
 
 def callback(r: "Reporter", context: "Context") -> None:
     """Prepare reporting of resource extraction."""
-    # TODO Generalize the following 3 blocks. Similar operations will be used in other
-    # cases.
+    # TODO Generalize the following 3 calls to Reporter.add(), e.g. by a utility
+    #      function. Similar operations will be used in other cases.
 
     # Apply missing units to model contents
     r.add(K.hist[0], "apply_units", K.hist, units="GWa/year")
@@ -35,11 +35,8 @@ def callback(r: "Reporter", context: "Context") -> None:
     # Add historical and current values together
     r.add(K.all[0], "add", K.hist[0], K.model[0])
 
-    # Convert to target units
-    r.add(K.all[1], "convert_units", K.all[0], units=CONV.unit)
-
     # Aggregate on 'c' dimension using groups from commodity_groups()
-    r.add(K.all[2], "aggregate", K.all[1], groups.c, keep=False)
+    r.add(K.all[1], "aggregate", K.all[0], groups.c, keep=False)
 
     # Add tasks to (a) transform to IAMC-structured data, (b) concatenate to all::iamc
     CONV.add_tasks(r)

message_ix_models/report/util.py

@@ -120,8 +120,13 @@ def add_tasks(self, c: "Computer") -> None:
             c.add(k.glb[1], "expand_dims", k.glb[0], coords.n_glb)
 
             # Add zeros to base data & update the base key for next steps
-            k.base += "glb"
-            c.add(k.base, "add", self.base, k.glb[1])
+            c.add(k.base[0], "add", self.base, k.glb[1])
+        else:
+            # Simple alias
+            c.add(k.base[0], k.base)
+
+        # Convert to target units
+        c.add(k.base[1], "convert_units", k.base[0], units=self.unit, sums=True)
 
         # Common keyword arguments for genno.compat.pyam.iamc
         args: dict = dict(rename=_RENAME, unit=self.unit)
@@ -151,7 +156,7 @@ def add_tasks(self, c: "Computer") -> None:
                 c,
                 args
                 | dict(
-                    base=k.base.drop(*dims),
+                    base=k.base[1].drop(*dims),
                     variable=f"{label} {i}",
                     collapse=dict(callback=collapse, var=var_parts),
                 ),