Ban implicit broadcasting over region dimension

src/muse/__main__.py

@@ -89,6 +89,15 @@ def patched_broadcast_compat_data(self, other):
             "`broadcast_timeslice` or `distribute_timeslice` (see `muse.timeslices`)."
         )
 
+    if (isinstance(other, Variable)) and ("region" in self.dims) != (
+        "region" in getattr(other, "dims", [])
+    ):
+        raise ValueError(
+            "Broadcasting along the 'region' dimension is required, but automatic "
+            "broadcasting is disabled. Please handle it explicitly using "
+            "`broadcast_regions` (see `muse.utilities`)."
+        )
+
     # 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/demand_share.py

@@ -503,6 +503,10 @@ def unmet_demand(
     assert "year" not in capacity.dims
     assert "year" not in demand.dims
 
+    # If there are no assets, no production is possible, so all demand is unmet.
+    if capacity.sizes["asset"] == 0:
+        return demand.copy()
+
     # Calculate maximum production by existing assets
     produced = maximum_production(
         capacity=capacity,

src/muse/mca.py

@@ -15,7 +15,7 @@
 from muse.readers import read_initial_market
 from muse.sectors import SECTORS_REGISTERED, AbstractSector, Sector
 from muse.timeslices import broadcast_timeslice, drop_timeslice
-from muse.utilities import future_propagation
+from muse.utilities import broadcast_regions, future_propagation
 
 
 class MCA:
@@ -298,12 +298,14 @@ def run(self) -> None:
                 getLogger(__name__).info(
                     f"Updating carbon price for year {investment_year}"
                 )
-                new_price = self.update_carbon_price(new_market)
+                new_price: float = self.update_carbon_price(new_market)
                 future_price = DataArray(new_price, coords=dict(year=investment_year))
                 new_market.prices.loc[dict(commodity=self.carbon_commodities)] = (
                     future_propagation(
                         new_market.prices.sel(commodity=self.carbon_commodities),
-                        broadcast_timeslice(future_price),
+                        broadcast_timeslice(
+                            broadcast_regions(future_price, new_market.region)
+                        ),
                     )
                 )
 

src/muse/timeslices.py

@@ -21,6 +21,8 @@
 import pandas as pd
 from xarray import DataArray
 
+from muse.utilities import broadcast_regions
+
 TIMESLICE: DataArray = None  # type: ignore
 
 
@@ -156,6 +158,8 @@ def distribute_timeslice(
     broadcasted = broadcast_timeslice(data, ts=ts)
     timeslice_sum = ts.sum("timeslice").clip(1e-6)  # prevents zero division
     timeslice_fractions = ts / broadcast_timeslice(timeslice_sum, ts=ts)
+    if "region" in data.dims:
+        timeslice_fractions = broadcast_regions(timeslice_fractions, data)
     return broadcasted * timeslice_fractions
 
 

src/muse/utilities.py

@@ -687,3 +687,23 @@ def camel_to_snake(name: str) -> str:
     result = result.replace("n2_o", "N2O")
     result = result.replace("f-gases", "F-gases")
     return result
+
+
+def broadcast_regions(data: xr.DataArray, template: xr.DataArray) -> xr.DataArray:
+    """Convert a non-regional array to a regional array by broadcasting.
+
+    If data is already regioned in the appropriate scheme, it will be returned
+    unchanged.
+
+    Args:
+        data: Array to broadcast.
+        template: Dataarray with region coordinates to broadcast to.
+
+    """
+    # If data already has regions, check that it matches the template regions.
+    if "region" in data.dims:
+        if data.region.equals(template.region):
+            return data
+        raise ValueError("Data is already regioned, but does not match the reference.")
+
+    return data.expand_dims(region=template.region)

tests/conftest.py

@@ -13,6 +13,7 @@
 
 from muse.__main__ import patched_broadcast_compat_data
 from muse.agents import Agent
+from muse.utilities import broadcast_regions
 
 
 @contextmanager
@@ -241,7 +242,9 @@ def var(*dims, factor=100.0):
         return dims, (rand(*shape) * factor).astype(type(factor))
 
     result["agent_share"] = var("technology", "region", "year")
-    result["agent_share"] /= sum(result.agent_share)
+    result["agent_share"] /= broadcast_regions(
+        sum(result.agent_share), result.agent_share
+    )
     result["agent_share_zero"] = result["agent_share"] * 0
 
     # first create a mask so each tech will have consistent inputs/outputs across years
@@ -273,12 +276,20 @@ def var(*dims, factor=100.0):
             fout.loc[{"technology": tech, "commodity": i}] = 1
 
     # expand along year and region, and fill with random numbers
-    ones = (result.year == result.year) * (result.region == result.region)
-    result["fixed_inputs"] = result.fixed_inputs * ones
+    ones = broadcast_regions(result.year == result.year, result.region) * (
+        result.region == result.region
+    )
+    result["fixed_inputs"] = (
+        broadcast_regions(result.fixed_inputs, result.region) * ones
+    )
     result.fixed_inputs[:] *= rand(*result.fixed_inputs.shape)
-    result["flexible_inputs"] = result.flexible_inputs * ones
+    result["flexible_inputs"] = (
+        broadcast_regions(result.flexible_inputs, result.region) * ones
+    )
     result.flexible_inputs[:] *= rand(*result.flexible_inputs.shape)
-    result["fixed_outputs"] = result.fixed_outputs * ones
+    result["fixed_outputs"] = (
+        broadcast_regions(result.fixed_outputs, result.region) * ones
+    )
     result.fixed_outputs[:] *= rand(*result.fixed_outputs.shape)
 
     result["total_capacity_limit"] = var("technology", "region", "year")