Add interpolate_capacity function

Author: tsmbland

src/muse/agents/agent.py

@@ -297,7 +297,7 @@ def next(
         """
         from logging import getLogger
 
-        from muse.utilities import reduce_assets
+        from muse.utilities import interpolate_capacity, reduce_assets
 
         # Check inputs
         assert len(market.year) == 2
@@ -340,9 +340,10 @@ def next(
         search = search.sel(asset=condtechs)
 
         # Calculate capacity in current and investment year
-        capacity = reduce_assets(
-            self.assets.capacity, coords=("technology", "region")
-        ).interp(year=[current_year, investment_year])
+        capacity = interpolate_capacity(
+            reduce_assets(self.assets.capacity, coords=("technology", "region")),
+            year=[current_year, investment_year],
+        )
 
         # Calculate constraints
         constraints = self.constraints(

src/muse/demand_share.py

@@ -246,21 +246,27 @@ def new_and_retro(
 
     from muse.commodities import is_enduse
     from muse.quantities import maximum_production
-    from muse.utilities import agent_concatenation, broadcast_over_assets, reduce_assets
+    from muse.utilities import (
+        agent_concatenation,
+        broadcast_over_assets,
+        interpolate_capacity,
+        reduce_assets,
+    )
 
     current_year, investment_year = map(int, market.year.values)
 
     def decommissioning(capacity, technologies):
         return decommissioning_demand(
             technologies=technologies,
-            capacity=capacity.interp(
-                year=[current_year, investment_year], kwargs={"fill_value": 0.0}
+            capacity=interpolate_capacity(
+                capacity, year=[current_year, investment_year]
             ),
             timeslice_level=timeslice_level,
         )
 
-    capacity = reduce_assets([u.assets.capacity for u in agents]).interp(
-        year=[current_year, investment_year], kwargs={"fill_value": 0.0}
+    capacity = interpolate_capacity(
+        reduce_assets([u.assets.capacity for u in agents]),
+        year=[current_year, investment_year],
     )
 
     # Select technodata for assets
@@ -288,8 +294,8 @@ def decommissioning(capacity, technologies):
     id_to_share: MutableMapping[Hashable, xr.DataArray] = {}
     for region in demands.region.values:
         retro_capacity: MutableMapping[Hashable, xr.DataArray] = {
-            agent.uuid: agent.assets.capacity.interp(
-                year=[current_year, investment_year]
+            agent.uuid: interpolate_capacity(
+                agent.assets.capacity, year=[current_year, investment_year]
             )
             for agent in agents
             if agent.category == "retrofit" and agent.region == region
@@ -382,15 +388,20 @@ def standard_demand(
 
     from muse.commodities import is_enduse
     from muse.quantities import maximum_production
-    from muse.utilities import agent_concatenation, broadcast_over_assets, reduce_assets
+    from muse.utilities import (
+        agent_concatenation,
+        broadcast_over_assets,
+        interpolate_capacity,
+        reduce_assets,
+    )
 
     current_year, investment_year = map(int, market.year.values)
 
     def decommissioning(capacity, technologies):
         return decommissioning_demand(
             technologies=technologies,
-            capacity=capacity.interp(
-                year=[current_year, investment_year], kwargs={"fill_value": 0.0}
+            capacity=interpolate_capacity(
+                capacity, year=[current_year, investment_year]
             ),
             timeslice_level=timeslice_level,
         )
@@ -401,8 +412,9 @@ def decommissioning(capacity, technologies):
             raise RetrofitAgentInStandardDemandShare()
 
     # Calculate existing capacity
-    capacity = reduce_assets([agent.assets.capacity for agent in agents]).interp(
-        year=[current_year, investment_year], kwargs={"fill_value": 0.0}
+    capacity = interpolate_capacity(
+        reduce_assets([agent.assets.capacity for agent in agents]),
+        year=[current_year, investment_year],
     )
 
     # Select technodata for assets
@@ -425,8 +437,8 @@ def decommissioning(capacity, technologies):
     for region in demands.region.values:
         # Calculate current capacity
         current_capacity: MutableMapping[Hashable, xr.DataArray] = {
-            agent.uuid: agent.assets.capacity.interp(
-                year=[current_year, investment_year]
+            agent.uuid: interpolate_capacity(
+                agent.assets.capacity, year=[current_year, investment_year]
             )
             for agent in agents
             if agent.region == region
@@ -481,16 +493,21 @@ def unmet_forecasted_demand(
 ) -> xr.DataArray:
     """Forecast demand that cannot be serviced by non-decommissioned current assets."""
     from muse.commodities import is_enduse
-    from muse.utilities import broadcast_over_assets, reduce_assets
+    from muse.utilities import (
+        broadcast_over_assets,
+        interpolate_capacity,
+        reduce_assets,
+    )
 
     current_year, investment_year = map(int, market.year.values)
 
     comm_usage = technologies.comm_usage.sel(commodity=market.commodity)
     smarket: xr.Dataset = market.where(is_enduse(comm_usage), 0)
 
     # Calculate existing capacity
-    capacity = reduce_assets([agent.assets.capacity for agent in agents]).interp(
-        year=[current_year, investment_year], kwargs={"fill_value": 0.0}
+    capacity = interpolate_capacity(
+        reduce_assets([agent.assets.capacity for agent in agents]),
+        year=[current_year, investment_year],
     )
 
     # Select data for future years

src/muse/sectors/sector.py

@@ -275,10 +275,10 @@ def market_variables(self, market: xr.Dataset, technologies: xr.Dataset) -> Any:
         from muse.commodities import is_pollutant
         from muse.costs import levelized_cost_of_energy, supply_cost
         from muse.quantities import consumption
-        from muse.utilities import broadcast_over_assets
+        from muse.utilities import broadcast_over_assets, interpolate_capacity
 
         years = market.year.values
-        capacity = self.capacity.interp(year=years)
+        capacity = interpolate_capacity(self.capacity, year=years)
 
         # Select technology data for each asset
         # Each asset uses the technology data from the year it was installed
@@ -350,7 +350,7 @@ def capacity(self) -> xr.DataArray:
         dimensions: asset (technology, installation date,
         region), year.
         """
-        from muse.utilities import reduce_assets
+        from muse.utilities import interpolate_capacity, reduce_assets
 
         traded = [
             u.assets.capacity
@@ -372,13 +372,14 @@ def capacity(self) -> xr.DataArray:
             ]
             flat_list = [item for sublist in full_list for item in sublist]
             years = sorted(list(set(flat_list)))
-            return reduce_assets(
+            capacity = reduce_assets(
                 [
                     u.assets.capacity
                     for u in self.agents
                     if "dst_region" not in u.assets.capacity.dims
                 ]
-            ).interp(year=years, kwargs={"fill_value": 0.0})
+            )
+            return interpolate_capacity(capacity, year=years)
 
         # Only traded assets
         elif not nontraded:
@@ -389,25 +390,27 @@ def capacity(self) -> xr.DataArray:
             ]
             flat_list = [item for sublist in full_list for item in sublist]
             years = sorted(list(set(flat_list)))
-            return reduce_assets(
+            capacity = reduce_assets(
                 [
                     u.assets.capacity
                     for u in self.agents
                     if "dst_region" in u.assets.capacity.dims
                 ]
-            ).interp(year=years, kwargs={"fill_value": 0.0})
+            )
+            return interpolate_capacity(capacity, year=years)
 
         # Both traded and nontraded assets
         else:
             traded_results = reduce_assets(traded)
             nontraded_results = reduce_assets(nontraded)
-            return reduce_assets(
+            capacity = reduce_assets(
                 [
                     traded_results,
                     nontraded_results
                     * (nontraded_results.region == traded_results.dst_region),
                 ]
             )
+            return interpolate_capacity(capacity, year=years)
 
     @property
     def agents(self) -> Iterator[AbstractAgent]:

src/muse/utilities.py

@@ -355,21 +355,20 @@ def lexical_comparison(
 def merge_assets(
     capa_a: xr.DataArray,
     capa_b: xr.DataArray,
-    interpolation: str = "linear",
     dimension: str = "asset",
 ) -> xr.DataArray:
     """Merge two capacity arrays."""
     # Interpolate capacity arrays to a common time framework
     years = sorted(set(capa_a.year.values).union(capa_b.year.values))
     if len(capa_a.year) == 1:
         capa_a_interp = capa_a
-        capa_b_interp = capa_b.interp(year=years, method=interpolation).fillna(0)
+        capa_b_interp = interpolate_capacity(capa_b, year=years)
     elif len(capa_b.year) == 1:
-        capa_a_interp = capa_a.interp(year=years, method=interpolation).fillna(0)
+        capa_a_interp = interpolate_capacity(capa_a, year=years)
         capa_b_interp = capa_b
     else:
-        capa_a_interp = capa_a.interp(year=years, method=interpolation).fillna(0)
-        capa_b_interp = capa_b.interp(year=years, method=interpolation).fillna(0)
+        capa_a_interp = interpolate_capacity(capa_a, year=years)
+        capa_b_interp = interpolate_capacity(capa_b, year=years)
 
     # Concatenate the two capacity arrays
     result = xr.concat((capa_a_interp, capa_b_interp), dim=dimension)
@@ -402,6 +401,21 @@ def avoid_repetitions(data: xr.DataArray, dim: str = "year") -> xr.DataArray:
     return data.year[years]
 
 
+def interpolate_capacity(
+    data: xr.DataArray, year: int | Sequence[int] | xr.DataArray
+) -> xr.DataArray:
+    """Interpolates capacity data to the given years.
+
+    Capacity between years is interpolated linearly. Capacity beyond the final year is
+    set to zero.
+    """
+    return data.interp(
+        year=year,
+        method="linear",
+        kwargs={"fill_value": 0.0},
+    )
+
+
 def nametuple_to_dict(nametup: Mapping | NamedTuple) -> Mapping:
     """Transforms a nametuple of type GenericDict into an OrderDict."""
     from collections import OrderedDict

tests/test_constraints.py

@@ -6,7 +6,7 @@
 from pytest import approx, fixture
 
 from muse.timeslices import drop_timeslice
-from muse.utilities import reduce_assets
+from muse.utilities import interpolate_capacity, reduce_assets
 
 CURRENT_YEAR = 2020
 INVESTMENT_YEAR = 2025
@@ -50,8 +50,9 @@ def assets(residential):
 
 @fixture
 def capacity(assets):
-    return reduce_assets(assets.capacity, coords=("technology", "region")).interp(
-        year=[CURRENT_YEAR, INVESTMENT_YEAR], method="linear"
+    return interpolate_capacity(
+        reduce_assets(assets.capacity, coords=("technology", "region")),
+        year=[CURRENT_YEAR, INVESTMENT_YEAR],
     )
 
 

tests/test_demand_share.py

@@ -2,14 +2,15 @@
 from pytest import approx, fixture, raises
 
 from muse.timeslices import drop_timeslice
+from muse.utilities import interpolate_capacity
 
 CURRENT_YEAR = 2010
 INVESTMENT_YEAR = 2015
 
 
 @fixture
 def _capacity(stock):
-    return stock.capacity.interp(year=[CURRENT_YEAR, INVESTMENT_YEAR])
+    return interpolate_capacity(stock.capacity, year=[CURRENT_YEAR, INVESTMENT_YEAR])
 
 
 @fixture

tests/test_utilities.py

@@ -204,7 +204,7 @@ def test_lexical_nobin(order):
 def test_merge_assets():
     from numpy import arange
 
-    from muse.utilities import merge_assets
+    from muse.utilities import interpolate_capacity, merge_assets
 
     def fake(year, order=("installed", "technology")):
         result = xr.Dataset()
@@ -244,17 +244,17 @@ def fake(year, order=("installed", "technology")):
         ab_side = actual.sel(
             asset=((actual.installed == inst) & (actual.technology == tech))
         ).squeeze("asset")
-        a_side = (
-            capa_a.sel(asset=((capa_a.installed == inst) & (capa_a.technology == tech)))
-            .sum("asset")
-            .interp(year=ab_side.year, method="linear")
-            .fillna(0)
+        a_side = interpolate_capacity(
+            capa_a.sel(
+                asset=((capa_a.installed == inst) & (capa_a.technology == tech))
+            ).sum("asset"),
+            year=ab_side.year,
         )
-        b_side = (
-            capa_b.sel(asset=((capa_b.installed == inst) & (capa_b.technology == tech)))
-            .sum("asset")
-            .interp(year=ab_side.year, method="linear")
-            .fillna(0)
+        b_side = interpolate_capacity(
+            capa_b.sel(
+                asset=((capa_b.installed == inst) & (capa_b.technology == tech))
+            ).sum("asset"),
+            year=ab_side.year,
         )
         assert (ab_side.capacity == approx((a_side + b_side).values)).all()