Ban implicit broadcasting over year dimension

src/muse/__main__.py

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

@@ -8,6 +8,7 @@
 import xarray as xr
 
 from muse.timeslices import drop_timeslice
+from muse.utilities import broadcast_years
 
 
 class AbstractAgent(ABC):
@@ -494,7 +495,9 @@ def retirement_profile(
             investments = investments.reindex_like(profile, method="ffill")
 
         # Apply the retirement profile to the investments
-        new_assets = (investments * profile).rename(replacement="asset")
+        new_assets = (broadcast_years(investments, profile) * profile).rename(
+            replacement="asset"
+        )
         new_assets["installed"] = "asset", [investment_year] * len(new_assets.asset)
 
         # The new assets have picked up quite a few coordinates along the way.

src/muse/agents/factories.py

@@ -10,6 +10,7 @@
 
 from muse.agents.agent import Agent, InvestingAgent
 from muse.errors import AgentShareNotDefined, TechnologyNotDefined
+from muse.utilities import broadcast_years
 
 
 def create_standard_agent(
@@ -252,8 +253,8 @@ def _shared_capacity(
     techs = (existing > 0) & (shares > 0)
     techs = techs.any([u for u in techs.dims if u != "asset"])
     if not any(techs):
-        return (capacity * shares).copy()
-    return (capacity * shares).sel(asset=techs.values).copy()
+        return (capacity * broadcast_years(shares, capacity)).copy()
+    return (capacity * broadcast_years(shares, capacity)).sel(asset=techs.values).copy()
 
 
 def _standardize_inputs(

src/muse/costs.py

@@ -55,6 +55,7 @@
 from muse.commodities import is_enduse, is_fuel, is_material, is_pollutant
 from muse.quantities import production_amplitude
 from muse.timeslices import broadcast_timeslice, distribute_timeslice, get_level
+from muse.utilities import broadcast_years
 
 
 def cost(func):
@@ -625,10 +626,11 @@ def annual_to_lifetime(costs: xr.DataArray, technologies: xr.Dataset):
     rates = discount_factor(
         years=years,
         interest_rate=technologies.interest_rate,
-        mask=years <= life,
+        mask=years <= broadcast_years(life, years),
     )
     if "timeslice" in costs.dims:
         rates = broadcast_timeslice(rates, level=get_level(costs))
+    costs = broadcast_years(costs, years)
     return (costs * rates).sum("year")
 
 
@@ -648,7 +650,7 @@ def discount_factor(
     assert "year" not in interest_rate.dims
 
     # Calculate discount factor over the years
-    df = 1 / (1 + interest_rate) ** years
+    df = 1 / (1 + broadcast_years(interest_rate, years)) ** years
 
     # Apply mask
     if mask is not None:

src/muse/dispatch.py

@@ -49,6 +49,7 @@ def production(
 import xarray as xr
 
 from muse.registration import registrator
+from muse.utilities import check_dimensions
 
 
 class PRODUCTION_SIGNATURE(Protocol):
@@ -141,8 +142,11 @@ def share_based_production(
     from muse.quantities import emission, maximum_production, minimum_production
     from muse.utilities import broadcast_over_assets
 
-    assert "asset" not in demand.dims
-    assert "asset" in capacity.dims
+    check_dimensions(demand, ["timeslice", "commodity"], optional=["region"])
+    check_dimensions(capacity, ["asset"], optional=["dst_region"])
+    check_dimensions(
+        technologies, ["asset", "commodity"], optional=["timeslice", "dst_region"]
+    )
 
     # Maximum and minimum production for each asset
     maxprod = maximum_production(
@@ -234,9 +238,15 @@ def merit_order_production(
     )
     from muse.utilities import broadcast_over_assets
 
-    assert "asset" not in demand.dims
-    assert "asset" in capacity.dims
-    assert "asset" not in prices.dims
+    if "dst_region" in technologies.dims:
+        raise ValueError(
+            "Merit-order dispatch is not currently compatible with trade models."
+        )
+
+    check_dimensions(demand, ["timeslice", "commodity"], optional=["region"])
+    check_dimensions(capacity, ["asset"])
+    check_dimensions(technologies, ["asset", "commodity"], optional=["timeslice"])
+    check_dimensions(prices, ["timeslice", "commodity"], optional=["region"])
 
     # Normalise demand/prices dataarrays to ensure they have a region dimension
     # Multi-region models will already have a region dimension
@@ -283,51 +293,36 @@ def merit_order_production(
     # Initialise result with zeros
     result = xr.zeros_like(maxprod)
 
-    for y in maxprod.year.values:
-        prices_y = prices.sel(year=y)
-        maxprod_y = maxprod.sel(year=y)
-        minprod_y = minprod.sel(year=y)
-        maxcons_y = maxcons.sel(year=y)
-
-        for region in demand.region.values:
-            maxprod_region = maxprod_y.sel(
-                asset=maxprod_y.asset[maxprod_y.region == region]
-            )
-            techs_region = technologies.sel(
-                asset=technologies.asset[technologies.region == region]
-            )
-            minprod_region = minprod_y.sel(
-                asset=minprod_y.asset[minprod_y.region == region]
-            )
-            maxcons_region = maxcons_y.sel(
-                asset=maxcons_y.asset[maxcons_y.region == region]
-            )
+    for region in demand.region.values:
+        # Select data for this region
+        maxprod_region = maxprod.sel(asset=maxprod.asset[maxprod.region == region])
+        techs_region = technologies.sel(
+            asset=technologies.asset[technologies.region == region]
+        )
+        minprod_region = minprod.sel(asset=minprod.asset[minprod.region == region])
+        maxcons_region = maxcons.sel(asset=maxcons.asset[maxcons.region == region])
+
+        # Calculate timeslice-level costs for each asset assuming full
+        # dispatch. We use LCOE excluding capital costs.
+        technology_costs = marginal_cost(
+            techs_region,
+            prices.sel(region=region),
+            production=maxprod_region,
+            consumption=maxcons_region,
+        )
 
-            # Calculate timeslice-level costs for each asset in this year assuming full
-            # dispatch. We use LCOE excluding capital costs.
-            technology_costs = marginal_cost(
-                techs_region,
-                prices_y.sel(region=region),
-                production=maxprod_region,
-                consumption=maxcons_region,
+        # Calculate production by dispatching assets in order of
+        # increasing cost until demand is met
+        for ts in maxprod.timeslice.values:
+            dispatch = dispatch_by_merit_order(
+                demand=demand.sel(timeslice=ts, region=region),
+                minprod=minprod_region.sel(timeslice=ts),
+                maxprod=maxprod_region.sel(timeslice=ts),
+                technology_costs=technology_costs.sel(timeslice=ts),
             )
-
-            # Calculate production for this year by dispatching assets in order of
-            # increasing cost until demand is met
-            for ts in maxprod_y.timeslice.values:
-                dispatch = dispatch_by_merit_order(
-                    demand=demand.sel(year=y, timeslice=ts, region=region),
-                    minprod=minprod_region.sel(timeslice=ts),
-                    maxprod=maxprod_region.sel(timeslice=ts),
-                    technology_costs=technology_costs.sel(timeslice=ts),
-                )
-                result.loc[
-                    dict(
-                        year=y,
-                        timeslice=ts,
-                        asset=result.asset[result.region == region],
-                    )
-                ] = dispatch
+            result.loc[
+                dict(timeslice=ts, asset=result.asset[result.region == region])
+            ] = dispatch
 
     # Add production of environmental pollutants
     env = is_pollutant(technologies.comm_usage)

src/muse/examples.py

@@ -39,6 +39,7 @@
 from muse.mca import MCA
 from muse.sectors import AbstractSector
 from muse.timeslices import drop_timeslice
+from muse.utilities import broadcast_years
 
 __all__ = ["model", "technodata"]
 
@@ -271,6 +272,7 @@ def matching_market(sector: str, model: str = "default") -> xr.Dataset:
 
     market = xr.Dataset()
     techs = broadcast_over_assets(loaded_sector.technologies, assets.capacity)
+    techs = broadcast_years(techs, assets.capacity)
     production = maximum_production(techs, assets.capacity)
     market["supply"] = production.sum("asset")
     if "dst_region" in market.dims:

src/muse/investments.py

@@ -65,6 +65,7 @@ def investment(
 from muse.outputs.cache import cache_quantity
 from muse.registration import registrator
 from muse.timeslices import timeslice_max
+from muse.utilities import broadcast_years
 
 INVESTMENT_SIGNATURE = Callable[
     [xr.DataArray, xr.DataArray, xr.Dataset, list[Constraint], KwArg(Any)],
@@ -203,7 +204,7 @@ def cliff_retirement_profile(
         dims="year",
         coords={"year": range(investment_year, max_year + 1)},
     )
-    profile = allyears < (investment_year + technical_life)  # type: ignore
+    profile = allyears < broadcast_years(investment_year + technical_life, allyears)  # type: ignore
 
     # Minimize the number of years needed to represent the profile fully
     # This is done by removing the central year of any three repeating years, ensuring

src/muse/readers/toml.py

@@ -16,6 +16,7 @@
 import xarray as xr
 
 from muse.defaults import DATA_DIRECTORY
+from muse.utilities import broadcast_years
 
 DEFAULT_SETTINGS_PATH = DATA_DIRECTORY / "default_settings.toml"
 """Default settings path."""
@@ -721,7 +722,9 @@ def read_correlation_consumption(sector_conf: Any) -> xr.Dataset:
     # Split by timeslice
     if sector_conf.timeslice_shares_path is not None:
         shares = read_timeslice_shares(sector_conf.timeslice_shares_path)
-        consumption = broadcast_timeslice(consumption) * shares
+        consumption = broadcast_timeslice(consumption) * broadcast_years(
+            shares, consumption.year
+        )
     else:
         consumption = distribute_timeslice(consumption)
 

src/muse/regressions.py

@@ -9,6 +9,8 @@
 
 from xarray import DataArray, Dataset
 
+from muse.utilities import broadcast_years
+
 __all__ = [
     "Exponential",
     "ExponentialAdj",
@@ -356,8 +358,9 @@ def Exponential(
 ) -> DataArray:
     from numpy import exp
 
-    factor = 1e6 * self.coeffs.a * population
-    return factor * exp(self.coeffs.b * population / gdp)
+    coeffs = broadcast_years(self.coeffs, gdp.year)
+    factor = 1e6 * coeffs.a * population
+    return factor * exp(coeffs.b * population / gdp)
 
 
 @register_regression
@@ -377,10 +380,11 @@ def ExponentialAdj(
     if year is None:
         year = self.base_year
 
-    factor = 1e6 * self.coeffs.a * population
-    unadjusted = factor * exp(self.coeffs.b * population / gdp)
+    coeffs = broadcast_years(self.coeffs, gdp.year)
+    factor = 1e6 * coeffs.a * population
+    unadjusted = factor * exp(coeffs.b * population / gdp)
     p = power(year + forecast - self.base_year, n)
-    return unadjusted * (1 + self.coeffs.w * p) / (1 + p)
+    return unadjusted * (1 + coeffs.w * p) / (1 + p)
 
 
 @register_regression
@@ -394,7 +398,8 @@ def Logistic(
     """
     from numpy import exp, power
 
-    a, b, c, w = self.coeffs.a, self.coeffs.b, self.coeffs.c, self.coeffs.w
+    coeffs = broadcast_years(self.coeffs, gdp.year)
+    a, b, c, w = coeffs.a, coeffs.b, coeffs.c, coeffs.w
     p = power(forecast, n)
     factor = 1e6 * a * population * (1 + w * p) / (1 + p)
     return factor / (1 + b * exp(gdp * c / population))
@@ -406,8 +411,9 @@ def Loglog(self, gdp: DataArray, population: DataArray, *args, **kwargs) -> Data
     """1e6 * e^a * population * (gpd/population)^b."""
     from numpy import exp, power
 
-    factor = 1e6 * exp(self.coeffs.a) * population
-    return factor * power(gdp / population, self.coeffs.b)
+    coeffs = broadcast_years(self.coeffs, gdp.year)
+    factor = 1e6 * exp(coeffs.a) * population
+    return factor * power(gdp / population, coeffs.b)
 
 
 @register_regression
@@ -424,6 +430,7 @@ def LogisticSigmoid(
 ) -> DataArray:
     """0.001 * (constant * pop + gdp * c / sqrt(1 + (gdp * scale / pop)^2)."""
     from numpy import power
+    from xarray import ones_like
 
     constant = self.coeffs.a
     c = self.coeffs.c
@@ -442,8 +449,11 @@ def LogisticSigmoid(
         # fmt: enable
         scale = self.coeffs.b0.where(years < 2015, self.coeffs.b1)
     else:
-        scale = 1
+        scale = ones_like(self.coeffs.b0)
 
+    scale = broadcast_years(scale, gdp.year)
+    constant = broadcast_years(constant, gdp.year)
+    c = broadcast_years(c, gdp.year)
     p = power(1 + power(gdp * scale / population, 2), 0.5)
     return 0.001 * (constant * population + gdp * c / p)
 
@@ -498,7 +508,10 @@ def __call__(
 
         if year is not None and "year" in data.dims:
             data = data.interp(year=year, method=self.interpolation)
-        return coeffs.a * data.population + scale * (
+        a = broadcast_years(coeffs.a, data.year)
+        scale = broadcast_years(scale, data.year)
+        gdpcap_offset = broadcast_years(gdpcap_offset, data.year)
+        return a * data.population + scale * (
             data.gdp - gdpcap_offset * data.population
         )