From b2c3f611fd1902155f7860875634afc9158bb87d Mon Sep 17 00:00:00 2001
From: Paul Natsuo Kishimoto <mail@paul.kishimoto.name>
Date: Fri, 17 Feb 2023 14:44:58 +0100
Subject: [PATCH] Add compound_growth() genno computation, tests

---
 message_ix_models/report/computations.py      | 14 +++++++++
 .../tests/report/test_computations.py         | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 message_ix_models/tests/report/test_computations.py

diff --git a/message_ix_models/report/computations.py b/message_ix_models/report/computations.py
index 374cd58a44..99d2abc512 100644
--- a/message_ix_models/report/computations.py
+++ b/message_ix_models/report/computations.py
@@ -5,6 +5,7 @@
 
 import ixmp
 import pandas as pd
+from genno.computations import pow
 from iam_units import convert_gwp
 from iam_units.emissions import SPECIES
 from ixmp.reporting import Quantity
@@ -22,6 +23,19 @@
 ]
 
 
+def compound_growth(qty: Quantity, dim: str) -> Quantity:
+    """Compute compound growth along `dim` of `qty`."""
+    # Compute intervals along `dim`
+    # The value at index d is the duration between d and the next index d+1
+    c = qty.coords[dim]
+    dur = (c - c.shift({dim: 1})).fillna(0).shift({dim: -1})
+    # - Raise the values of `qty` to the power of the duration.
+    # - Compute cumulative product along `dim` from the first index.
+    # - Shift, so the value at index d is the growth relative to the prior index d-1
+    # - Fill in 1.0 for the first index.
+    return pow(qty, Quantity(dur)).cumprod(dim).shift({dim: 1}).fillna(1.0)
+
+
 def get_ts(
     scenario: ixmp.Scenario,
     filters: Optional[dict] = None,
diff --git a/message_ix_models/tests/report/test_computations.py b/message_ix_models/tests/report/test_computations.py
new file mode 100644
index 0000000000..1ff65b8abf
--- /dev/null
+++ b/message_ix_models/tests/report/test_computations.py
@@ -0,0 +1,30 @@
+import xarray as xr
+from genno import Quantity
+
+from message_data.reporting.computations import compound_growth
+
+
+def test_compound_growth():
+    """:func:`.compound_growth` on a 2-D quantity."""
+    qty = Quantity(
+        xr.DataArray(
+            [
+                [1.01, 1.0, 1.02, 1e6],  # Varying growth rates for x=x1
+                [1.0, 1.0, 1.0, 1.0],  # No rates/constant for x=x2
+            ],
+            coords=(["x1", "x2"], [2020, 2021, 2030, 2035]),
+            dims=("x", "t"),
+        )
+    )
+
+    # Function runs
+    result = compound_growth(qty, "t")
+
+    # Results have expected values
+    r1 = result.sel(x="x1")
+    assert all(1.0 == r1.sel(t=2020))
+    assert all(1.01 == r1.sel(t=2021) / r1.sel(t=2020))
+    assert all(1.0 == r1.sel(t=2030) / r1.sel(t=2021))
+    assert all(1.02**5 == r1.sel(t=2035) / r1.sel(t=2030))
+
+    assert all(1.0 == result.sel(x="x2"))