Merge branch 'techs' into foresight

Author: tsmbland

src/muse/__main__.py

@@ -89,6 +89,11 @@ def patched_broadcast_compat_data(self, other):
             "`broadcast_timeslice` or `distribute_timeslice` (see `muse.timeslices`)."
         )
 
+    if (isinstance(other, Variable)) and {"technology", "asset"}.issubset(
+        set(self.dims) | set(getattr(other, "dims", []))
+    ):
+        raise ValueError()
+
     # The rest of the function is copied directly from
     # xarray.core.variable._broadcast_compat_data
     if all(hasattr(other, attr) for attr in ["dims", "data", "shape", "encoding"]):

src/muse/utilities.py

@@ -180,12 +180,9 @@ def operation(x):
 
 
 def broadcast_techs(
-    technologies: xr.Dataset | xr.DataArray,
+    x: xr.Dataset | xr.DataArray,
     template: xr.DataArray | xr.Dataset,
     dimension: str = "asset",
-    interpolation: str = "linear",
-    installed_as_year: bool = True,
-    **kwargs,
 ) -> xr.Dataset | xr.DataArray:
     """Broadcasts technologies to the shape of template in given dimension.
 
@@ -201,39 +198,49 @@ def broadcast_techs(
     This function broadcast the first representation to the shape and coordinates
     of the second.
 
-    This function does not support technologies with a 'year' dimension. Please select
-    technology data for a specific year before calling this function.
+    This function does not support arrays with a 'year' dimension. Please select
+    data for a specific year before calling this function.
 
     Arguments:
-        technologies: The dataset to broadcast
+        x: 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.
-    """
-    assert "year" not in technologies.dims
+        dimension: TODO
+
+    Example:
+        Define the example array:
+        >>> import xarray as xr
+        >>> x = xr.DataArray(
+        ...     data=[[1, 2, 3], [4, 5, 6]],
+        ...     dims=['technology', 'region'],
+        ...     coords={'technology': ['gasboiler', 'heatpump'],
+        ...             'region': ['R1', 'R2', 'R3']},
+        ... )
 
-    # 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,)]
-    # the first selection reduces the size of technologies without affecting the
-    # dimensions.
-    first_sel = {
-        n: technologies[n].isin(template[n])
-        for n in names
-        if n in technologies.dims and n != "year"
-    }
-    first_sel.update({k: v for k, v in kwargs.items() if k != "year"})
-    techs = technologies.sel(first_sel)
+        Define the assets template:
+        >>> template = xr.DataArray(
+        ...     data=[0, 0],
+        ...     dims=["asset"],
+        ...     coords={
+        ...         "region": (["asset"], ["R1", "R2"]),
+        ...         "technology": (["asset"], ["gasboiler", "heatpump"]),
+        ...         "installed": (["asset"], [2020, 2025])},
+        ... )
 
-    second_sel = {n: template[n] for n in template.coords if n in techs.dims}
+        Reshape the data to match the template:
+        >>> broadcast_techs(x, template)
+        <xarray.DataArray (asset: 2)> Size: 16B
+        array([1, 5])
+        Coordinates:
+            technology  (asset) <U9 72B 'gasboiler' 'heatpump'
+            region      (asset) <U2 16B 'R1' 'R2'
+            installed   (asset) int64 16B 2020 2025
+        Dimensions without coordinates: asset
+    """
+    assert "year" not in x.dims
+    assert "installed" not in x.dims
 
-    return techs.sel(second_sel)
+    sel = {n: template[n] for n in template.coords if n in x.dims}
+    return x.sel(sel)
 
 
 def clean_assets(assets: xr.Dataset, years: int | Sequence[int]):
@@ -313,7 +320,7 @@ def filter_with_template(
         `data` transformed to match the form of `template`
     """
     if asset_dimension in template.dims:
-        return broadcast_techs(data, template, dimension=asset_dimension, **kwargs)
+        return broadcast_techs(data, template, dimension=asset_dimension)
 
     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"}

tests/test_utilities.py

@@ -43,11 +43,11 @@ def test_reduce_assets_with_zero_size(capacity: xr.DataArray):
 def test_broadcast_tech(technologies, capacity):
     from muse.utilities import broadcast_techs
 
+    technologies = technologies.sel(year=2010)
     regions = make_array(technologies.region)
     commodities = make_array(technologies.commodity)
-    years = make_array(technologies.year)
     techs = make_array(technologies.technology)
-    technologies["fixed_outputs"] = regions * commodities * years * techs
+    technologies["fixed_outputs"] = regions * commodities * techs
 
     actual = broadcast_techs(technologies.fixed_outputs, capacity)
 
@@ -57,16 +57,15 @@ def test_broadcast_tech(technologies, capacity):
 
     for asset in capacity.asset:
         region = regions.sel(region=asset.region)
-        year = years.interp(year=asset.installed, method="linear")
         tech = techs.sel(technology=asset.technology)
-        expected = region * year * tech * commodities
+        expected = region * tech * commodities
         assert actual.isel(asset=int(asset)).values == approx(expected.values)
 
 
 def test_broadcast_tech_idempotent(technologies, capacity):
     from muse.utilities import broadcast_techs
 
-    first = broadcast_techs(technologies, capacity)
+    first = broadcast_techs(technologies.sel(year=2010), capacity)
     second = broadcast_techs(first, capacity)
     assert (first == second).all()