Skip to content

Commit

Permalink
Add compound_growth() genno computation, tests
Browse files Browse the repository at this point in the history
  • Loading branch information
khaeru committed Aug 30, 2023
1 parent e52c72a commit b2c3f61
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
14 changes: 14 additions & 0 deletions message_ix_models/report/computations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
30 changes: 30 additions & 0 deletions message_ix_models/tests/report/test_computations.py
Original file line number Diff line number Diff line change
@@ -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"))

0 comments on commit b2c3f61

Please sign in to comment.