Xarray patch: Check array dimensions (#647)

* Check array dimensions after multiplication

* Fix tests in costs module

* Fix objectives tests

* Fix demand_share tests

* Fix quantities tests, delete gross_margin

* Fix most remaining tests

* Remove some unnecessary casting

* Doctest, hardcode "asset"

* Better docstring

* Fix remaining test (?)

* Error message and better docstring

* Fix doctest

Author: tsmbland

src/muse/__main__.py

@@ -102,6 +102,22 @@ def patched_broadcast_compat_data(self, other):
         self_data = self.data
         other_data = other
         dims = self.dims
+
+    # Check output dimensions
+    if "asset" in dims and any(
+        dim in dims for dim in ["region", "technology", "installed"]
+    ):
+        raise ValueError(
+            "DataArrays with an 'asset' dimension cannot be broadcasted along "
+            "'region', 'technology', or 'installed' dimensions. "
+            "This error is usually raised when attempting to combine asset-level data "
+            "(e.g. a capacity dataset with an 'asset' dimension) with a fully explicit "
+            "technology dataset (e.g. a technology dataset with 'region' and "
+            "'technology' dimensions). "
+            "Please use `broadcast_techs` on the latter object before performing this "
+            "operation."
+        )
+
     return self_data, other_data, dims
 
 

src/muse/objectives.py

@@ -179,7 +179,7 @@ def decorated_objective(
             demand, ["asset", "timeslice", "commodity"], optional=["region"]
         )
         check_dimensions(
-            technologies, ["replacement", "commodity"], optional=["timeslice"]
+            technologies, ["replacement", "commodity"], optional=["timeslice", "asset"]
         )
 
         # Calculate objective

src/muse/quantities.py

@@ -9,8 +9,6 @@
 
 from __future__ import annotations
 
-from typing import cast
-
 import numpy as np
 import xarray as xr
 
@@ -143,98 +141,15 @@ def emission(
     from muse.utilities import broadcast_techs
 
     # just in case we are passed a technologies dataset, like in other functions
-    fouts = broadcast_techs(
-        getattr(fixed_outputs, "fixed_outputs", fixed_outputs), production
-    )
+    fixed_outputs = getattr(fixed_outputs, "fixed_outputs", fixed_outputs)
+    fouts = broadcast_techs(fixed_outputs, production)
     envs = is_pollutant(fouts.comm_usage)
     enduses = is_enduse(fouts.comm_usage)
     return production.sel(commodity=enduses).sum("commodity") * broadcast_timeslice(
         fouts.sel(commodity=envs), level=timeslice_level
     )
 
 
-def gross_margin(
-    technologies: xr.Dataset,
-    capacity: xr.DataArray,
-    prices: xr.Dataset,
-    timeslice_level: str | None = None,
-) -> xr.DataArray:
-    """The percentage of revenue after direct expenses have been subtracted.
-
-    .. _reference:
-    https://www.investopedia.com/terms/g/grossmargin.asp
-    We first calculate the revenues, which depend on prices
-    We then deduct the direct expenses
-    - energy commodities INPUTS are related to fuel costs
-    - environmental commodities OUTPUTS are related to environmental costs
-    - variable costs is given as technodata inputs
-    - non-environmental commodities OUTPUTS are related to revenues.
-    """
-    from muse.commodities import is_enduse, is_pollutant
-    from muse.utilities import broadcast_techs
-
-    tech = broadcast_techs(  # type: ignore
-        cast(
-            xr.Dataset,
-            technologies[
-                [
-                    "technical_life",
-                    "interest_rate",
-                    "var_par",
-                    "var_exp",
-                    "fixed_outputs",
-                    "fixed_inputs",
-                ]
-            ],
-        ),
-        capacity,
-    )
-
-    var_par = tech.var_par
-    var_exp = tech.var_exp
-    fixed_outputs = tech.fixed_outputs
-    fixed_inputs = tech.fixed_inputs
-    # We separate the case where we have one or more regions
-    caparegions = np.array(capacity.region.values).reshape(-1)
-    if len(caparegions) > 1:
-        prices.sel(region=capacity.region)
-    else:
-        prices = prices.where(prices.region == capacity.region, drop=True)
-    prices = prices.interp(year=capacity.year.values)
-
-    # Filters for pollutants and output commodities
-    environmentals = is_pollutant(technologies.comm_usage)
-    enduses = is_enduse(technologies.comm_usage)
-
-    # Variable costs depend on factors such as labour
-    variable_costs = distribute_timeslice(
-        var_par * ((fixed_outputs.sel(commodity=enduses)).sum("commodity")) ** var_exp,
-        level=timeslice_level,
-    )
-
-    # The individual prices are selected
-    # costs due to consumables, direct inputs
-    consumption_costs = (
-        prices * distribute_timeslice(fixed_inputs, level=timeslice_level)
-    ).sum("commodity")
-    # costs due to pollutants
-    production_costs = prices * distribute_timeslice(
-        fixed_outputs, level=timeslice_level
-    )
-    environmental_costs = (production_costs.sel(commodity=environmentals)).sum(
-        "commodity"
-    )
-    # revenues due to product sales
-    revenues = (production_costs.sel(commodity=enduses)).sum("commodity")
-
-    # Gross margin is the net between revenues and all costs
-    result = revenues - environmental_costs - variable_costs - consumption_costs
-
-    # Gross margin is defined as a ratio on revenues and as a percentage
-    result *= 100 / revenues
-    return result
-
-
 def consumption(
     technologies: xr.Dataset,
     production: xr.DataArray,
@@ -352,8 +267,8 @@ def maximum_production(
     capa = filter_input(
         capacity, **{k: v for k, v in filters.items() if k in capacity.dims}
     )
-    btechs = broadcast_techs(  # type: ignore
-        cast(xr.Dataset, technologies[["fixed_outputs", "utilization_factor"]]), capa
+    btechs = broadcast_techs(
+        technologies[["fixed_outputs", "utilization_factor"]], capa
     )
     ftechs = filter_input(
         btechs, **{k: v for k, v in filters.items() if k in btechs.dims}
@@ -470,12 +385,8 @@ def minimum_production(
     if "minimum_service_factor" not in technologies:
         return broadcast_timeslice(xr.zeros_like(capa), level=timeslice_level)
 
-    btechs = broadcast_techs(  # type: ignore
-        cast(
-            xr.Dataset,
-            technologies[["fixed_outputs", "minimum_service_factor"]],
-        ),
-        capa,
+    btechs = broadcast_techs(
+        technologies[["fixed_outputs", "minimum_service_factor"]], capa
     )
     ftechs = filter_input(
         btechs, **{k: v for k, v in filters.items() if k in btechs.dims}

src/muse/sectors/sector.py

@@ -314,7 +314,7 @@ def market_variables(self, market: xr.Dataset, technologies: xr.Dataset) -> Any:
 
         # Calculate LCOE
         # We select data for the second year, which corresponds to the investment year
-        technodata = cast(xr.Dataset, broadcast_techs(technologies, supply))
+        technodata = broadcast_techs(technologies, supply)
         lcoe = levelized_cost_of_energy(
             prices=market.prices.sel(region=supply.region).isel(year=1),
             technologies=technodata,

src/muse/utilities.py

@@ -182,7 +182,6 @@ def operation(x):
 def broadcast_techs(
     technologies: xr.Dataset | xr.DataArray,
     template: xr.DataArray | xr.Dataset,
-    dimension: str = "asset",
     interpolation: str = "linear",
     installed_as_year: bool = True,
     **kwargs,
@@ -204,18 +203,57 @@ def broadcast_techs(
     Arguments:
         technologies: The dataset to broadcast
         template: the dataset or data-array to use as a template
-        dimension: the name of the dimensiom from `template` over which to
-            broadcast
         interpolation: interpolation method used across `year`
         installed_as_year: if the coordinate `installed` exists, then it is
             applied to the `year` dimension of the technologies dataset
         kwargs: further arguments are used initial filters over the
             `technologies` dataset.
+
+    Example:
+        Define the technology array:
+        >>> import xarray as xr
+        >>> technologies = xr.DataArray(
+        ...     data=[[1, 2, 3], [4, 5, 6]],
+        ...     dims=['technology', 'region'],
+        ...     coords={'technology': ['gasboiler', 'heatpump'],
+        ...             'region': ['R1', 'R2', 'R3']},
+        ... )
+
+        This array contains a value for every combination of technology and region (e.g.
+        this could refer to the efficiency of each technology in each region).
+
+        Define the assets template:
+        >>> assets = xr.DataArray(
+        ...     data=[10, 50],
+        ...     dims=["asset"],
+        ...     coords={
+        ...         "region": (["asset"], ["R1", "R2"]),
+        ...         "technology": (["asset"], ["gasboiler", "heatpump"])},
+        ... )
+
+        We have two assets: a gas boiler in region R1 and a heat pump in region R2. In
+        this case the values don't matter, but could correspond to the installed
+        capacity of each asset, for example.
+
+        We want to select the values from the technology array that correspond to each
+        asset in the template. To do this, we perform `broadcast_techs` on
+        `technologies` using `assets` as a template:
+        >>> broadcast_techs(technologies, assets)
+        <xarray.DataArray (asset: 2)> Size: 16B
+        array([1, 5])
+        Coordinates:
+            technology  (asset) <U9 72B 'gasboiler' 'heatpump'
+            region      (asset) <U2 16B 'R1' 'R2'
+        Dimensions without coordinates: asset
+
+        The output array has the same shape as the assets template. Each value in the
+        output is the value in the original technology array that matches the
+        technology & region of each asset.
     """
     # this assert will trigger if 'year' is changed to 'installed' in
     # technologies, because then this function should be modified.
     assert "installed" not in technologies.dims
-    names = [u for u in template.coords if template[u].dims == (dimension,)]
+    names = [u for u in template.coords if template[u].dims == ("asset",)]
     # the first selection reduces the size of technologies without affecting the
     # dimensions.
     first_sel = {
@@ -293,7 +331,6 @@ def filter_input(
 def filter_with_template(
     data: xr.Dataset | xr.DataArray,
     template: xr.DataArray | xr.Dataset,
-    asset_dimension: str = "asset",
     **kwargs,
 ):
     """Filters data to match template.
@@ -313,8 +350,8 @@ def filter_with_template(
     Returns:
         `data` transformed to match the form of `template`
     """
-    if asset_dimension in template.dims:
-        return broadcast_techs(data, template, dimension=asset_dimension, **kwargs)
+    if "asset" in template.dims:
+        return broadcast_techs(data, template)
 
     match_indices = set(data.dims).intersection(template.dims) - set(kwargs)
     match = {d: template[d].isin(data[d]).values for d in match_indices if d != "year"}