Merge branch 'develop' into rc2

Author: tsmbland

docs/advanced-guide/extending-muse.ipynb

@@ -99,7 +99,6 @@
     "def consumption_zero(\n",
     "    market: Dataset,\n",
     "    capacity: DataArray,\n",
-    "    technologies: Dataset,\n",
     "):\n",
     "    \"\"\"Current consumption.\"\"\"\n",
     "    result = (\n",
@@ -451,7 +450,6 @@
     "def consumption_zero(  # noqa: F811\n",
     "    market: Dataset,\n",
     "    capacity: DataArray,\n",
-    "    technologies: Dataset,\n",
     "    sum_over: Optional[list[str]] = None,\n",
     "    drop: Optional[list[str]] = None,\n",
     "    rounding: int = 4,\n",

src/muse/mca.py

@@ -13,7 +13,7 @@
 
 from muse.outputs.cache import OutputCache
 from muse.readers import read_initial_market
-from muse.sectors import SECTORS_REGISTERED, AbstractSector
+from muse.sectors import SECTORS_REGISTERED, AbstractSector, Sector
 from muse.timeslices import drop_timeslice
 from muse.utilities import future_propagation
 
@@ -342,6 +342,11 @@ def run(self) -> None:
 
             _, new_market, self.sectors = self.find_equilibrium(new_market)
 
+            # Save sector outputs
+            for sector in self.sectors:
+                if type(sector) is Sector:
+                    sector.save_outputs()
+
             # If we need to account for the carbon budget, we might need to change
             # the budget for the future, too.
             if check_carbon_budget and shoots and year_idx < nyear - 2:

src/muse/outputs/sector.py

@@ -9,7 +9,6 @@
     def quantity(
         capacity: xr.DataArray,
         market: xr.Dataset,
-        technologies: xr.Dataset
     ) -> Union[xr.DataArray, DataFrame]:
         pass
 
@@ -142,7 +141,6 @@ def factory(
 def capacity(
     market: xr.Dataset,
     capacity: xr.DataArray,
-    technologies: xr.Dataset,
     rounding: int = 4,
 ) -> pd.DataFrame:
     """Current capacity."""
@@ -179,7 +177,6 @@ def market_quantity(
 def consumption(
     market: xr.Dataset,
     capacity: xr.DataArray,
-    technologies: xr.Dataset,
     sum_over: Optional[list[str]] = None,
     drop: Optional[list[str]] = None,
     rounding: int = 4,
@@ -200,7 +197,6 @@ def consumption(
 def supply(
     market: xr.Dataset,
     capacity: xr.DataArray,
-    technologies: xr.Dataset,
     sum_over: Optional[list[str]] = None,
     drop: Optional[list[str]] = None,
     rounding: int = 4,
@@ -221,7 +217,6 @@ def supply(
 def costs(
     market: xr.Dataset,
     capacity: xr.DataArray,
-    technologies: xr.Dataset,
     sum_over: Optional[list[str]] = None,
     drop: Optional[list[str]] = None,
     rounding: int = 4,

src/muse/sectors/sector.py

@@ -151,6 +151,8 @@ def __init__(
         It can be anything registered with
         :py:func:`@register_production<muse.production.register_production>`.
         """
+        self.output_data: xr.Dataset
+        """Full supply, consumption and costs data for the most recent year."""
 
     @property
     def forecast(self):
@@ -232,29 +234,19 @@ def group_assets(x: xr.DataArray) -> xr.DataArray:
                 technologies, market, time_period=time_period, current_year=current_year
             )
 
-        # > output to mca
+        # Full output data
         supply, consume, costs = self.market_variables(market, technologies)
-
-        output_data = xr.Dataset(
+        self.output_data = xr.Dataset(
             dict(
                 supply=supply,
                 consumption=consume,
                 costs=costs,
             )
         )
 
-        # < output to mca
-        self.outputs(output_data, self.capacity, technologies)
-
+        # Output data for MCA (aggregated over assets)
         if len(supply.region.dims) == 0:
-            output_data = xr.Dataset(
-                dict(
-                    supply=supply,
-                    consumption=consume,
-                    costs=costs,
-                )
-            )
-            output_data = output_data.sum("asset")
+            output_data = self.output_data.sum("asset")
             output_data = output_data.expand_dims(region=[output_data.region.values])
         else:
             output_data = xr.Dataset(
@@ -265,9 +257,7 @@ def group_assets(x: xr.DataArray) -> xr.DataArray:
                 )
             )
 
-        # > to mca timeslices
         result = output_data.copy(deep=True)
-
         if "dst_region" in result:
             exclude = ["dst_region", "commodity", "year", "timeslice"]
             prices = market.prices.expand_dims(dst_region=market.prices.region.values)
@@ -299,9 +289,12 @@ def group_assets(x: xr.DataArray) -> xr.DataArray:
         result = self.convert_market_timeslice(result, mca_market.timeslice)
         result["comm_usage"] = technologies.comm_usage.sel(commodity=result.commodity)
         result.set_coords("comm_usage")
-        # < to mca timeslices
         return result
 
+    def save_outputs(self) -> None:
+        """Calls the outputs function with the current output data."""
+        self.outputs(self.output_data, self.capacity)
+
     def market_variables(self, market: xr.Dataset, technologies: xr.Dataset) -> Any:
         """Computes resulting market: production, consumption, and costs."""
         from muse.commodities import is_pollutant