diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py index 8146711b9..be7aac7cd 100644 --- a/src/muse/agents/agent.py +++ b/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"] diff --git a/src/muse/constraints.py b/src/muse/constraints.py index 16c146f9e..c69fd5a18 100644 --- a/src/muse/constraints.py +++ b/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}) diff --git a/src/muse/demand_matching.py b/src/muse/demand_matching.py index 0a9c94d5f..9640bc74e 100644 --- a/src/muse/demand_matching.py +++ b/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))) diff --git a/src/muse/examples.py b/src/muse/examples.py index 115f39d28..e75823db7 100644 --- a/src/muse/examples.py +++ b/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( diff --git a/src/muse/mca.py b/src/muse/mca.py index e7716e562..8af0954ae 100644 --- a/src/muse/mca.py +++ b/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( diff --git a/src/muse/objectives.py b/src/muse/objectives.py index 87af8f9a3..998437fdf 100644 --- a/src/muse/objectives.py +++ b/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 diff --git a/src/muse/outputs/mca.py b/src/muse/outputs/mca.py index c93480bf7..1c0166515 100644 --- a/src/muse/outputs/mca.py +++ b/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 diff --git a/src/muse/sectors/preset_sector.py b/src/muse/sectors/preset_sector.py index 2bdbbdec8..8028d8d51 100644 --- a/src/muse/sectors/preset_sector.py +++ b/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 diff --git a/src/muse/sectors/subsector.py b/src/muse/sectors/subsector.py index 846b379f5..c6d35ef01 100644 --- a/src/muse/sectors/subsector.py +++ b/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() diff --git a/src/muse/timeslices.py b/src/muse/timeslices.py index 8f3e791c5..547d09be1 100644 --- a/src/muse/timeslices.py +++ b/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) diff --git a/tests/test_constraints.py b/tests/test_constraints.py index 90794ea5d..6f8d1712b 100644 --- a/tests/test_constraints.py +++ b/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 diff --git a/tests/test_demand_share.py b/tests/test_demand_share.py index 6e835f7ab..549407896 100644 --- a/tests/test_demand_share.py +++ b/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 diff --git a/tests/test_mca.py b/tests/test_mca.py index ee07782bb..8db98378b 100644 --- a/tests/test_mca.py +++ b/tests/test_mca.py @@ -1,6 +1,7 @@ from collections.abc import Sequence from muse.commodities import CommodityUsage +from muse.timeslices import drop_timeslice from xarray import Dataset @@ -16,20 +17,14 @@ def test_check_equilibrium(market: Dataset): new_market = market.copy(deep=True) assert check_equilibrium(new_market, market, tol, equilibrium_variable) - new_market["supply"] = ( - new_market["supply"].drop_vars(["timeslice", "month", "day", "hour"]) - + tol * 1.5 - ) + new_market["supply"] = drop_timeslice(new_market["supply"]) + tol * 1.5 assert not check_equilibrium(new_market, market, tol, equilibrium_variable) equilibrium_variable = "prices" assert check_equilibrium(new_market, market, tol, equilibrium_variable) - new_market["prices"] = ( - new_market["prices"].drop_vars(["timeslice", "month", "day", "hour"]) - + tol * 1.5 - ) + new_market["prices"] = drop_timeslice(new_market["prices"]) + tol * 1.5 assert not check_equilibrium(new_market, market, tol, equilibrium_variable) @@ -39,17 +34,12 @@ def test_check_demand_fulfillment(market): tolerance_unmet_demand = -0.1 - market["supply"] = market.consumption.copy(deep=True).drop_vars( - ["timeslice", "month", "day", "hour"] - ) + market["supply"] = drop_timeslice(market.consumption.copy(deep=True)) assert check_demand_fulfillment( market, tolerance_unmet_demand, ) - market["supply"] = ( - market["supply"].drop_vars(["timeslice", "month", "day", "hour"]) - + tolerance_unmet_demand * 1.5 - ) + market["supply"] = drop_timeslice(market["supply"]) + tolerance_unmet_demand * 1.5 assert not check_demand_fulfillment( market, tolerance_unmet_demand, diff --git a/tests/test_quantities.py b/tests/test_quantities.py index 74439b1b0..686a2c14b 100644 --- a/tests/test_quantities.py +++ b/tests/test_quantities.py @@ -2,6 +2,7 @@ import numpy as np import xarray as xr +from muse.timeslices import drop_timeslice from pytest import approx, fixture @@ -460,7 +461,7 @@ def test_costed_production_exact_match(market, capacity, technologies): market, QuantityType.EXTENSIVE, ) - market["consumption"] = maxdemand.drop_vars(["timeslice", "month", "day", "hour"]) + market["consumption"] = drop_timeslice(maxdemand) result = costed_production(market.consumption, costs, capacity, technologies) assert isinstance(result, xr.DataArray) actual = xr.Dataset(dict(r=result)).groupby("region").sum("asset").r @@ -487,9 +488,7 @@ def test_costed_production_single_region(market, capacity, technologies): market, QuantityType.EXTENSIVE, ) - market["consumption"] = (0.9 * maxdemand).drop_vars( - ["timeslice", "month", "day", "hour"] - ) + market["consumption"] = drop_timeslice(0.9 * maxdemand) technodata = broadcast_techs(technologies, capacity) costs = annual_levelized_cost_of_energy( market.prices.sel(region=technodata.region), technodata @@ -522,9 +521,7 @@ def test_costed_production_single_year(market, capacity, technologies): market, QuantityType.EXTENSIVE, ) - market["consumption"] = (0.9 * maxdemand).drop_vars( - ["timeslice", "month", "day", "hour"] - ) + market["consumption"] = drop_timeslice(0.9 * maxdemand) technodata = broadcast_techs(technologies, capacity) costs = annual_levelized_cost_of_energy( market.prices.sel(region=technodata.region), technodata @@ -560,9 +557,7 @@ def test_costed_production_over_capacity(market, capacity, technologies): market, QuantityType.EXTENSIVE, ) - market["consumption"] = (maxdemand * 0.9).drop_vars( - ["timeslice", "month", "day", "hour"] - ) + market["consumption"] = drop_timeslice(maxdemand * 0.9) technodata = broadcast_techs(technologies, capacity) costs = annual_levelized_cost_of_energy( market.prices.sel(region=technodata.region), technodata @@ -598,9 +593,7 @@ def test_costed_production_with_minimum_service(market, capacity, technologies, ) minprod = maxprod * broadcast_techs(technologies.minimum_service_factor, maxprod) maxdemand = xr.Dataset(dict(mp=minprod)).groupby("region").sum("asset").mp - market["consumption"] = (maxdemand * 0.9).drop_vars( - ["timeslice", "month", "day", "hour"] - ) + market["consumption"] = drop_timeslice(maxdemand * 0.9) technodata = broadcast_techs(technologies, capacity) costs = annual_levelized_cost_of_energy( market.prices.sel(region=technodata.region), technodata diff --git a/tests/test_timeslices.py b/tests/test_timeslices.py index 49d548f05..7134d8c22 100644 --- a/tests/test_timeslices.py +++ b/tests/test_timeslices.py @@ -40,7 +40,7 @@ def transforms(toml, reference): @fixture -def rough(reference): +def timeslice_dataarray(reference): from pandas import MultiIndex return DataArray( @@ -59,28 +59,34 @@ def rough(reference): ) -def test_convert_extensive_timeslice(reference, rough, transforms): - z = convert_timeslice(rough, reference, finest=reference, transforms=transforms) +def test_convert_extensive_timeslice(reference, timeslice_dataarray, transforms): + z = convert_timeslice( + timeslice_dataarray, reference, finest=reference, transforms=transforms + ) assert z.shape == reference.shape assert z.values == approx( [ - float(rough[0] * reference[0] / (reference[0] + reference[1])), - float(rough[0] * reference[1] / (reference[0] + reference[1])), + float( + timeslice_dataarray[0] * reference[0] / (reference[0] + reference[1]) + ), + float( + timeslice_dataarray[0] * reference[1] / (reference[0] + reference[1]) + ), 0, 0, - float(rough[1]), + float(timeslice_dataarray[1]), 0, 0, 0, - float(rough[2]), + float(timeslice_dataarray[2]), 0, ] ) -def test_convert_intensive_timeslice(reference, rough, transforms): +def test_convert_intensive_timeslice(reference, timeslice_dataarray, transforms): z = convert_timeslice( - rough, + timeslice_dataarray, reference, finest=reference, transforms=transforms, @@ -89,15 +95,15 @@ def test_convert_intensive_timeslice(reference, rough, transforms): assert z.values == approx( [ - float(rough[0]), - float(rough[0]), + float(timeslice_dataarray[0]), + float(timeslice_dataarray[0]), 0, 0, - float(rough[1]), + float(timeslice_dataarray[1]), 0, 0, 0, - float(rough[2]), + float(timeslice_dataarray[2]), 0, ] ) @@ -241,3 +247,17 @@ def to_bitstring(x): assert ( to_bitstring(vectors[("springautumn", "week", "night")]) == "0101000001010000" ) + + +def test_drop_timeslice(timeslice_dataarray): + from muse.timeslices import drop_timeslice + + dropped = drop_timeslice(timeslice_dataarray) + coords_to_check = {"timeslice", "semester", "week", "day"} + assert coords_to_check.issubset(timeslice_dataarray.coords) + assert not coords_to_check.intersection(dropped.coords) + + # Test on arrays without timeslice data + data_without_timeslice = DataArray([1, 2, 3], dims=["x"]) + assert drop_timeslice(data_without_timeslice).equals(data_without_timeslice) + assert drop_timeslice(dropped).equals(dropped)