Simplify trade model (not yet working)

src/muse/agents/__init__.py

@@ -2,10 +2,9 @@
     "AbstractAgent",
     "Agent",
     "InvestingAgent",
-    "factory",
     "agents_factory",
     "create_agent",
 ]
 
 from muse.agents.agent import AbstractAgent, Agent, InvestingAgent
-from muse.agents.factories import agents_factory, create_agent, factory
+from muse.agents.factories import agents_factory, create_agent

src/muse/agents/agent.py

@@ -80,14 +80,9 @@ def next(
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        time_period: int = 1,
-    ):
-        """Iterates agent one turn.
-
-        The goal is to figure out from market variables which technologies to invest in
-        and by how much.
-        """
-        pass
+        time_period: int,
+    ) -> None:
+        """Increments agent to the next time point (e.g. performing investments)."""
 
     def __repr__(self):
         return (
@@ -98,10 +93,7 @@ def __repr__(self):
 
 
 class Agent(AbstractAgent):
-    """Agent that is capable of computing a search-space and a cost metric.
-
-    This agent will not perform any investment itself.
-    """
+    """Standard agent that does not perform investments."""
 
     def __init__(
         self,
@@ -124,7 +116,7 @@ def __init__(
         spend_limit: int = 0,
         **kwargs,
     ):
-        """Creates a standard buildings agent.
+        """Creates a standard agent.
 
         Arguments:
             name: Name of the agent, used for cross-refencing external tables
@@ -166,10 +158,7 @@ def __init__(
         )
 
         self.year = year
-        """ Current year.
-
-        The year is incremented by one every time next is called.
-        """
+        """ Current year. Incremented by one every time next is called."""
         self.forecast = forecast
         """Number of years to look into the future for forecating purposed."""
         if search_rules is None:
@@ -250,8 +239,46 @@ def next(
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        time_period: int = 1,
-    ) -> Optional[xr.Dataset]:
+        time_period: int,
+    ) -> None:
+        self.year += time_period
+
+
+class InvestingAgent(Agent):
+    """Agent that performs investment for itself."""
+
+    def __init__(
+        self,
+        *args,
+        constraints: Optional[Callable] = None,
+        investment: Optional[Callable] = None,
+        **kwargs,
+    ):
+        """Creates an investing agent.
+
+        Arguments:
+            *args: See :py:class:`~muse.agents.agent.Agent`
+            constraints: Set of constraints limiting investment
+            investment: A function to perform investments
+            **kwargs: See :py:class:`~muse.agents.agent.Agent`
+        """
+        from muse.constraints import factory as csfactory
+        from muse.investments import factory as ifactory
+
+        super().__init__(*args, **kwargs)
+
+        self.invest = investment or ifactory()
+        """Method to use when fulfilling demand from rated set of techs."""
+        self.constraints = constraints or csfactory()
+        """Creates a set of constraints limiting investment."""
+
+    def next(
+        self,
+        technologies: xr.Dataset,
+        market: xr.Dataset,
+        demand: xr.DataArray,
+        time_period: int,
+    ) -> None:
         """Iterates agent one turn.
 
         The goal is to figure out from market variables which technologies to
@@ -263,22 +290,75 @@ def next(
         """
         from logging import getLogger
 
-        # dataset with intermediate computational results from search
-        # makes it easier to pass intermediate results to functions, as well as
-        # filter them when inside a function
+        current_year = self.year
+
+        # Skip forward if demand is zero
         if demand.size == 0 or demand.sum() < 1e-12:
             self.year += time_period
             return None
 
+        # Calculate the search space
         search_space = (
             self.search_rules(self, demand, technologies, market).fillna(0).astype(int)
         )
 
+        # Skip forward if the search space is empty
         if any(u == 0 for u in search_space.shape):
             getLogger(__name__).critical("Search space is empty")
             self.year += time_period
             return None
 
+        # Calculate the decision metric
+        decision = self.compute_decision(technologies, market, demand, search_space)
+        search = xr.Dataset(dict(search_space=search_space, decision=decision))
+        if "timeslice" in search.dims:
+            search["demand"] = drop_timeslice(demand)
+        else:
+            search["demand"] = demand
+
+        # Filter assets with demand
+        not_assets = [u for u in search.demand.dims if u != "asset"]
+        condtechs = (
+            search.demand.sum(not_assets) > getattr(self, "tolerance", 1e-8)
+        ).values
+        search = search.sel(asset=condtechs)
+
+        # Calculate constraints
+        constraints = self.constraints(
+            search.demand,
+            self.assets,
+            search.search_space,
+            market,
+            technologies,
+            year=current_year,
+        )
+
+        # Calculate investments
+        investments = self.invest(
+            search[["search_space", "decision"]],
+            technologies,
+            constraints,
+            year=current_year,
+        )
+
+        # Add investments
+        self.add_investments(
+            technologies,
+            investments,
+            current_year=current_year,
+            time_period=time_period,
+        )
+
+        # Increment the year
+        self.year += time_period
+
+    def compute_decision(
+        self,
+        technologies: xr.Dataset,
+        market: xr.Dataset,
+        demand: xr.DataArray,
+        search_space: xr.DataArray,
+    ) -> xr.DataArray:
         # Filter technologies according to the search space, forecast year and region
         techs = self.filter_input(
             technologies,
@@ -297,23 +377,12 @@ def next(
         # Filter prices according to the region
         prices = self.filter_input(market.prices)
 
-        # Compute the objective
-        decision = self._compute_objective(
+        # Compute the objectives
+        objectives = self.objectives(
             technologies=techs, demand=reduced_demand, prices=prices
         )
 
-        self.year += time_period
-        return xr.Dataset(dict(search_space=search_space, decision=decision))
-
-    def _compute_objective(
-        self,
-        technologies: xr.Dataset,
-        demand: xr.DataArray,
-        prices: xr.DataArray,
-    ) -> xr.DataArray:
-        objectives = self.objectives(
-            technologies=technologies, demand=demand, prices=prices
-        )
+        # Compute the decision metric
         decision = self.decision(objectives)
         return decision
 
@@ -323,19 +392,20 @@ def add_investments(
         investments: xr.DataArray,
         current_year: int,
         time_period: int,
-    ):
+    ) -> None:
         """Add new assets to the agent."""
+        # Calculate retirement profile of new assets
         new_capacity = self.retirement_profile(
             technologies, investments, current_year, time_period
         )
-
         if new_capacity is None:
             return
         new_capacity = new_capacity.drop_vars(
             set(new_capacity.coords) - set(self.assets.coords)
         )
         new_assets = xr.Dataset(dict(capacity=new_capacity))
 
+        # Merge new assets with existing assets
         self.assets = self.merge_transform(self.assets, new_assets)
 
     def retirement_profile(
@@ -347,10 +417,13 @@ def retirement_profile(
     ) -> Optional[xr.DataArray]:
         from muse.investments import cliff_retirement_profile
 
+        # Sum investments
         if "asset" in investments.dims:
             investments = investments.sum("asset")
         if "agent" in investments.dims:
             investments = investments.squeeze("agent", drop=True)
+
+        # Filter out investments below the threshold
         investments = investments.sel(
             replacement=(investments > self.asset_threshold).any(
                 [d for d in investments.dims if d != "replacement"]
@@ -359,22 +432,22 @@ def retirement_profile(
         if investments.size == 0:
             return None
 
-        # figures out the retirement profile for the new investments
+        # Calculate the retirement profile for new investments
+        # Note: technical life must be at least the length of the time period
         lifetime = self.filter_input(
             technologies.technical_life,
             year=current_year,
             technology=investments.replacement,
-        )
+        ).clip(min=time_period)
         profile = cliff_retirement_profile(
-            lifetime.clip(min=time_period),
-            current_year=current_year + time_period,
-            protected=max(self.forecast - time_period - 1, 0),
+            lifetime,
+            investment_year=current_year + time_period,
         )
         if "dst_region" in investments.coords:
             investments = investments.reindex_like(profile, method="ffill")
 
+        # Apply the retirement profile to the investments
         new_assets = (investments * profile).rename(replacement="asset")
-
         new_assets["installed"] = "asset", [current_year] * len(new_assets.asset)
 
         # The new assets have picked up quite a few coordinates along the way.
@@ -383,89 +456,3 @@ def retirement_profile(
             new, old = new_assets.dims, self.assets.dims
             raise RuntimeError(f"Asset dimensions do not match: {new} vs {old}")
         return new_assets
-
-
-class InvestingAgent(Agent):
-    """Agent that performs investment for itself."""
-
-    def __init__(
-        self,
-        *args,
-        constraints: Optional[Callable] = None,
-        investment: Optional[Callable] = None,
-        **kwargs,
-    ):
-        """Creates a standard buildings agent.
-
-        Arguments:
-            *args: See :py:class:`~muse.agents.agent.Agent`
-            constraints: Set of constraints limiting investment
-            investment: A function to perform investments
-            **kwargs: See :py:class:`~muse.agents.agent.Agent`
-        """
-        from muse.constraints import factory as csfactory
-        from muse.investments import factory as ifactory
-
-        super().__init__(*args, **kwargs)
-
-        if investment is None:
-            investment = ifactory()
-        self.invest = investment
-        """Method to use when fulfilling demand from rated set of techs."""
-        if not callable(constraints):
-            constraints = csfactory()
-        self.constraints = constraints
-        """Creates a set of constraints limiting investment."""
-
-    def next(
-        self,
-        technologies: xr.Dataset,
-        market: xr.Dataset,
-        demand: xr.DataArray,
-        time_period: int = 1,
-    ):
-        """Iterates agent one turn.
-
-        The goal is to figure out from market variables which technologies to
-        invest in and by how much.
-
-        This function will modify `self.assets` and increment `self.year`.
-        Other attributes are left unchanged. Arguments to the function are
-        never modified.
-        """
-        current_year = self.year
-        search = super().next(technologies, market, demand, time_period=time_period)
-        if search is None:
-            return None
-
-        if "timeslice" in search.dims:
-            search["demand"] = drop_timeslice(demand)
-        else:
-            search["demand"] = demand
-        not_assets = [u for u in search.demand.dims if u != "asset"]
-        condtechs = (
-            search.demand.sum(not_assets) > getattr(self, "tolerance", 1e-8)
-        ).values
-        search = search.sel(asset=condtechs)
-        constraints = self.constraints(
-            search.demand,
-            self.assets,
-            search.search_space,
-            market,
-            technologies,
-            year=current_year,
-        )
-
-        investments = self.invest(
-            search[["search_space", "decision"]],
-            technologies,
-            constraints,
-            year=current_year,
-        )
-
-        self.add_investments(
-            technologies,
-            investments,
-            current_year=self.year - time_period,
-            time_period=time_period,
-        )