"""Pre and post hooks on agents."""
__all__ = [
"register_initial_asset_transform",
"register_final_asset_transform",
"noop",
"clean",
"old_assets_only",
"merge_assets",
"new_assets_only",
"housekeeping_factory",
"asset_merge_factory",
]
from typing import Callable, Mapping, MutableMapping, Text, Union
from xarray import Dataset
from muse.agents import Agent
from muse.registration import registrator
INITIAL_ASSET_TRANSFORM: MutableMapping[Text, Callable] = {}
""" Transform at the start of each step. """
FINAL_ASSET_TRANSFORM: MutableMapping[Text, Callable] = {}
""" Transform at the end of each step, including new assets. """
[docs]
def housekeeping_factory(settings: Union[Text, Mapping] = "noop") -> Callable:
"""Returns a function for performing initial housekeeping.
For instance, remove technologies with no capacity now or in the future.
Available housekeeping functions should be registered with
:py:func:`@register_initial_asset_transform<register_initial_asset_transform>`.
"""
from muse.agents import AbstractAgent
if isinstance(settings, Text):
name = settings
params: Mapping = {}
else:
params = {k: v for k, v in settings.items() if k != "name"}
name = settings["name"]
transform = INITIAL_ASSET_TRANSFORM[name]
def initial_assets_transform(agent: AbstractAgent, assets: Dataset) -> Dataset:
return transform(agent, assets, **params)
return initial_assets_transform
[docs]
def asset_merge_factory(settings: Union[Text, Mapping] = "new") -> Callable:
"""Returns a function for merging new investments into assets.
Available merging functions should be registered with
:py:func:`@register_final_asset_transform<register_final_asset_transform>`.
"""
"""Returns a function for performing initial housekeeping.
For instance, remove technologies with no capacity now or in the future.
Available housekeeping functions should be registered with
:py:func:`@register_initial_asset_transform<register_initial_asset_transform>`.
"""
if isinstance(settings, Text):
name = settings
params: Mapping = {}
else:
params = {k: v for k, v in settings.items() if k != "name"}
name = settings["name"]
transform = FINAL_ASSET_TRANSFORM[name]
def final_assets_transform(old_assets: Dataset, new_assets):
return transform(old_assets, new_assets, **params)
final_assets_transform.__name__ = name
return final_assets_transform
[docs]
@register_initial_asset_transform(name="default")
def noop(agent: Agent, assets: Dataset) -> Dataset:
"""Return assets as they are."""
return assets
[docs]
@register_initial_asset_transform
def clean(agent: Agent, assets: Dataset) -> Dataset:
"""Removes empty assets."""
from muse.utilities import clean_assets
years = [agent.year, agent.forecast_year]
return clean_assets(assets, years)
[docs]
@register_final_asset_transform(name="new")
def new_assets_only(old_assets: Dataset, new_assets: Dataset) -> Dataset:
"""Returns newly invested assets and ignores old assets."""
return new_assets
[docs]
@register_final_asset_transform(name="old")
def old_assets_only(old_assets: Dataset, new_assets: Dataset) -> Dataset:
"""Returns old assets and ignores newly invested assets."""
return old_assets
[docs]
@register_final_asset_transform(name="merge")
def merge_assets(old_assets: Dataset, new_assets: Dataset) -> Dataset:
"""Adds new assets to old along asset dimension.
New assets are assumed to be nonequivalent to any old_assets. Indeed,
it is expected that the asset dimension does not have coordinates (i.e.
it is a combination of coordinates, such as technology and installation
year).
After merging the new assets, quantities are back-filled along the year
dimension. Further missing values (i.e. future years the old_assets
did not take into account) are set to zero.
"""
from logging import getLogger
from muse.utilities import merge_assets
assert "asset" not in old_assets
if "asset" not in new_assets.dims and "replacement" in new_assets.dims:
new_assets = new_assets.rename(replacement="asset")
if len(new_assets.capacity) == 0:
getLogger(__name__).critical("there are no new assets")
return old_assets
return merge_assets(old_assets, new_assets)