Helper function for dropping timeslice

src/muse/agents/agent.py

@@ -6,6 +6,8 @@
 
 import xarray as xr
 
+from muse.timeslices import drop_timeslice
+
 
 class AbstractAgent(ABC):
     """Base class for all agents."""
@@ -416,7 +418,7 @@ def next(
             return None
 
         if "timeslice" in search.dims:
-            search["demand"] = demand.drop_vars(["timeslice", "month", "day", "hour"])
+            search["demand"] = drop_timeslice(demand)
         else:
             search["demand"] = demand
         not_assets = [u for u in search.demand.dims if u != "asset"]

src/muse/constraints.py

@@ -115,6 +115,7 @@ def constraints(
 from mypy_extensions import KwArg
 
 from muse.registration import registrator
+from muse.timeslices import drop_timeslice
 
 CAPACITY_DIMS = "asset", "replacement", "region"
 """Default dimensions for capacity decision variables."""
@@ -691,7 +692,7 @@ def lp_costs(
     production = zeros_like(ts_costs * fouts)
     for dim in production.dims:
         if isinstance(production.get_index(dim), pd.MultiIndex):
-            production = production.drop_vars(["timeslice", "month", "day", "hour"])
+            production = drop_timeslice(production)
             production[dim] = pd.Index(production.get_index(dim), tupleize_cols=False)
 
     return xr.Dataset(dict(capacity=costs, production=production))
@@ -758,7 +759,7 @@ def lp_constraint(constraint: Constraint, lpcosts: xr.Dataset) -> Constraint:
     constraint = constraint.copy(deep=False)
     for dim in constraint.dims:
         if isinstance(constraint.get_index(dim), pd.MultiIndex):
-            constraint = constraint.drop_vars(["timeslice", "month", "day", "hour"])
+            constraint = drop_timeslice(constraint)
             constraint[dim] = pd.Index(constraint.get_index(dim), tupleize_cols=False)
     b = constraint.b.drop_vars(set(constraint.b.coords) - set(constraint.b.dims))
     b = b.rename({k: f"c({k})" for k in b.dims})

src/muse/demand_matching.py

@@ -51,6 +51,8 @@
 import pandas as pd
 from xarray import DataArray
 
+from muse.timeslices import drop_timeslice
+
 
 def demand_matching(
     demand: DataArray,
@@ -248,7 +250,7 @@ def demand_matching(
 
     if len(multics) > 0:
         for k in multics:
-            ds = ds.drop_vars(["timeslice", "month", "day", "hour"])
+            ds = drop_timeslice(ds)
             ds[k] = pd.Index(constraint.get_index(k), tupleize_cols=False)
         result = demand_matching(  # type: ignore
             ds.demand, ds.cost, *(ds[f"constraint{i}"] for i in range(len(constraints)))

src/muse/examples.py

@@ -34,6 +34,7 @@
 
 from muse.mca import MCA
 from muse.sectors import AbstractSector
+from muse.timeslices import drop_timeslice
 
 __all__ = ["model", "technodata"]
 
@@ -194,12 +195,8 @@ def mca_market(model: str = "default") -> xr.Dataset:
             .sel(region=settings.regions)
             .interp(year=settings.time_framework, method=settings.interpolation_mode)
         )
-        market["supply"] = zeros_like(market.exports).drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
-        market["consumption"] = zeros_like(market.exports).drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
+        market["supply"] = drop_timeslice(zeros_like(market.exports))
+        market["consumption"] = drop_timeslice(zeros_like(market.exports))
 
         return cast(xr.Dataset, market)
 
@@ -265,12 +262,12 @@ def matching_market(sector: str, model: str = "default") -> xr.Dataset:
         market = market.rename(dst_region="region")
     if market.region.dims:
         consump = consumption(loaded_sector.technologies, production)
-        market["consumption"] = (
+        market["consumption"] = drop_timeslice(
             consump.groupby("region").sum(
                 {"asset", "dst_region"}.intersection(consump.dims)
             )
             + market.supply
-        ).drop_vars(["timeslice", "month", "day", "hour"])
+        )
     else:
         market["consumption"] = (
             consumption(loaded_sector.technologies, production).sum(

src/muse/mca.py

@@ -14,6 +14,7 @@
 from muse.outputs.cache import OutputCache
 from muse.readers import read_initial_market
 from muse.sectors import SECTORS_REGISTERED, AbstractSector
+from muse.timeslices import drop_timeslice
 
 
 class MCA:
@@ -59,12 +60,8 @@ def factory(cls, settings: str | Path | Mapping | Any) -> MCA:
             ).sel(region=settings.regions)
         ).interp(year=settings.time_framework, method=settings.interpolation_mode)
 
-        market["supply"] = zeros_like(market.exports).drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
-        market["consumption"] = zeros_like(market.exports).drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
+        market["supply"] = drop_timeslice(zeros_like(market.exports))
+        market["consumption"] = drop_timeslice(zeros_like(market.exports))
 
         # We create the sectors
         sectors = []
@@ -461,9 +458,7 @@ def single_year_iteration(
     sectors = deepcopy(sectors)
     market = market.copy(deep=True)
     if "updated_prices" not in market.data_vars:
-        market["updated_prices"] = market.prices.copy().drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
+        market["updated_prices"] = drop_timeslice(market.prices.copy())
 
     # eventually, the first market should be one that creates the initial demand
     for sector in sectors:
@@ -551,16 +546,12 @@ def find_equilibrium(
     else:
         included = ones(len(market.commodity), dtype=bool)
 
-    market["updated_prices"] = market.prices.copy().drop_vars(
-        ["timeslice", "month", "day", "hour"]
-    )
+    market["updated_prices"] = drop_timeslice(market.prices.copy())
     prior_market = market.copy(deep=True)
     converged = False
     equilibrium_sectors = sectors
     for i in range(maxiter):
-        market["prices"] = market.updated_prices.drop_vars(
-            ["timeslice", "month", "day", "hour"]
-        )
+        market["prices"] = drop_timeslice(market.updated_prices)
         prior_market, market = market, prior_market
         market.consumption[:] = 0.0
         market.supply[:] = 0.0
@@ -582,9 +573,11 @@ def find_equilibrium(
             new_price.loc[dict(commodity=included)] = market.updated_prices.sel(
                 commodity=included, year=market.year[1]
             )
-            market["prices"] = future_propagation(  # type: ignore
-                market["prices"], new_price
-            ).drop_vars(["timeslice", "month", "day", "hour"])
+            market["prices"] = drop_timeslice(
+                future_propagation(  # type: ignore
+                    market["prices"], new_price
+                )
+            )
 
             break
 
@@ -599,9 +592,11 @@ def find_equilibrium(
                     dict(year=market.year[1], commodity=included)
                 ]
             )
-            market["prices"] = future_propagation(  # type: ignore
-                market["prices"], new_price
-            ).drop_vars(["timeslice", "month", "day", "hour"])
+            market["prices"] = drop_timeslice(
+                future_propagation(  # type: ignore
+                    market["prices"], new_price
+                )
+            )
         if not equilibrium:
             equilibrium_reached = True
             converged = True
@@ -626,9 +621,11 @@ def find_equilibrium(
             0.2
             * market.updated_prices.loc[dict(year=market.year[1], commodity=included)]
         )
-        market["prices"] = future_propagation(  # type: ignore
-            market["prices"], new_price
-        ).drop_vars(["timeslice", "month", "day", "hour"])
+        market["prices"] = drop_timeslice(
+            future_propagation(  # type: ignore
+                market["prices"], new_price
+            )
+        )
         getLogger(__name__).critical(msg)
 
     return FindEquilibriumResults(

src/muse/objectives.py

@@ -81,6 +81,7 @@ def comfort(
 from muse.agents import Agent
 from muse.outputs.cache import cache_quantity
 from muse.registration import registrator
+from muse.timeslices import drop_timeslice
 
 OBJECTIVE_SIGNATURE = Callable[
     [Agent, xr.DataArray, xr.DataArray, xr.Dataset, xr.Dataset, KwArg(Any)],
@@ -142,7 +143,7 @@ def objectives(
         for name, objective in functions:
             obj = objective(agent, demand, search_space, *args, **kwargs)
             if "timeslice" in obj.dims and "timeslice" in result.dims:
-                obj = obj.drop_vars(["timeslice", "month", "day", "hour"])
+                obj = drop_timeslice(obj)
             result[name] = obj
         return result
 

src/muse/outputs/mca.py

@@ -34,7 +34,7 @@ def quantity(
 
 from muse.registration import registrator
 from muse.sectors import AbstractSector
-from muse.timeslices import QuantityType, convert_timeslice
+from muse.timeslices import QuantityType, convert_timeslice, drop_timeslice
 
 OUTPUT_QUANTITY_SIGNATURE = Callable[
     [xr.Dataset, list[AbstractSector], KwArg(Any)], Union[xr.DataArray, pd.DataFrame]
@@ -362,9 +362,9 @@ def sector_supply(sector: AbstractSector, market: xr.Dataset, **kwargs) -> pd.Da
             capacity = a.filter_input(a.assets.capacity, year=output_year).fillna(0.0)
             technologies = a.filter_input(techs, year=output_year).fillna(0.0)
             agent_market = market.sel(year=output_year).copy()
-            agent_market["consumption"] = (
+            agent_market["consumption"] = drop_timeslice(
                 agent_market.consumption * a.quantity
-            ).drop_vars(["timeslice", "month", "day", "hour"])
+            )
             included = [
                 i
                 for i in agent_market["commodity"].values
@@ -607,9 +607,9 @@ def sector_consumption(
             capacity = a.filter_input(a.assets.capacity, year=output_year).fillna(0.0)
             technologies = a.filter_input(techs, year=output_year).fillna(0.0)
             agent_market = market.sel(year=output_year).copy()
-            agent_market["consumption"] = (
+            agent_market["consumption"] = drop_timeslice(
                 agent_market.consumption * a.quantity
-            ).drop_vars(["timeslice", "month", "day", "hour"])
+            )
             included = [
                 i
                 for i in agent_market["commodity"].values

src/muse/sectors/preset_sector.py

@@ -7,6 +7,7 @@
 from xarray import DataArray, Dataset
 
 from muse.sectors.register import AbstractSector, register_sector
+from muse.timeslices import drop_timeslice
 
 
 @register_sector(name=("preset", "presets"))
@@ -77,15 +78,15 @@ def factory(cls, name: str, settings: Any) -> PresetSector:
                 assert consumption.region.isin(shares.region).all()
                 if "timeslice" in shares.dims:
                     ts = shares.timeslice
-                    shares = shares.drop_vars(["timeslice", "month", "day", "hour"])
+                    shares = drop_timeslice(shares)
                     consumption = (shares * consumption).assign_coords(timeslice=ts)
                 else:
                     consumption = consumption * shares.sel(
                         region=consumption.region, commodity=consumption.commodity
                     )
-            presets["consumption"] = consumption.drop_vars(
-                ["timeslice", "month", "day", "hour"]
-            ).assign_coords(timeslice=timeslice)
+            presets["consumption"] = drop_timeslice(consumption).assign_coords(
+                timeslice=timeslice
+            )
 
         if getattr(sector_conf, "supply_path", None) is not None:
             supply = read_csv_outputs(sector_conf.supply_path)
@@ -115,9 +116,7 @@ def factory(cls, name: str, settings: Any) -> PresetSector:
         for component in components:
             others = components.intersection(presets.data_vars).difference({component})
             if component not in presets and len(others) > 0:
-                presets[component] = zeros_like(presets[others.pop()]).drop_vars(
-                    ["timeslice", "month", "day", "hour"]
-                )
+                presets[component] = drop_timeslice(zeros_like(presets[others.pop()]))
         # add timeslice, if missing
         for component in {"supply", "consumption"}:
             if "timeslice" not in presets[component].dims:
@@ -166,9 +165,9 @@ def next(self, mca_market: Dataset) -> Dataset:
             mca_market.timeslice,
             QuantityType.EXTENSIVE,
         )
-        result["costs"] = convert_timeslice(
-            costs, mca_market.timeslice, QuantityType.INTENSIVE
-        ).drop_vars(["timeslice", "month", "day", "hour"])
+        result["costs"] = drop_timeslice(
+            convert_timeslice(costs, mca_market.timeslice, QuantityType.INTENSIVE)
+        )
         assert isinstance(result, Dataset)
         return result
 

src/muse/sectors/subsector.py

@@ -11,6 +11,7 @@
 import xarray as xr
 
 from muse.agents import Agent
+from muse.timeslices import drop_timeslice
 
 
 class Subsector:
@@ -57,9 +58,9 @@ def invest(
             current_year = market.year.min()
         if self.expand_market_prices:
             market = market.copy()
-            market["prices"] = np.maximum(
-                market.prices, market.prices.rename(region="dst_region")
-            ).drop_vars(["timeslice", "month", "day", "hour"])
+            market["prices"] = drop_timeslice(
+                np.maximum(market.prices, market.prices.rename(region="dst_region"))
+            )
 
         for agent in self.agents:
             agent.asset_housekeeping()

src/muse/timeslices.py

@@ -596,4 +596,15 @@ def represent_hours(
     return convert_timeslice(DataArray([nhours]), timeslices).squeeze()
 
 
+def drop_timeslice(data: DataArray) -> DataArray:
+    """Drop the timeslice variable from a DataArray.
+
+    If the array doesn't contain the timeslice variable, return the input unchanged.
+    """
+    if "timeslice" not in data.dims:
+        return data
+
+    return data.drop_vars(data.timeslice.indexes)
+
+
 setup_module(DEFAULT_TIMESLICE_DESCRIPTION)

tests/test_constraints.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
+from muse.timeslices import drop_timeslice
 from pytest import approx, fixture
 
 
@@ -419,9 +420,7 @@ def _as_list(data: Union[xr.DataArray, xr.Dataset]) -> Union[xr.DataArray, xr.Da
             data.get_index("timeslice"), names=("month", "day", "hour")
         )
         mindex_coords = xr.Coordinates.from_pandas_multiindex(index, "timeslice")
-        data = data.drop_vars(["timeslice", "month", "day", "hour"]).assign_coords(
-            mindex_coords
-        )
+        data = drop_timeslice(data).assign_coords(mindex_coords)
     return data
 
 

tests/test_demand_share.py

@@ -1,4 +1,5 @@
 import xarray as xr
+from muse.timeslices import drop_timeslice
 from pytest import approx, fixture, raises
 
 
@@ -25,9 +26,9 @@ def _matching_market(technologies, stock, timeslice):
         QuantityType.EXTENSIVE,
     )
     market["supply"] = production.sum("asset")
-    market["consumption"] = (
+    market["consumption"] = drop_timeslice(
         consumption(technologies, production).sum("asset") + market.supply
-    ).drop_vars(["timeslice", "month", "day", "hour"])
+    )
     market["prices"] = market.supply.dims, random(market.supply.shape)
 
     return market