diff --git a/src/muse/__main__.py b/src/muse/__main__.py index f5de4d1fa..80f6abee1 100644 --- a/src/muse/__main__.py +++ b/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"]): diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py index 746ca3e76..81ae37405 100644 --- a/src/muse/demand_share.py +++ b/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, diff --git a/src/muse/mca.py b/src/muse/mca.py index badaeb3ca..183bedafc 100644 --- a/src/muse/mca.py +++ b/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) + ), ) ) diff --git a/src/muse/timeslices.py b/src/muse/timeslices.py index c7603f968..4017ce2b8 100644 --- a/src/muse/timeslices.py +++ b/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 diff --git a/src/muse/utilities.py b/src/muse/utilities.py index 66843382e..0c1728c05 100644 --- a/src/muse/utilities.py +++ b/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) diff --git a/tests/conftest.py b/tests/conftest.py index 3d1f3b434..63061a2c8 100644 --- a/tests/conftest.py +++ b/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")