From fa59017b69d00f2f545ca3aaeb512e89ccdbbdaa Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Wed, 2 Oct 2024 17:10:56 +0100
Subject: [PATCH 01/18] Add comments and simplify sector.next

---
 src/muse/mca.py            | 16 +++++------
 src/muse/sectors/sector.py | 59 +++++++++++---------------------------
 2 files changed, 25 insertions(+), 50 deletions(-)

diff --git a/src/muse/mca.py b/src/muse/mca.py
index 55014c70..7fa2cd48 100644
--- a/src/muse/mca.py
+++ b/src/muse/mca.py
@@ -310,6 +310,7 @@ def run(self) -> None:
                 )
                 self.carbon_price = future_propagation(self.carbon_price, future_price)
 
+            # Solve the market
             _, new_market, self.sectors = self.find_equilibrium(new_market)
 
             # Save sector outputs
@@ -324,17 +325,19 @@ def run(self) -> None:
                     new_market, year_idx
                 )
 
+            # Update the market
             dims = {i: new_market[i] for i in new_market.dims}
             self.market.supply.loc[dims] = new_market.supply
             self.market.consumption.loc[dims] = new_market.consumption
-
             dims = {i: new_market[i] for i in new_market.prices.dims if i != "year"}
             self.market.prices.loc[dims] = future_propagation(
                 self.market.prices.sel(dims), new_market.prices.sel(year=years[1])
             )
 
+            # Global outputs
             self.outputs(self.market, self.sectors, year=self.time_framework[year_idx])  # type: ignore
             self.outputs_cache.consolidate_cache(year=self.time_framework[year_idx])
+
             getLogger(__name__).info(
                 f"Finish simulation year {years[0]} ({year_idx+1}/{nyear})!"
             )
@@ -429,29 +432,26 @@ def single_year_iteration(
     if "updated_prices" not in market.data_vars:
         market["updated_prices"] = drop_timeslice(market.prices.copy())
 
-    # eventually, the first market should be one that creates the initial demand
     for sector in sectors:
+        # Solve the sector
         sector_market = sector.next(
             market[["supply", "consumption", "prices"]]  # type:ignore
         )
-
         sector_market = sector_market.sel(year=market.year)
 
+        # Calculate net consumption
         dims = {i: sector_market[i] for i in sector_market.consumption.dims}
-
         sector_market.consumption.loc[dims] = (
             sector_market.consumption.loc[dims] - sector_market.supply.loc[dims]
         ).clip(min=0.0, max=None)
 
+        # Update market supply and consumption
         market.consumption.loc[dims] += sector_market.consumption
-
         dims = {i: sector_market[i] for i in sector_market.supply.dims}
         market.supply.loc[dims] += sector_market.supply
 
+        # Update market prices
         costs = sector_market.costs.sel(commodity=is_enduse(sector_market.comm_usage))
-
-        # do not write costs lower than 1e-4
-        # should correspond to rounding value
         if len(costs.commodity) > 0:
             costs = costs.where(costs > 1e-4, 0)
             dims = {i: costs[i] for i in costs.dims}
diff --git a/src/muse/sectors/sector.py b/src/muse/sectors/sector.py
index b25a7455..0a6328ed 100644
--- a/src/muse/sectors/sector.py
+++ b/src/muse/sectors/sector.py
@@ -26,7 +26,6 @@ def factory(cls, name: str, settings: Any) -> Sector:
         from muse.interactions import factory as interaction_factory
         from muse.outputs.sector import factory as ofactory
         from muse.production import factory as pfactory
-        from muse.readers import read_timeslices
         from muse.readers.toml import read_technodata
         from muse.utilities import nametuple_to_dict
 
@@ -34,10 +33,6 @@ def factory(cls, name: str, settings: Any) -> Sector:
         for attribute in ("name", "type", "priority", "path"):
             sector_settings.pop(attribute, None)
 
-        timeslices = read_timeslices(
-            sector_settings.pop("timeslice_levels", None)
-        ).get_index("timeslice")
-
         technologies = read_technodata(settings, name, settings.time_framework)
 
         if "subsectors" not in sector_settings:
@@ -81,7 +76,6 @@ def factory(cls, name: str, settings: Any) -> Sector:
             name,
             technologies,
             subsectors=subsectors,
-            timeslices=timeslices,
             supply_prod=supply,
             outputs=outputs,
             interactions=interactions,
@@ -93,8 +87,6 @@ def __init__(
         name: str,
         technologies: xr.Dataset,
         subsectors: Sequence[Subsector] = [],
-        timeslices: pd.MultiIndex | None = None,
-        technodata_timeslices: xr.Dataset = None,
         interactions: Callable[[Sequence[AbstractAgent]], None] | None = None,
         interpolation: str = "linear",
         outputs: Callable | None = None,
@@ -110,11 +102,6 @@ def __init__(
         """Subsectors controlled by this object."""
         self.technologies: xr.Dataset = technologies
         """Parameters describing the sector's technologies."""
-        self.timeslices: pd.MultiIndex | None = timeslices
-        """Timeslice at which this sector operates.
-
-        If None, it will operate using the timeslice of the input market.
-        """
         self.interpolation: Mapping[str, Any] = {
             "method": interpolation,
             "kwargs": {"fill_value": "extrapolate"},
@@ -201,41 +188,25 @@ def group_assets(x: xr.DataArray) -> xr.DataArray:
             current_year = int(mca_market.year.min())
         getLogger(__name__).info(f"Running {self.name} for year {current_year}")
 
-        # > to sector timeslice
-        market = self.convert_market_timeslice(
-            mca_market.sel(
-                commodity=self.technologies.commodity, region=self.technologies.region
-            ).interp(
-                year=sorted(
-                    {
-                        current_year,
-                        current_year + time_period,
-                        current_year + self.forecast,
-                    }
-                ),
-                **self.interpolation,
-            ),
-            self.timeslices,
-        )
-        # > agent interactions
+        # Agent interactions
         self.interactions(list(self.agents))
-        # > investment
-        years = sorted(
-            set(
-                market.year.data.tolist()
-                + self.capacity.installed.data.tolist()
-                + self.technologies.year.data.tolist()
-            )
+
+        # Select appropriate data from the market
+        market = mca_market.sel(
+            commodity=self.technologies.commodity, region=self.technologies.region
         )
-        technologies = self.technologies.interp(year=years, **self.interpolation)
 
+        # Investments
         for subsector in self.subsectors:
             subsector.invest(
-                technologies, market, time_period=time_period, current_year=current_year
+                self.technologies,
+                market,
+                time_period=time_period,
+                current_year=current_year,
             )
 
         # Full output data
-        supply, consume, costs = self.market_variables(market, technologies)
+        supply, consume, costs = self.market_variables(market, self.technologies)
         self.output_data = xr.Dataset(
             dict(
                 supply=supply,
@@ -287,7 +258,9 @@ def group_assets(x: xr.DataArray) -> xr.DataArray:
                 dict(supply=supply, consumption=consumption, costs=costs)
             )
         result = self.convert_market_timeslice(result, mca_market.timeslice)
-        result["comm_usage"] = technologies.comm_usage.sel(commodity=result.commodity)
+        result["comm_usage"] = self.technologies.comm_usage.sel(
+            commodity=result.commodity
+        )
         result.set_coords("comm_usage")
         return result
 
@@ -306,15 +279,17 @@ def market_variables(self, market: xr.Dataset, technologies: xr.Dataset) -> Any:
         years = market.year.values
         capacity = self.capacity.interp(year=years, **self.interpolation)
 
+        # Calculate supply
         supply = self.supply_prod(
             market=market, capacity=capacity, technologies=technologies
         )
-
         if "timeslice" in market.prices.dims and "timeslice" not in supply.dims:
             supply = convert_timeslice(supply, market.timeslice, QuantityType.EXTENSIVE)
 
+        # Calculate consumption
         consume = consumption(technologies, supply, market.prices)
 
+        # Calculate commodity prices
         technodata = cast(xr.Dataset, broadcast_techs(technologies, supply))
         costs = supply_cost(
             supply.where(~is_pollutant(supply.comm_usage), 0),

From 697ff3f30a2f262a6f1cac39c0de7853298e7b54 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Thu, 3 Oct 2024 11:27:03 +0100
Subject: [PATCH 02/18] Simplify agent module, more comments

---
 src/muse/agents/agent.py      | 82 +++++++++++++-----------------
 src/muse/sectors/sector.py    | 15 ++++--
 src/muse/sectors/subsector.py | 93 ++++++++++++++++++-----------------
 3 files changed, 94 insertions(+), 96 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index fb34dffb..a6d45e30 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -245,40 +245,13 @@ def asset_housekeeping(self):
         #  state.
         self.assets = self._housekeeping(self, self.assets)
 
-    def next(
+    def compute_decision(
         self,
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        time_period: int = 1,
-    ) -> Optional[xr.Dataset]:
-        """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.
-        """
-        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
-        if demand.size == 0 or demand.sum() < 1e-12:
-            self.year += time_period
-            return None
-
-        search_space = (
-            self.search_rules(self, demand, technologies, market).fillna(0).astype(int)
-        )
-
-        if any(u == 0 for u in search_space.shape):
-            getLogger(__name__).critical("Search space is empty")
-            self.year += time_period
-            return None
-
+        search_space,
+    ):
         # Filter technologies according to the search space, forecast year and region
         techs = self.filter_input(
             technologies,
@@ -297,23 +270,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
 
@@ -433,20 +395,42 @@ def next(
         Other attributes are left unchanged. Arguments to the function are
         never modified.
         """
+        from logging import getLogger
+
         current_year = self.year
-        search = super().next(technologies, market, demand, time_period=time_period)
-        if search is None:
+
+        # 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,
@@ -456,6 +440,7 @@ def next(
             year=current_year,
         )
 
+        # Calculate investments
         investments = self.invest(
             search[["search_space", "decision"]],
             technologies,
@@ -463,9 +448,12 @@ def next(
             year=current_year,
         )
 
+        # Add investments
         self.add_investments(
             technologies,
             investments,
-            current_year=self.year - time_period,
+            current_year=current_year,
             time_period=time_period,
         )
+
+        self.year += time_period
diff --git a/src/muse/sectors/sector.py b/src/muse/sectors/sector.py
index 0a6328ed..1dc42dd5 100644
--- a/src/muse/sectors/sector.py
+++ b/src/muse/sectors/sector.py
@@ -29,16 +29,19 @@ def factory(cls, name: str, settings: Any) -> Sector:
         from muse.readers.toml import read_technodata
         from muse.utilities import nametuple_to_dict
 
+        # Read sector settings
         sector_settings = getattr(settings.sectors, name)._asdict()
         for attribute in ("name", "type", "priority", "path"):
             sector_settings.pop(attribute, None)
-
-        technologies = read_technodata(settings, name, settings.time_framework)
-
         if "subsectors" not in sector_settings:
             raise RuntimeError(f"Missing 'subsectors' section in sector {name}")
         if len(sector_settings["subsectors"]._asdict()) == 0:
             raise RuntimeError(f"Empty 'subsectors' section in sector {name}")
+
+        # Read technologies
+        technologies = read_technodata(settings, name, settings.time_framework)
+
+        # Create subsectors
         subsectors = [
             Subsector.factory(
                 subsec_settings,
@@ -51,14 +54,18 @@ def factory(cls, name: str, settings: Any) -> Sector:
             ._asdict()
             .items()
         ]
+
+        # Check that subsector commodities are disjoint
         are_disjoint_commodities = sum(len(s.commodities) for s in subsectors) == len(
             set().union(*(set(s.commodities) for s in subsectors))  # type: ignore
         )
         if not are_disjoint_commodities:
             raise RuntimeError("Subsector commodities are not disjoint")
 
+        # Create outputs
         outputs = ofactory(*sector_settings.pop("outputs", []), sector_name=name)
 
+        #
         supply_args = sector_settings.pop(
             "supply", sector_settings.pop("dispatch_production", {})
         )
@@ -68,8 +75,10 @@ def factory(cls, name: str, settings: Any) -> Sector:
             supply_args = nametuple_to_dict(supply_args)
         supply = pfactory(**supply_args)
 
+        # Create interactions
         interactions = interaction_factory(sector_settings.pop("interactions", None))
 
+        # Create sector
         for attr in ("technodata", "commodities_out", "commodities_in"):
             sector_settings.pop(attr, None)
         return cls(
diff --git a/src/muse/sectors/subsector.py b/src/muse/sectors/subsector.py
index 26ba1cb2..d3704517 100644
--- a/src/muse/sectors/subsector.py
+++ b/src/muse/sectors/subsector.py
@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-from collections.abc import Hashable, MutableMapping, Sequence
+from collections.abc import Sequence
 from typing import (
     Any,
     Callable,
-    cast,
 )
 
 import numpy as np
@@ -51,35 +50,35 @@ def invest(
         self,
         technologies: xr.Dataset,
         market: xr.Dataset,
-        time_period: int = 5,
-        current_year: int | None = None,
+        time_period: int,
+        current_year: int,
     ) -> None:
-        if current_year is None:
-            current_year = market.year.min()
+        # Expand prices to include destination region (for trade models)
         if self.expand_market_prices:
             market = market.copy()
             market["prices"] = drop_timeslice(
                 np.maximum(market.prices, market.prices.rename(region="dst_region"))
             )
 
+        # Agent housekeeping
         for agent in self.agents:
             agent.asset_housekeeping()
 
-        lp_problem = self.aggregate_lp(
-            technologies, market, time_period, current_year=current_year
-        )
-        if lp_problem is None:
-            return
+        # Perform the investment
+        self.aggregate_lp(technologies, market, time_period, current_year=current_year)
+        # if lp_problem is None:
+        #     return
 
-        years = technologies.year
-        techs = technologies.interp(year=years)
-        techs = techs.sel(year=current_year + time_period)
+        # # If there is a problem with the LP...
+        # years = technologies.year
+        # techs = technologies.interp(year=years)
+        # techs = techs.sel(year=current_year + time_period)
 
-        solution = self.investment(
-            search=lp_problem[0], technologies=techs, constraints=lp_problem[1]
-        )
+        # solution = self.investment(
+        #     search=lp_problem[0], technologies=techs, constraints=lp_problem[1]
+        # )
 
-        self.assign_back_to_agents(technologies, solution, current_year, time_period)
+        # self.assign_back_to_agents(technologies, solution, current_year, time_period)
 
     def assign_back_to_agents(
         self,
@@ -99,14 +98,12 @@ def aggregate_lp(
         self,
         technologies: xr.Dataset,
         market: xr.Dataset,
-        time_period: int = 5,
-        current_year: int | None = None,
-    ) -> tuple[xr.Dataset, Sequence[xr.Dataset]] | None:
+        time_period,
+        current_year,
+    ):
         from muse.utilities import agent_concatenation, reduce_assets
 
-        if current_year is None:
-            current_year = market.year.min()
-
+        # Split demand across agents
         demands = self.demand_share(
             self.agents,
             market,
@@ -122,42 +119,46 @@ def aggregate_lp(
                 dimension.
             """
             raise ValueError(msg)
-        agent_market = market.copy()
+
+        # Concatenate assets
         assets = agent_concatenation(
             {agent.uuid: agent.assets for agent in self.agents}
         )
+
+        # Calculate existing capacity
+        agent_market = market.copy()
         agent_market["capacity"] = (
             reduce_assets(assets.capacity, coords=("region", "technology"))
             .interp(year=market.year, method="linear", kwargs={"fill_value": 0.0})
             .swap_dims(dict(asset="technology"))
         )
 
-        agent_lps: MutableMapping[Hashable, xr.Dataset] = {}
+        # agent_lps: MutableMapping[Hashable, xr.Dataset] = {}
         for agent in self.agents:
             if "agent" in demands.coords:
                 share = demands.sel(asset=demands.agent == agent.uuid)
             else:
                 share = demands
-            result = agent.next(
-                technologies, agent_market, share, time_period=time_period
-            )
-            if result is not None:
-                agent_lps[agent.uuid] = result
-
-        if len(agent_lps) == 0:
-            return None
-
-        lps = cast(xr.Dataset, agent_concatenation(agent_lps, dim="agent"))
-        coords = {"agent", "technology", "region"}.intersection(assets.asset.coords)
-        constraints = self.constraints(
-            demand=demands,
-            assets=reduce_assets(assets, coords=coords).set_coords(coords),
-            search_space=lps.search_space,
-            market=market,
-            technologies=technologies,
-            year=current_year,
-        )
-        return lps, constraints
+
+            # Compute investments for the agent
+            agent.next(technologies, agent_market, share, time_period=time_period)
+        #     if result is not None:
+        #         agent_lps[agent.uuid] = result
+
+        # if len(agent_lps) == 0:
+        #     return None
+
+        # lps = cast(xr.Dataset, agent_concatenation(agent_lps, dim="agent"))
+        # coords = {"agent", "technology", "region"}.intersection(assets.asset.coords)
+        # constraints = self.constraints(
+        #     demand=demands,
+        #     assets=reduce_assets(assets, coords=coords).set_coords(coords),
+        #     search_space=lps.search_space,
+        #     market=market,
+        #     technologies=technologies,
+        #     year=current_year,
+        # )
+        # return lps, constraints
 
     @classmethod
     def factory(

From b47c811957c92548257c107108229a969c8fd477 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Thu, 3 Oct 2024 14:15:51 +0100
Subject: [PATCH 03/18] Simplify retirment profile code

---
 src/muse/agents/agent.py      | 166 ++++++++++++++++++----------------
 src/muse/investments.py       |  35 +++----
 src/muse/sectors/subsector.py |  36 +-------
 tests/test_investments.py     |   2 +-
 4 files changed, 103 insertions(+), 136 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index a6d45e30..8efdd607 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -80,14 +80,9 @@ def next(
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        time_period: int = 1,
+        time_period: int,
     ):
-        """Iterates agent one turn.
-
-        The goal is to figure out from market variables which technologies to invest in
-        and by how much.
-        """
-        pass
+        """Increments agent to the next time point (e.g. performing investments)."""
 
     def __repr__(self):
         return (
@@ -124,7 +119,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,9 +161,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."""
@@ -245,6 +238,15 @@ def asset_housekeeping(self):
         #  state.
         self.assets = self._housekeeping(self, self.assets)
 
+    def next(
+        self,
+        technologies: xr.Dataset,
+        market: xr.Dataset,
+        demand: xr.DataArray,
+        time_period: int,
+    ):
+        self.year += time_period
+
     def compute_decision(
         self,
         technologies: xr.Dataset,
@@ -279,73 +281,6 @@ def compute_decision(
         decision = self.decision(objectives)
         return decision
 
-    def add_investments(
-        self,
-        technologies: xr.Dataset,
-        investments: xr.DataArray,
-        current_year: int,
-        time_period: int,
-    ):
-        """Add new assets to the agent."""
-        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))
-
-        self.assets = self.merge_transform(self.assets, new_assets)
-
-    def retirement_profile(
-        self,
-        technologies: xr.Dataset,
-        investments: xr.DataArray,
-        current_year: int,
-        time_period: int,
-    ) -> Optional[xr.DataArray]:
-        from muse.investments import cliff_retirement_profile
-
-        if "asset" in investments.dims:
-            investments = investments.sum("asset")
-        if "agent" in investments.dims:
-            investments = investments.squeeze("agent", drop=True)
-        investments = investments.sel(
-            replacement=(investments > self.asset_threshold).any(
-                [d for d in investments.dims if d != "replacement"]
-            )
-        )
-        if investments.size == 0:
-            return None
-
-        # figures out the retirement profile for the new investments
-        lifetime = self.filter_input(
-            technologies.technical_life,
-            year=current_year,
-            technology=investments.replacement,
-        )
-        profile = cliff_retirement_profile(
-            lifetime.clip(min=time_period),
-            current_year=current_year + time_period,
-            protected=max(self.forecast - time_period - 1, 0),
-        )
-        if "dst_region" in investments.coords:
-            investments = investments.reindex_like(profile, method="ffill")
-
-        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.
-        # we try and keep only those that were there originally.
-        if set(new_assets.dims) != set(self.assets.dims):
-            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."""
@@ -357,7 +292,7 @@ def __init__(
         investment: Optional[Callable] = None,
         **kwargs,
     ):
-        """Creates a standard buildings agent.
+        """Creates an investing agent.
 
         Arguments:
             *args: See :py:class:`~muse.agents.agent.Agent`
@@ -384,7 +319,7 @@ def next(
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        time_period: int = 1,
+        time_period: int,
     ):
         """Iterates agent one turn.
 
@@ -456,4 +391,75 @@ def next(
             time_period=time_period,
         )
 
+        # Increment the year
         self.year += time_period
+
+    def add_investments(
+        self,
+        technologies: xr.Dataset,
+        investments: xr.DataArray,
+        current_year: int,
+        time_period: int,
+    ):
+        """Add new assets to the agent."""
+        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))
+
+        self.assets = self.merge_transform(self.assets, new_assets)
+
+    def retirement_profile(
+        self,
+        technologies: xr.Dataset,
+        investments: xr.DataArray,
+        current_year: int,
+        time_period: int,
+    ) -> 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"]
+            )
+        )
+        if investments.size == 0:
+            return None
+
+        # 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,
+            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.
+        # we try and keep only those that were there originally.
+        if set(new_assets.dims) != set(self.assets.dims):
+            new, old = new_assets.dims, self.assets.dims
+            raise RuntimeError(f"Asset dimensions do not match: {new} vs {old}")
+        return new_assets
diff --git a/src/muse/investments.py b/src/muse/investments.py
index 7d92e5c7..08a000b1 100644
--- a/src/muse/investments.py
+++ b/src/muse/investments.py
@@ -175,8 +175,7 @@ def compute_investment(
 
 def cliff_retirement_profile(
     technical_life: xr.DataArray,
-    current_year: int = 0,
-    protected: int = 0,
+    investment_year: int,
     interpolation: str = "linear",
     **kwargs,
 ) -> xr.DataArray:
@@ -186,19 +185,13 @@ def cliff_retirement_profile(
     Assets with a technical life smaller than the input time-period should automatically
     be renewed.
 
-    Hence, if ``technical_life <= protected``, then effectively, the technical life is
-    rewritten as ``technical_life * n`` with ``n = int(protected // technical_life) +
-    1``.
-
     We could just return an array where each year is represented. Instead, to save
     memory, we return a compact view of the same where years where no change happens are
     removed.
 
     Arguments:
         technical_life: lifetimes for each technology
-        current_year: current year
-        protected: The technologies are assumed to be renewed between years
-            `current_year` and `current_year + protected`
+        investment_year: The year in which the investment is made
         interpolation: Interpolation type
         **kwargs: arguments by which to filter technical_life, if any.
 
@@ -211,26 +204,26 @@ def cliff_retirement_profile(
     if kwargs:
         technical_life = technical_life.sel(**kwargs)
     if "year" in technical_life.dims:
-        technical_life = technical_life.interp(year=current_year, method=interpolation)
-    technical_life = (1 + protected // technical_life) * technical_life  # type:ignore
+        technical_life = technical_life.interp(
+            year=investment_year, method=interpolation
+        )
 
+    # Create profile across all years
     if len(technical_life) > 0:
-        max_year = int(current_year + technical_life.max())
+        max_year = int(investment_year + technical_life.max())
     else:
-        max_year = int(current_year + protected)
+        max_year = investment_year
     allyears = xr.DataArray(
-        range(current_year, max_year + 1),
+        range(investment_year, max_year + 1),
         dims="year",
-        coords={"year": range(current_year, max_year + 1)},
+        coords={"year": range(investment_year, max_year + 1)},
     )
+    profile = allyears < (investment_year + technical_life)  # type: ignore
 
-    profile = allyears < (current_year + technical_life)  # type: ignore
-
-    # now we minimize the number of years needed to represent the profile fully
-    # this is done by removing the central year of any three repeating year, ensuring
-    # the removed year can be recovered by a linear interpolation.
+    # 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
+    # the removed year can be recovered by linear interpolation.
     goodyears = avoid_repetitions(profile.astype(int))
-
     return profile.sel(year=goodyears).astype(bool)
 
 
diff --git a/src/muse/sectors/subsector.py b/src/muse/sectors/subsector.py
index d3704517..3beb7efa 100644
--- a/src/muse/sectors/subsector.py
+++ b/src/muse/sectors/subsector.py
@@ -64,21 +64,8 @@ def invest(
         for agent in self.agents:
             agent.asset_housekeeping()
 
-        # Perform the investment
+        # Perform the investments
         self.aggregate_lp(technologies, market, time_period, current_year=current_year)
-        # if lp_problem is None:
-        #     return
-
-        # # If there is a problem with the LP...
-        # years = technologies.year
-        # techs = technologies.interp(year=years)
-        # techs = techs.sel(year=current_year + time_period)
-
-        # solution = self.investment(
-        #     search=lp_problem[0], technologies=techs, constraints=lp_problem[1]
-        # )
-
-        # self.assign_back_to_agents(technologies, solution, current_year, time_period)
 
     def assign_back_to_agents(
         self,
@@ -133,32 +120,13 @@ def aggregate_lp(
             .swap_dims(dict(asset="technology"))
         )
 
-        # agent_lps: MutableMapping[Hashable, xr.Dataset] = {}
+        # Increment each agent (perform investments)
         for agent in self.agents:
             if "agent" in demands.coords:
                 share = demands.sel(asset=demands.agent == agent.uuid)
             else:
                 share = demands
-
-            # Compute investments for the agent
             agent.next(technologies, agent_market, share, time_period=time_period)
-        #     if result is not None:
-        #         agent_lps[agent.uuid] = result
-
-        # if len(agent_lps) == 0:
-        #     return None
-
-        # lps = cast(xr.Dataset, agent_concatenation(agent_lps, dim="agent"))
-        # coords = {"agent", "technology", "region"}.intersection(assets.asset.coords)
-        # constraints = self.constraints(
-        #     demand=demands,
-        #     assets=reduce_assets(assets, coords=coords).set_coords(coords),
-        #     search_space=lps.search_space,
-        #     market=market,
-        #     technologies=technologies,
-        #     year=current_year,
-        # )
-        # return lps, constraints
 
     @classmethod
     def factory(
diff --git a/tests/test_investments.py b/tests/test_investments.py
index 92b9abbb..415f0c9f 100644
--- a/tests/test_investments.py
+++ b/tests/test_investments.py
@@ -75,7 +75,7 @@ def test_cliff_retirement_random_profile(protected):
 
     current = 5
     profile = cliff_retirement_profile(
-        lifetime, current_year=current, protected=protected
+        lifetime, investment_year=current, protected=protected
     )
     assert profile.year.min() == current
     assert profile.year.max() <= current + effective_lifetime.max() + 1

From a86e07f7f5f3cc93e9868587e5be4a2bd247ecfa Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Thu, 3 Oct 2024 14:56:04 +0100
Subject: [PATCH 04/18] Simplify merge_assets

---
 src/muse/agents/agent.py |  3 ++-
 src/muse/utilities.py    | 50 +++++++++++++++-------------------------
 2 files changed, 21 insertions(+), 32 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index 8efdd607..6a014e34 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -402,10 +402,10 @@ def add_investments(
         time_period: int,
     ):
         """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(
@@ -413,6 +413,7 @@ def add_investments(
         )
         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(
diff --git a/src/muse/utilities.py b/src/muse/utilities.py
index a2afd009..6a8e4e09 100644
--- a/src/muse/utilities.py
+++ b/src/muse/utilities.py
@@ -385,41 +385,29 @@ def merge_assets(
     dimension: str = "asset",
 ) -> xr.DataArray:
     """Merge two capacity arrays."""
+    # Interpolate capacity arrays to a common time framework
     years = sorted(set(capa_a.year.values).union(capa_b.year.values))
-
     if len(capa_a.year) == 1:
-        result = xr.concat(
-            (
-                capa_a,
-                capa_b.interp(year=years, method=interpolation).fillna(0),
-            ),
-            dim=dimension,
-        ).fillna(0)
+        capa_a_interp = capa_a
+        capa_b_interp = capa_b.interp(year=years, method=interpolation).fillna(0)
     elif len(capa_b.year) == 1:
-        result = xr.concat(
-            (
-                capa_a.interp(year=years, method=interpolation).fillna(0),
-                capa_b,
-            ),
-            dim=dimension,
-        ).fillna(0)
+        capa_a_interp = capa_a.interp(year=years, method=interpolation).fillna(0)
+        capa_b_interp = capa_b
     else:
-        result = xr.concat(
-            (
-                capa_a.interp(year=years, method=interpolation).fillna(0),
-                capa_b.interp(year=years, method=interpolation).fillna(0),
-            ),
-            dim=dimension,
-        )
-    forgroup = result.pipe(coords_to_multiindex, dimension=dimension)
-    if len(forgroup[dimension]) != len(set(forgroup[dimension].values)):
-        result = (
-            forgroup.groupby(dimension)
-            .sum(dimension)
-            .clip(min=0)
-            .pipe(multiindex_to_coords, dimension=dimension)
-        )
-    return result
+        capa_a_interp = capa_a.interp(year=years, method=interpolation).fillna(0)
+        capa_b_interp = capa_b.interp(year=years, method=interpolation).fillna(0)
+
+    # Concatenate the two capacity arrays
+    result = xr.concat((capa_a_interp, capa_b_interp), dim=dimension)
+
+    # forgroup = result.pipe(coords_to_multiindex, dimension=dimension)
+    # result = (
+    #     forgroup.groupby(dimension)
+    #     .sum(dimension)
+    #     .clip(min=0)
+    #     .pipe(multiindex_to_coords, dimension=dimension)
+    # )
+    return result.clip(min=0)
 
 
 def avoid_repetitions(data: xr.DataArray, dim: str = "year") -> xr.DataArray:

From 3e513111cbc1c15ac2acefc0bd442034300d7735 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 10:06:02 +0100
Subject: [PATCH 05/18] Revert change to merge_assets

---
 src/muse/utilities.py | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/muse/utilities.py b/src/muse/utilities.py
index 6a8e4e09..120949ee 100644
--- a/src/muse/utilities.py
+++ b/src/muse/utilities.py
@@ -141,6 +141,8 @@ def reduce_assets(
     """
     from copy import copy
 
+    assets = copy(assets)
+
     if operation is None:
 
         def operation(x):
@@ -148,23 +150,31 @@ def operation(x):
 
     assert operation is not None
 
+    # Concatenate assets if a sequence is given
     if not isinstance(assets, (xr.Dataset, xr.DataArray)):
         assets = xr.concat(assets, dim=dim)
     assert isinstance(assets, (xr.Dataset, xr.DataArray))
+
+    # If there are no assets, nothing needs to be done
     if assets[dim].size == 0:
         return assets
+
+    # Coordinates to reduce over (e.g. technology, installed)
     if coords is None:
         coords = [cast(str, k) for k, v in assets.coords.items() if v.dims == (dim,)]
     elif isinstance(coords, str):
         coords = (coords,)
     coords = [k for k in coords if k in assets.coords and assets[k].dims == (dim,)]
-    assets = copy(assets)
+
+    # Create a new dimension to group by
     dtypes = [(d, assets[d].dtype) for d in coords]
     grouper = np.array(
         list(zip(*(cast(Iterator, assets[d].values) for d in coords))), dtype=dtypes
     )
     assert "grouper" not in assets.coords
     assets["grouper"] = "asset", grouper
+
+    # Perform the operation
     result = operation(assets.groupby("grouper")).rename(grouper=dim)
     for i, d in enumerate(coords):
         result[d] = dim, [u[i] for u in result[dim].values]
@@ -400,14 +410,16 @@ def merge_assets(
     # Concatenate the two capacity arrays
     result = xr.concat((capa_a_interp, capa_b_interp), dim=dimension)
 
-    # forgroup = result.pipe(coords_to_multiindex, dimension=dimension)
-    # result = (
-    #     forgroup.groupby(dimension)
-    #     .sum(dimension)
-    #     .clip(min=0)
-    #     .pipe(multiindex_to_coords, dimension=dimension)
-    # )
-    return result.clip(min=0)
+    #
+    forgroup = result.pipe(coords_to_multiindex, dimension=dimension)
+    if len(forgroup[dimension]) != len(set(forgroup[dimension].values)):
+        result = (
+            forgroup.groupby(dimension)
+            .sum(dimension)
+            .clip(min=0)
+            .pipe(multiindex_to_coords, dimension=dimension)
+        )
+    return result
 
 
 def avoid_repetitions(data: xr.DataArray, dim: str = "year") -> xr.DataArray:

From 941a5a6824a162411bd1cdf0cc055649cf4fc3aa Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 10:07:36 +0100
Subject: [PATCH 06/18] Delete unused factory

---
 src/muse/agents/__init__.py  |  2 +-
 src/muse/agents/factories.py | 91 ------------------------------------
 2 files changed, 1 insertion(+), 92 deletions(-)

diff --git a/src/muse/agents/__init__.py b/src/muse/agents/__init__.py
index 069f1c2e..926d1ee1 100644
--- a/src/muse/agents/__init__.py
+++ b/src/muse/agents/__init__.py
@@ -8,4 +8,4 @@
 ]
 
 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
diff --git a/src/muse/agents/factories.py b/src/muse/agents/factories.py
index efd7056f..236eb24e 100644
--- a/src/muse/agents/factories.py
+++ b/src/muse/agents/factories.py
@@ -7,7 +7,6 @@
 import xarray as xr
 
 from muse.agents.agent import Agent, InvestingAgent
-from muse.defaults import DEFAULT_SECTORS_DIRECTORY
 from muse.errors import RetrofitAgentNotDefined, TechnologyNotDefined
 
 
@@ -173,96 +172,6 @@ def create_agent(agent_type: str, **kwargs) -> Agent:
     return method(**kwargs)  # type: ignore
 
 
-def factory(
-    existing_capacity_path: Optional[Union[Path, str]] = None,
-    agent_parameters_path: Optional[Union[Path, str]] = None,
-    technodata_path: Optional[Union[Path, str]] = None,
-    technodata_timeslices_path: Optional[Union[str, Path]] = None,
-    sector: Optional[str] = None,
-    sectors_directory: Union[str, Path] = DEFAULT_SECTORS_DIRECTORY,
-    baseyear: int = 2010,
-) -> list[Agent]:
-    """Reads list of agents from standard MUSE input files."""
-    from copy import deepcopy
-    from logging import getLogger
-    from textwrap import dedent
-
-    from muse.readers import (
-        read_csv_agent_parameters,
-        read_initial_assets,
-        read_technodata_timeslices,
-        read_technodictionary,
-    )
-    from muse.readers.csv import find_sectors_file
-
-    if sector is None:
-        assert existing_capacity_path is not None
-        assert agent_parameters_path is not None
-        assert technodata_path is not None
-
-    if existing_capacity_path is None:
-        existing_capacity_path = find_sectors_file(
-            f"Existing{sector}.csv", sector, sectors_directory
-        )
-    if agent_parameters_path is None:
-        agent_parameters_path = find_sectors_file(
-            f"BuildingAgent{sector}.csv", sector, sectors_directory
-        )
-    if technodata_path is None:
-        technodata_path = find_sectors_file(
-            f"technodata{sector}.csv", sector, sectors_directory
-        )
-
-    params = read_csv_agent_parameters(agent_parameters_path)
-    techno = read_technodictionary(technodata_path)
-    capa = read_initial_assets(existing_capacity_path)
-    if technodata_timeslices_path and isinstance(
-        technodata_timeslices_path, (str, Path)
-    ):
-        technodata_timeslices = read_technodata_timeslices(technodata_timeslices_path)
-    else:
-        technodata_timeslices = None
-    result = []
-    for param in params:
-        if param["agent_type"] == "retrofit":
-            param["technologies"] = techno.sel(region=param["region"])
-        if technodata_timeslices is not None:
-            param.drop_vars("utilization_factor")
-            param = param.merge(technodata_timeslices.sel(region=param["region"]))
-        param["category"] = param["agent_type"]
-        param["capacity"] = deepcopy(capa.sel(region=param["region"]))
-        param["year"] = baseyear
-        result.append(create_agent(**param))
-
-    nregs = len({u.region for u in result})
-    types = [u.name for u in result]
-    msg = dedent(
-        """\
-        Read agents for sector {name} from:
-            - agent parameter file {para}
-            - technologies data file {tech}
-            - initial capacity file {ini}
-
-        Found {n} agents across {nregs} regions{end}
-        """.format(
-            n=len(result),
-            name=sector,
-            para=agent_parameters_path,
-            tech=technodata_path,
-            ini=existing_capacity_path,
-            nregs=nregs,
-            end="." if len(result) == 0 else ", with:\n",
-        )
-    )
-    for t in set(types):
-        n = types.count(t)
-        msg += "    - {n} {t} agent{plural}\n".format(
-            n=n, t=t, plural="" if n == 1 else "s"
-        )
-    getLogger(__name__).info(msg)
-    return result
-
-
 def agents_factory(
     params_or_path: Union[str, Path, list],
     capacity: Union[xr.DataArray, str, Path],

From 9d2cac9c4be1729193e2483373ff5a35f83ce974 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 10:08:01 +0100
Subject: [PATCH 07/18] More comments added to code

---
 src/muse/demand_share.py | 93 +++++++++++++++++++++++++---------------
 src/muse/investments.py  | 14 +++++-
 2 files changed, 70 insertions(+), 37 deletions(-)

diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index f95e43f3..f913c5fb 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -363,8 +363,15 @@ def decommissioning(capacity):
     if current_year is None:
         current_year = market.year.min()
 
+    # Make sure there are no retrofit agents
+    for agent in agents:
+        if agent.category == "retrofit":
+            raise RetrofitAgentInStandardDemandShare()
+
+    # Calculate existing capacity
     capacity = reduce_assets([agent.assets.capacity for agent in agents])
 
+    # Calculate new and retrofit demands
     demands = new_and_retro_demands(
         capacity,
         market,
@@ -374,34 +381,32 @@ def decommissioning(capacity):
         forecast=forecast,
     )
 
+    # Only consider end-use commodities
     demands = demands.where(
         is_enduse(technologies.comm_usage.sel(commodity=demands.commodity)), 0
     )
 
-    for agent in agents:
-        if agent.category == "retrofit":
-            raise RetrofitAgentInStandardDemandShare()
-
     id_to_share: MutableMapping[Hashable, xr.DataArray] = {}
     for region in demands.region.values:
+        # Calculate current capacity
         current_capacity: MutableMapping[Hashable, xr.DataArray] = {
             agent.uuid: agent.assets.capacity
             for agent in agents
             if agent.region == region
         }
+
+        # Split demands between agents
         id_to_quantity = {
             agent.uuid: (agent.name, agent.region, agent.quantity)
             for agent in agents
             if agent.region == region
         }
-
         retro_demands: MutableMapping[Hashable, xr.DataArray] = _inner_split(
             current_capacity,
             demands.retrofit.sel(region=region),
             decommissioning,
             id_to_quantity,
         )
-
         new_demands = _inner_split(
             current_capacity,
             demands.new.sel(region=region),
@@ -413,6 +418,7 @@ def decommissioning(capacity):
             id_to_quantity,
         )
 
+        # Sum new and retrofit demands
         total_demands = {
             k: new_demands[k] + retro_demands[k] for k in new_demands.keys()
         }
@@ -527,14 +533,21 @@ def unmet_demand(
 
     prod_method = production if callable(production) else prod_factory(production)
     assert callable(prod_method)
+
+    # Calculate production by existing assets
     produced = prod_method(market=market, capacity=capacity, technologies=technologies)
+
+    # Total commodity production by summing over assets
     if "dst_region" in produced.dims:
         produced = produced.sum("asset").rename(dst_region="region")
     elif "region" in produced.coords and produced.region.dims:
         produced = produced.groupby("region").sum("asset")
     else:
         produced = produced.sum("asset")
-    return (market.consumption - produced).clip(min=0)
+
+    # Unmet demand is the difference between the consumption and the production
+    unmet_demand = (market.consumption - produced).clip(min=0)
+    return unmet_demand
 
 
 def new_consumption(
@@ -565,20 +578,23 @@ def new_consumption(
     if current_year is None:
         current_year = market.year.min()
 
-    ts_capa = convert_timeslice(
-        capacity.interp(year=current_year), market.timeslice, QuantityType.EXTENSIVE
-    )
+    # Interpolate market to forecast year
+    market = market.interp(year=[current_year, current_year + forecast])
+    current = market.sel(year=current_year, drop=True)
+    forecasted = market.sel(year=current_year + forecast, drop=True)
+
+    # Calculate the increase in consumption over the forecast period
+    delta = (forecasted.consumption - current.consumption).clip(min=0)
+
+    # Capacity in the forecast year
     ts_capa = convert_timeslice(
         capacity.interp(year=current_year + forecast),
         market.timeslice,
         QuantityType.EXTENSIVE,
     )
     assert isinstance(ts_capa, xr.DataArray)
-    market = market.interp(year=[current_year, current_year + forecast])
-    current = market.sel(year=current_year, drop=True)
-    forecasted = market.sel(year=current_year + forecast, drop=True)
 
-    delta = (forecasted.consumption - current.consumption).clip(min=0)
+    #
     missing = unmet_demand(current, ts_capa, technologies)
     consumption = minimum(delta, missing)
     return consumption
@@ -612,23 +628,28 @@ def new_and_retro_demands(
     if current_year is None:
         current_year = market.year.min()
 
+    # Interpolate market to forecast year
     smarket: xr.Dataset = market.interp(year=[current_year, current_year + forecast])
+
+    # Split capacity between timeslices
     ts_capa = convert_timeslice(
         capacity.interp(year=[current_year, current_year + forecast]),
         market.timeslice,
         QuantityType.EXTENSIVE,
     )
-
     assert isinstance(ts_capa, xr.DataArray)
+
     if hasattr(ts_capa, "region") and ts_capa.region.dims == ():
         ts_capa["region"] = "asset", [str(ts_capa.region.values)] * len(ts_capa.asset)
 
+    # Calculate demand to allocate to "new" agents
     new_demand = new_consumption(
         ts_capa, smarket, technologies, current_year=current_year, forecast=forecast
     )
     if "year" in new_demand.dims:
         new_demand = new_demand.squeeze("year")
 
+    # Total production in the forecast year by existing assets
     service = (
         production_method(
             smarket.sel(year=current_year + forecast),
@@ -638,37 +659,39 @@ def new_and_retro_demands(
         .groupby("region")
         .sum("asset")
     )
-    # existing asset should not execute beyond demand
+
+    # Existing asset should not execute beyond demand
     service = minimum(
         service, smarket.consumption.sel(year=current_year + forecast, drop=True)
     )
+
+    # Leftover demand that cannot be serviced by existing assets or "new" agents
     retro_demand = (
         smarket.consumption.sel(year=current_year + forecast, drop=True)
         - new_demand
         - service
     ).clip(min=0)
-
     if "year" in retro_demand.dims:
         retro_demand = retro_demand.squeeze("year")
 
     return xr.Dataset({"new": new_demand, "retrofit": retro_demand})
 
 
-def new_demand(
-    capacity: xr.DataArray,
-    market: xr.Dataset,
-    technologies: xr.Dataset,
-    production: Union[str, Mapping, Callable] = "maximum_production",
-    current_year: Optional[int] = None,
-    forecast: int = 5,
-) -> xr.DataArray:
-    """Calculates the new demand that needs to be covered.
-
-    It groups the demand related to an increase in consumption as well as the existing
-    demand associated with decommissoned assets. Internally, it just calls
-    `new_and_retro` demands and adds together both components.
-    """
-    demand = new_and_retro_demands(
-        capacity, market, technologies, production, current_year, forecast
-    )
-    return (demand["new"] + demand["retrofit"]).rename("demand")
+# def new_demand(
+#     capacity: xr.DataArray,
+#     market: xr.Dataset,
+#     technologies: xr.Dataset,
+#     production: Union[str, Mapping, Callable] = "maximum_production",
+#     current_year: Optional[int] = None,
+#     forecast: int = 5,
+# ) -> xr.DataArray:
+#     """Calculates the new demand that needs to be covered.
+
+#     It groups the demand related to an increase in consumption as well as the existing
+#     demand associated with decommissoned assets. Internally, it just calls
+#     `new_and_retro` demands and adds together both components.
+#     """
+#     demand = new_and_retro_demands(
+#         capacity, market, technologies, production, current_year, forecast
+#     )
+#     return (demand["new"] + demand["retrofit"]).rename("demand")
diff --git a/src/muse/investments.py b/src/muse/investments.py
index 08a000b1..8d8cc707 100644
--- a/src/muse/investments.py
+++ b/src/muse/investments.py
@@ -154,6 +154,7 @@ def compute_investment(
         """
         from numpy import zeros
 
+        # Skip the investment step if no assets or replacements are available
         if any(u == 0 for u in search.decision.shape):
             return xr.DataArray(
                 zeros((len(search.asset), len(search.replacement))),
@@ -161,6 +162,7 @@ def compute_investment(
                 dims=("asset", "replacement"),
             )
 
+        # Otherwise, compute the investment
         return investment(
             search.decision,
             search.search_space,
@@ -305,18 +307,24 @@ def scipy_match_demand(
 
     if "timeslice" in costs.dims and timeslice_op is not None:
         costs = timeslice_op(costs)
+
+    timeslice = next(cs.timeslice for cs in constraints if "timeslice" in cs.dims)
+
+    # Select technodata for the current year
     if "year" in technologies.dims and year is None:
         raise ValueError("Missing year argument")
     elif "year" in technologies.dims:
         techs = technologies.sel(year=year).drop_vars("year")
     else:
         techs = technologies
-    timeslice = next(cs.timeslice for cs in constraints if "timeslice" in cs.dims)
 
+    # Run scipy optimization with highs solver
     adapter = ScipyAdapter.factory(
         techs, cast(np.ndarray, costs), timeslice, *constraints
     )
     res = linprog(**adapter.kwargs, method="highs")
+
+    # Backup: try with highs-ipm
     if not res.success and (res.status != 0):
         res = linprog(
             **adapter.kwargs,
@@ -338,7 +346,9 @@ def scipy_match_demand(
             getLogger(__name__).critical(msg)
             raise GrowthOfCapacityTooConstrained
 
-    return cast(Callable[[np.ndarray], xr.Dataset], adapter.to_muse)(res.x)
+    # Convert results to a MUSE friendly format
+    result = cast(Callable[[np.ndarray], xr.Dataset], adapter.to_muse)(res.x)
+    return result
 
 
 @register_investment(name=["cvxopt"])

From e386da12460a0c11a249ec0d3a3571d0af8544d0 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 12:10:10 +0100
Subject: [PATCH 08/18] Revert some changes to fix tests

---
 src/muse/demand_share.py   |  9 +++++----
 src/muse/investments.py    |  2 +-
 src/muse/quantities.py     |  8 +++++---
 src/muse/sectors/sector.py | 19 ++++++++++++++++++-
 tests/test_agents.py       |  2 +-
 tests/test_investments.py  | 14 +++++++-------
 tests/test_subsector.py    | 11 +++++++++--
 7 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index f913c5fb..01af3214 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -473,6 +473,7 @@ def _inner_split(
     """
     from numpy import logical_and
 
+    # Find decrease in capacity production by each asset over time
     shares: Mapping[Hashable, xr.DataArray] = {
         key: method(capacity=capacity)
         .groupby("technology")
@@ -480,24 +481,24 @@ def _inner_split(
         .rename(technology="asset")
         for key, capacity in assets.items()
     }
+
+    # Total decrease in production across assets
     try:
         summed_shares: xr.DataArray = xr.concat(shares.values(), dim="concat_dim").sum(
             "concat_dim"
         )
-
-        # Calculates the total demand assigned in the previous step with the "method"
-        # function across agents and assets.
         total: xr.DataArray = summed_shares.sum("asset")
     except AttributeError:
         raise AgentWithNoAssetsInDemandShare()
 
     # Calculates the demand divided by the number of assets times the number of agents
     # if the demand is bigger than zero and the total demand assigned with the "method"
-    # function is zero.
+    # function is zero (i.e. no decrease in production).
     unassigned = (demand / (len(shares) * len(summed_shares))).where(
         logical_and(demand > 1e-12, total <= 1e-12), 0
     )
 
+    # ???
     totals = {
         key: (share / share.sum("asset")).fillna(0) for key, share in shares.items()
     }
diff --git a/src/muse/investments.py b/src/muse/investments.py
index 8d8cc707..19598f49 100644
--- a/src/muse/investments.py
+++ b/src/muse/investments.py
@@ -177,7 +177,7 @@ def compute_investment(
 
 def cliff_retirement_profile(
     technical_life: xr.DataArray,
-    investment_year: int,
+    investment_year: int = 0,
     interpolation: str = "linear",
     **kwargs,
 ) -> xr.DataArray:
diff --git a/src/muse/quantities.py b/src/muse/quantities.py
index 7daf87b1..66df60c3 100644
--- a/src/muse/quantities.py
+++ b/src/muse/quantities.py
@@ -250,9 +250,11 @@ def decommissioning_demand(
     baseyear = min(year)
     dyears = [u for u in year if u != baseyear]
 
-    return maximum_production(
-        technologies, capacity.sel(year=baseyear) - capacity.sel(year=dyears)
-    ).clip(min=0)
+    # Calculate the decrease in capacity from the current year to future years
+    capacity_decrease = capacity.sel(year=baseyear) - capacity.sel(year=dyears)
+
+    # Calculate production associated with this capacity
+    return maximum_production(technologies, capacity_decrease).clip(min=0)
 
 
 def consumption(
diff --git a/src/muse/sectors/sector.py b/src/muse/sectors/sector.py
index 1dc42dd5..31e04dae 100644
--- a/src/muse/sectors/sector.py
+++ b/src/muse/sectors/sector.py
@@ -26,6 +26,7 @@ def factory(cls, name: str, settings: Any) -> Sector:
         from muse.interactions import factory as interaction_factory
         from muse.outputs.sector import factory as ofactory
         from muse.production import factory as pfactory
+        from muse.readers import read_timeslices
         from muse.readers.toml import read_technodata
         from muse.utilities import nametuple_to_dict
 
@@ -38,6 +39,11 @@ def factory(cls, name: str, settings: Any) -> Sector:
         if len(sector_settings["subsectors"]._asdict()) == 0:
             raise RuntimeError(f"Empty 'subsectors' section in sector {name}")
 
+        # Timeslices
+        timeslices = read_timeslices(
+            sector_settings.pop("timeslice_levels", None)
+        ).get_index("timeslice")
+
         # Read technologies
         technologies = read_technodata(settings, name, settings.time_framework)
 
@@ -79,12 +85,18 @@ def factory(cls, name: str, settings: Any) -> Sector:
         interactions = interaction_factory(sector_settings.pop("interactions", None))
 
         # Create sector
-        for attr in ("technodata", "commodities_out", "commodities_in"):
+        for attr in (
+            "technodata",
+            "commodities_out",
+            "commodities_in",
+            "technodata_timeslices",
+        ):
             sector_settings.pop(attr, None)
         return cls(
             name,
             technologies,
             subsectors=subsectors,
+            timeslices=timeslices,
             supply_prod=supply,
             outputs=outputs,
             interactions=interactions,
@@ -96,6 +108,7 @@ def __init__(
         name: str,
         technologies: xr.Dataset,
         subsectors: Sequence[Subsector] = [],
+        timeslices: pd.MultiIndex | None = None,
         interactions: Callable[[Sequence[AbstractAgent]], None] | None = None,
         interpolation: str = "linear",
         outputs: Callable | None = None,
@@ -111,6 +124,10 @@ def __init__(
         """Subsectors controlled by this object."""
         self.technologies: xr.Dataset = technologies
         """Parameters describing the sector's technologies."""
+        self.timeslices: pd.MultiIndex | None = timeslices
+        """Timeslice at which this sector operates.
+        If None, it will operate using the timeslice of the input market.
+        """
         self.interpolation: Mapping[str, Any] = {
             "method": interpolation,
             "kwargs": {"fill_value": "extrapolate"},
diff --git a/tests/test_agents.py b/tests/test_agents.py
index c767c2c9..3287e0e3 100644
--- a/tests/test_agents.py
+++ b/tests/test_agents.py
@@ -143,7 +143,7 @@ def test_run_retro_agent(retro_agent, technologies, agent_market, demand_share):
     technologies.max_capacity_addition[:] = retro_agent.assets.capacity.sum() * 100
     technologies.max_capacity_growth[:] = retro_agent.assets.capacity.sum() * 100
 
-    retro_agent.next(technologies, agent_market, demand_share)
+    retro_agent.next(technologies, agent_market, demand_share, time_period=5)
 
 
 def test_merge_assets(assets):
diff --git a/tests/test_investments.py b/tests/test_investments.py
index 415f0c9f..f330eda4 100644
--- a/tests/test_investments.py
+++ b/tests/test_investments.py
@@ -44,7 +44,7 @@ def test_cliff_retirement_known_profile():
         name="technical_life",
     )
 
-    profile = cliff_retirement_profile(lifetime)
+    profile = cliff_retirement_profile(technical_life=lifetime)
     expected = array(
         [
             [True, False, False, False],
@@ -73,12 +73,12 @@ def test_cliff_retirement_random_profile(protected):
     )
     effective_lifetime = (protected // lifetime + 1) * lifetime
 
-    current = 5
+    investment_year = 5
     profile = cliff_retirement_profile(
-        lifetime, investment_year=current, protected=protected
+        technical_life=lifetime, investment_year=investment_year, protected=protected
     )
-    assert profile.year.min() == current
-    assert profile.year.max() <= current + effective_lifetime.max() + 1
-    assert profile.astype(int).interp(year=current).all()
-    assert profile.astype(int).interp(year=current + protected).all()
+    assert profile.year.min() == investment_year
+    assert profile.year.max() <= investment_year + effective_lifetime.max() + 1
+    assert profile.astype(int).interp(year=investment_year).all()
+    assert profile.astype(int).interp(year=investment_year + protected).all()
     assert not profile.astype(int).interp(year=profile.year.max()).any()
diff --git a/tests/test_subsector.py b/tests/test_subsector.py
index 9c326f1f..3ec688c2 100644
--- a/tests/test_subsector.py
+++ b/tests/test_subsector.py
@@ -48,7 +48,12 @@ def test_subsector_investing_aggregation():
             subsector = Subsector(agents, commodities)
             initial_agents = deepcopy(agents)
             assert {agent.year for agent in agents} == {int(market.year.min())}
-            assert subsector.aggregate_lp(technologies, market) is None
+            assert (
+                subsector.aggregate_lp(
+                    technologies, market, time_period=5, current_year=5
+                )
+                is None
+            )
             assert {agent.year for agent in agents} == {int(market.year.min() + 5)}
             for initial, final in zip(initial_agents, agents):
                 assert initial.assets.sum() != final.assets.sum()
@@ -105,7 +110,9 @@ def test_subsector_noninvesting_aggregation(market, model, technologies, tmp_pat
         commodity=technologies.commodity, region=technologies.region
     ).interp(year=[2020, 2025])
     assert all(agent.year == 2020 for agent in agents)
-    result = subsector.aggregate_lp(technologies, market)
+    result = subsector.aggregate_lp(
+        technologies, market, time_period=5, current_year=2020
+    )
 
     assert result is not None
     assert len(result) == 2

From 018ca5c3bb8c32dc438a3edf4802de4a057165e5 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 12:32:59 +0100
Subject: [PATCH 09/18] Fix tests

---
 tests/test_investments.py |  4 ++--
 tests/test_subsector.py   | 18 +-----------------
 2 files changed, 3 insertions(+), 19 deletions(-)

diff --git a/tests/test_investments.py b/tests/test_investments.py
index f330eda4..a4fc8769 100644
--- a/tests/test_investments.py
+++ b/tests/test_investments.py
@@ -75,10 +75,10 @@ def test_cliff_retirement_random_profile(protected):
 
     investment_year = 5
     profile = cliff_retirement_profile(
-        technical_life=lifetime, investment_year=investment_year, protected=protected
+        technical_life=lifetime.clip(min=protected), investment_year=investment_year
     )
     assert profile.year.min() == investment_year
     assert profile.year.max() <= investment_year + effective_lifetime.max() + 1
     assert profile.astype(int).interp(year=investment_year).all()
-    assert profile.astype(int).interp(year=investment_year + protected).all()
+    assert profile.astype(int).interp(year=investment_year + protected - 1).all()
     assert not profile.astype(int).interp(year=profile.year.max()).any()
diff --git a/tests/test_subsector.py b/tests/test_subsector.py
index 3ec688c2..5d18358a 100644
--- a/tests/test_subsector.py
+++ b/tests/test_subsector.py
@@ -1,4 +1,3 @@
-from collections.abc import Sequence
 from unittest.mock import MagicMock, patch
 
 import xarray as xr
@@ -110,22 +109,7 @@ def test_subsector_noninvesting_aggregation(market, model, technologies, tmp_pat
         commodity=technologies.commodity, region=technologies.region
     ).interp(year=[2020, 2025])
     assert all(agent.year == 2020 for agent in agents)
-    result = subsector.aggregate_lp(
-        technologies, market, time_period=5, current_year=2020
-    )
-
-    assert result is not None
-    assert len(result) == 2
-
-    lpcosts, lpconstraints = result
-    assert isinstance(lpcosts, xr.Dataset)
-    assert {"search_space", "decision"} == set(lpcosts.data_vars)
-    assert "agent" in lpcosts.coords
-    assert isinstance(lpconstraints, Sequence)
-    assert len(lpconstraints) == 1
-    assert all(isinstance(u, xr.Dataset) for u in lpconstraints)
-    # makes sure agent investment got called
-    assert all(agent.year == 2025 for agent in agents)
+    subsector.aggregate_lp(technologies, market, time_period=5, current_year=2020)
 
 
 def test_factory_smoke_test(model, technologies, tmp_path):

From 4fd2e79a486e026f31412f8362fc7e54f3ac8535 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 14:02:07 +0100
Subject: [PATCH 10/18] Small fix to another test

---
 src/muse/agents/agent.py |  3 +--
 src/muse/demand_share.py | 20 --------------------
 tests/test_utilities.py  |  2 +-
 3 files changed, 2 insertions(+), 23 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index 6a014e34..836adbab 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -161,8 +161,7 @@ def __init__(
         )
 
         self.year = year
-        """ Current year. 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:
diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index 01af3214..a5187c4f 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -676,23 +676,3 @@ def new_and_retro_demands(
         retro_demand = retro_demand.squeeze("year")
 
     return xr.Dataset({"new": new_demand, "retrofit": retro_demand})
-
-
-# def new_demand(
-#     capacity: xr.DataArray,
-#     market: xr.Dataset,
-#     technologies: xr.Dataset,
-#     production: Union[str, Mapping, Callable] = "maximum_production",
-#     current_year: Optional[int] = None,
-#     forecast: int = 5,
-# ) -> xr.DataArray:
-#     """Calculates the new demand that needs to be covered.
-
-#     It groups the demand related to an increase in consumption as well as the existing
-#     demand associated with decommissoned assets. Internally, it just calls
-#     `new_and_retro` demands and adds together both components.
-#     """
-#     demand = new_and_retro_demands(
-#         capacity, market, technologies, production, current_year, forecast
-#     )
-#     return (demand["new"] + demand["retrofit"]).rename("demand")
diff --git a/tests/test_utilities.py b/tests/test_utilities.py
index f829ac6b..44d7bf0f 100644
--- a/tests/test_utilities.py
+++ b/tests/test_utilities.py
@@ -37,7 +37,7 @@ def test_reduce_assets_with_zero_size(capacity: xr.DataArray):
 
     x = capacity.sel(asset=[])
     actual = reduce_assets(x)
-    assert actual is x
+    assert (actual == x).all()
 
 
 def test_broadcast_tech(technologies, capacity):

From 3b0cb49bd6e1df925406f825c4e465fb6ed2f706 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Fri, 4 Oct 2024 15:51:45 +0100
Subject: [PATCH 11/18] Remove more redundant code

---
 src/muse/demand_share.py   | 34 ++++++++-------------------
 src/muse/sectors/sector.py | 47 ++++++++++++++++----------------------
 2 files changed, 30 insertions(+), 51 deletions(-)

diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index a5187c4f..9c7e4212 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -112,9 +112,9 @@ def new_and_retro(
     agents: Sequence[AbstractAgent],
     market: xr.Dataset,
     technologies: xr.Dataset,
+    current_year: int,
+    forecast: int,
     production: Union[str, Mapping, Callable] = "maximum_production",
-    current_year: Optional[int] = None,
-    forecast: int = 5,
 ) -> xr.DataArray:
     r"""Splits demand across new and retro agents.
 
@@ -236,9 +236,6 @@ def decommissioning(capacity):
             technologies, capacity, year=[current_year, current_year + forecast]
         ).squeeze("year")
 
-    if current_year is None:
-        current_year = market.year.min()
-
     capacity = reduce_assets([u.assets.capacity for u in agents])
 
     demands = new_and_retro_demands(
@@ -323,9 +320,9 @@ def standard_demand(
     agents: Sequence[AbstractAgent],
     market: xr.Dataset,
     technologies: xr.Dataset,
+    current_year: int,
+    forecast: int,
     production: Union[str, Mapping, Callable] = "maximum_production",
-    current_year: Optional[int] = None,
-    forecast: int = 5,
 ) -> xr.DataArray:
     r"""Splits demand across new agents.
 
@@ -360,9 +357,6 @@ def decommissioning(capacity):
             technologies, capacity, year=[current_year, current_year + forecast]
         ).squeeze("year")
 
-    if current_year is None:
-        current_year = market.year.min()
-
     # Make sure there are no retrofit agents
     for agent in agents:
         if agent.category == "retrofit":
@@ -433,18 +427,15 @@ def unmet_forecasted_demand(
     agents: Sequence[AbstractAgent],
     market: xr.Dataset,
     technologies: xr.Dataset,
-    current_year: Optional[int] = None,
+    current_year: int,
+    forecast: int,
     production: Union[str, Mapping, Callable] = "maximum_production",
-    forecast: int = 5,
 ) -> xr.DataArray:
     """Forecast demand that cannot be serviced by non-decommissioned current assets."""
     from muse.commodities import is_enduse
     from muse.timeslices import QuantityType, convert_timeslice
     from muse.utilities import reduce_assets
 
-    if current_year is None:
-        current_year = market.year.min()
-
     year = current_year + forecast
     comm_usage = technologies.comm_usage.sel(commodity=market.commodity)
     smarket: xr.Dataset = market.where(is_enduse(comm_usage), 0).interp(year=year)
@@ -555,8 +546,8 @@ def new_consumption(
     capacity: xr.DataArray,
     market: xr.Dataset,
     technologies: xr.Dataset,
-    current_year: Optional[int] = None,
-    forecast: int = 5,
+    current_year: int,
+    forecast: int,
 ) -> xr.DataArray:
     r"""Computes share of the demand attributed to new agents.
 
@@ -576,9 +567,6 @@ def new_consumption(
 
     from muse.timeslices import QuantityType, convert_timeslice
 
-    if current_year is None:
-        current_year = market.year.min()
-
     # Interpolate market to forecast year
     market = market.interp(year=[current_year, current_year + forecast])
     current = market.sel(year=current_year, drop=True)
@@ -605,9 +593,9 @@ def new_and_retro_demands(
     capacity: xr.DataArray,
     market: xr.Dataset,
     technologies: xr.Dataset,
+    current_year: int,
+    forecast: int,
     production: Union[str, Mapping, Callable] = "maximum_production",
-    current_year: Optional[int] = None,
-    forecast: int = 5,
 ) -> xr.Dataset:
     """Splits demand into *new* and *retrofit* demand.
 
@@ -626,8 +614,6 @@ def new_and_retro_demands(
 
     production_method = production if callable(production) else prod_factory(production)
     assert callable(production_method)
-    if current_year is None:
-        current_year = market.year.min()
 
     # Interpolate market to forecast year
     smarket: xr.Dataset = market.interp(year=[current_year, current_year + forecast])
diff --git a/src/muse/sectors/sector.py b/src/muse/sectors/sector.py
index 31e04dae..d5c2b517 100644
--- a/src/muse/sectors/sector.py
+++ b/src/muse/sectors/sector.py
@@ -174,30 +174,18 @@ def forecast(self):
         If no agents with a "forecast" attribute are found, defaults to 5. It cannot be
         lower than 1 year.
         """
-        forecasts = [
-            getattr(agent, "forecast")
-            for agent in self.agents
-            if hasattr(agent, "forecast")
-        ]
-        if len(forecasts) == 0:
-            return 5
+        forecasts = [getattr(agent, "forecast") for agent in self.agents]
         return max(1, max(forecasts))
 
     def next(
         self,
         mca_market: xr.Dataset,
-        time_period: int | None = None,
-        current_year: int | None = None,
     ) -> xr.Dataset:
         """Advance sector by one time period.
 
         Args:
             mca_market:
                 Market with ``demand``, ``supply``, and ``prices``.
-            time_period:
-                Length of the time period in the framework. Defaults to the range of
-                ``mca_market.year``.
-            current_year: Current year of the simulation
 
         Returns:
             A market containing the ``supply`` offered by the sector, it's attendant
@@ -208,10 +196,8 @@ def next(
         def group_assets(x: xr.DataArray) -> xr.DataArray:
             return xr.Dataset(dict(x=x)).groupby("region").sum("asset").x
 
-        if time_period is None:
-            time_period = int(mca_market.year.max() - mca_market.year.min())
-        if current_year is None:
-            current_year = int(mca_market.year.min())
+        time_period = int(mca_market.year.max() - mca_market.year.min())
+        current_year = int(mca_market.year.min())
         getLogger(__name__).info(f"Running {self.name} for year {current_year}")
 
         # Agent interactions
@@ -347,6 +333,8 @@ def capacity(self) -> xr.DataArray:
             for u in self.agents
             if "dst_region" not in u.assets.capacity.dims
         ]
+
+        # Only nontraded assets
         if not traded:
             full_list = [
                 list(nontraded[i].year.values)
@@ -361,7 +349,9 @@ def capacity(self) -> xr.DataArray:
                 if "dst_region" not in u.assets.capacity.dims
             ]
             return reduce_assets(nontraded)
-        if not nontraded:
+
+        # Only traded assets
+        elif not nontraded:
             full_list = [
                 list(traded[i].year.values)
                 for i in range(len(traded))
@@ -375,15 +365,18 @@ def capacity(self) -> xr.DataArray:
                 if "dst_region" in u.assets.capacity.dims
             ]
             return reduce_assets(traded)
-        traded_results = reduce_assets(traded)
-        nontraded_results = reduce_assets(nontraded)
-        return reduce_assets(
-            [
-                traded_results,
-                nontraded_results
-                * (nontraded_results.region == traded_results.dst_region),
-            ]
-        )
+
+        # Both traded and nontraded assets
+        else:
+            traded_results = reduce_assets(traded)
+            nontraded_results = reduce_assets(nontraded)
+            return reduce_assets(
+                [
+                    traded_results,
+                    nontraded_results
+                    * (nontraded_results.region == traded_results.dst_region),
+                ]
+            )
 
     @property
     def agents(self) -> Iterator[AbstractAgent]:

From cc9d2370d1f95d495372f8c92a91d08a3dfd661c Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Mon, 14 Oct 2024 16:22:05 +0100
Subject: [PATCH 12/18] Fix test

---
 tests/test_demand_share.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/test_demand_share.py b/tests/test_demand_share.py
index 9661a180..1d51a610 100644
--- a/tests/test_demand_share.py
+++ b/tests/test_demand_share.py
@@ -345,6 +345,7 @@ def test_unmet_forecast_demand(technologies, coords, timeslice, stock_factory):
     asia_market = _matching_market(technologies, asia_stock, timeslice)
     usa_market = _matching_market(technologies, usa_stock, timeslice)
     market = xr.concat((asia_market, usa_market), dim="region")
+    current_year = market.year[0]
 
     # spoof some agents
     @dataclass
@@ -357,7 +358,9 @@ class Agent:
         Agent(0.7 * usa_stock.squeeze("region")),
         Agent(asia_stock.squeeze("region")),
     ]
-    result = unmet_forecasted_demand(agents, market, technologies)
+    result = unmet_forecasted_demand(
+        agents, market, technologies, current_year=current_year, forecast=5
+    )
     assert set(result.dims) == set(market.consumption.dims) - {"year"}
     assert result.values == approx(0)
 
@@ -367,7 +370,9 @@ class Agent:
         Agent(0.8 * usa_stock.squeeze("region")),
         Agent(1.1 * asia_stock.squeeze("region")),
     ]
-    result = unmet_forecasted_demand(agents, market, technologies)
+    result = unmet_forecasted_demand(
+        agents, market, technologies, current_year=current_year, forecast=5
+    )
     assert set(result.dims) == set(market.consumption.dims) - {"year"}
     assert result.values == approx(0)
 
@@ -376,7 +381,9 @@ class Agent:
         Agent(0.5 * usa_stock.squeeze("region")),
         Agent(0.5 * asia_stock.squeeze("region")),
     ]
-    result = unmet_forecasted_demand(agents, market, technologies)
+    result = unmet_forecasted_demand(
+        agents, market, technologies, current_year=current_year, forecast=5
+    )
     comm_usage = technologies.comm_usage.sel(commodity=market.commodity)
     enduse = is_enduse(comm_usage)
     assert (result.commodity == comm_usage.commodity).all()

From d054a3ba9222f602f274c9886dce59975a14926f Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Thu, 17 Oct 2024 15:09:31 +0100
Subject: [PATCH 13/18] A few more tiny changes (e.g. typing)

---
 src/muse/agents/agent.py      | 12 ++++++------
 src/muse/demand_share.py      |  3 +--
 src/muse/sectors/sector.py    |  4 +---
 src/muse/sectors/subsector.py |  2 +-
 4 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index 836adbab..7e69d463 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -81,7 +81,7 @@ def next(
         market: xr.Dataset,
         demand: xr.DataArray,
         time_period: int,
-    ):
+    ) -> None:
         """Increments agent to the next time point (e.g. performing investments)."""
 
     def __repr__(self):
@@ -243,7 +243,7 @@ def next(
         market: xr.Dataset,
         demand: xr.DataArray,
         time_period: int,
-    ):
+    ) -> None:
         self.year += time_period
 
     def compute_decision(
@@ -251,8 +251,8 @@ def compute_decision(
         technologies: xr.Dataset,
         market: xr.Dataset,
         demand: xr.DataArray,
-        search_space,
-    ):
+        search_space: xr.DataArray,
+    ) -> xr.DataArray:
         # Filter technologies according to the search space, forecast year and region
         techs = self.filter_input(
             technologies,
@@ -319,7 +319,7 @@ def next(
         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
@@ -399,7 +399,7 @@ 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(
diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index 9c7e4212..69fcca89 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -484,7 +484,7 @@ def _inner_split(
 
     # Calculates the demand divided by the number of assets times the number of agents
     # if the demand is bigger than zero and the total demand assigned with the "method"
-    # function is zero (i.e. no decrease in production).
+    # function is zero.
     unassigned = (demand / (len(shares) * len(summed_shares))).where(
         logical_and(demand > 1e-12, total <= 1e-12), 0
     )
@@ -583,7 +583,6 @@ def new_consumption(
     )
     assert isinstance(ts_capa, xr.DataArray)
 
-    #
     missing = unmet_demand(current, ts_capa, technologies)
     consumption = minimum(delta, missing)
     return consumption
diff --git a/src/muse/sectors/sector.py b/src/muse/sectors/sector.py
index d5c2b517..8d0f082f 100644
--- a/src/muse/sectors/sector.py
+++ b/src/muse/sectors/sector.py
@@ -71,7 +71,6 @@ def factory(cls, name: str, settings: Any) -> Sector:
         # Create outputs
         outputs = ofactory(*sector_settings.pop("outputs", []), sector_name=name)
 
-        #
         supply_args = sector_settings.pop(
             "supply", sector_settings.pop("dispatch_production", {})
         )
@@ -171,8 +170,7 @@ def __init__(
     def forecast(self):
         """Maximum forecast horizon across agents.
 
-        If no agents with a "forecast" attribute are found, defaults to 5. It cannot be
-        lower than 1 year.
+        It cannot be lower than 1 year.
         """
         forecasts = [getattr(agent, "forecast") for agent in self.agents]
         return max(1, max(forecasts))
diff --git a/src/muse/sectors/subsector.py b/src/muse/sectors/subsector.py
index 3beb7efa..edf82a19 100644
--- a/src/muse/sectors/subsector.py
+++ b/src/muse/sectors/subsector.py
@@ -87,7 +87,7 @@ def aggregate_lp(
         market: xr.Dataset,
         time_period,
         current_year,
-    ):
+    ) -> None:
         from muse.utilities import agent_concatenation, reduce_assets
 
         # Split demand across agents

From 0c84ba9feed0040ee38092ee688bf1f2a12b574c Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Thu, 17 Oct 2024 15:12:34 +0100
Subject: [PATCH 14/18] Remove inline comment

---
 src/muse/demand_share.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index 69fcca89..eeae9caf 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -489,7 +489,6 @@ def _inner_split(
         logical_and(demand > 1e-12, total <= 1e-12), 0
     )
 
-    # ???
     totals = {
         key: (share / share.sum("asset")).fillna(0) for key, share in shares.items()
     }

From 2adde920990ec569692bf3fe8f4f00080c465cb4 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Mon, 28 Oct 2024 14:45:46 +0000
Subject: [PATCH 15/18] Small changes to tidy up

---
 src/muse/agents/__init__.py   |  1 -
 src/muse/agents/agent.py      | 73 +++++++++++++++++------------------
 src/muse/agents/factories.py  |  2 +-
 src/muse/sectors/subsector.py | 14 -------
 tests/test_subsector.py       |  7 +---
 5 files changed, 37 insertions(+), 60 deletions(-)

diff --git a/src/muse/agents/__init__.py b/src/muse/agents/__init__.py
index 926d1ee1..392b04e7 100644
--- a/src/muse/agents/__init__.py
+++ b/src/muse/agents/__init__.py
@@ -2,7 +2,6 @@
     "AbstractAgent",
     "Agent",
     "InvestingAgent",
-    "factory",
     "agents_factory",
     "create_agent",
 ]
diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index 7e69d463..fb8fef3b 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -93,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,
@@ -246,40 +243,6 @@ def next(
     ) -> None:
         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,
-            technology=search_space.replacement,
-            year=self.forecast_year,
-        ).drop_vars("technology")
-
-        # Reduce dimensions of the demand array
-        reduced_demand = demand.sel(
-            {
-                k: search_space[k]
-                for k in set(demand.dims).intersection(search_space.dims)
-            }
-        )
-
-        # Filter prices according to the region
-        prices = self.filter_input(market.prices)
-
-        # Compute the objectives
-        objectives = self.objectives(
-            technologies=techs, demand=reduced_demand, prices=prices
-        )
-
-        # Compute the decision metric
-        decision = self.decision(objectives)
-        return decision
-
 
 class InvestingAgent(Agent):
     """Agent that performs investment for itself."""
@@ -393,6 +356,40 @@ def next(
         # 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,
+            technology=search_space.replacement,
+            year=self.forecast_year,
+        ).drop_vars("technology")
+
+        # Reduce dimensions of the demand array
+        reduced_demand = demand.sel(
+            {
+                k: search_space[k]
+                for k in set(demand.dims).intersection(search_space.dims)
+            }
+        )
+
+        # Filter prices according to the region
+        prices = self.filter_input(market.prices)
+
+        # Compute the objectives
+        objectives = self.objectives(
+            technologies=techs, demand=reduced_demand, prices=prices
+        )
+
+        # Compute the decision metric
+        decision = self.decision(objectives)
+        return decision
+
     def add_investments(
         self,
         technologies: xr.Dataset,
diff --git a/src/muse/agents/factories.py b/src/muse/agents/factories.py
index 236eb24e..f912ecd1 100644
--- a/src/muse/agents/factories.py
+++ b/src/muse/agents/factories.py
@@ -19,7 +19,7 @@ def create_standard_agent(
     interpolation: str = "linear",
     **kwargs,
 ):
-    """Creates retrofit agent from muse primitives."""
+    """Creates standard (noninvesting) agent from muse primitives."""
     from muse.filters import factory as filter_factory
 
     if share is not None:
diff --git a/src/muse/sectors/subsector.py b/src/muse/sectors/subsector.py
index edf82a19..03798af0 100644
--- a/src/muse/sectors/subsector.py
+++ b/src/muse/sectors/subsector.py
@@ -67,20 +67,6 @@ def invest(
         # Perform the investments
         self.aggregate_lp(technologies, market, time_period, current_year=current_year)
 
-    def assign_back_to_agents(
-        self,
-        technologies: xr.Dataset,
-        solution: xr.DataArray,
-        current_year: int,
-        time_period: int,
-    ):
-        agents = {u.uuid: u for u in self.agents}
-
-        for uuid, assets in solution.groupby("agent"):
-            agents[uuid].add_investments(
-                technologies, assets, current_year, time_period
-            )
-
     def aggregate_lp(
         self,
         technologies: xr.Dataset,
diff --git a/tests/test_subsector.py b/tests/test_subsector.py
index 5d18358a..b94d94fa 100644
--- a/tests/test_subsector.py
+++ b/tests/test_subsector.py
@@ -47,12 +47,7 @@ def test_subsector_investing_aggregation():
             subsector = Subsector(agents, commodities)
             initial_agents = deepcopy(agents)
             assert {agent.year for agent in agents} == {int(market.year.min())}
-            assert (
-                subsector.aggregate_lp(
-                    technologies, market, time_period=5, current_year=5
-                )
-                is None
-            )
+            subsector.aggregate_lp(technologies, market, time_period=5, current_year=5)
             assert {agent.year for agent in agents} == {int(market.year.min() + 5)}
             for initial, final in zip(initial_agents, agents):
                 assert initial.assets.sum() != final.assets.sum()

From 7f3370a53d23828746e3db87e0dcfc5d5a3bd776 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Tue, 29 Oct 2024 14:08:39 +0000
Subject: [PATCH 16/18] Be more explicit in _inner_split

---
 src/muse/demand_share.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/muse/demand_share.py b/src/muse/demand_share.py
index b5a81add..36a3c648 100644
--- a/src/muse/demand_share.py
+++ b/src/muse/demand_share.py
@@ -492,7 +492,9 @@ def _inner_split(
     # Calculates the demand divided by the number of assets times the number of agents
     # if the demand is bigger than zero and the total demand assigned with the "method"
     # function is zero.
-    unassigned = (demand / (len(shares) * len(summed_shares))).where(
+    n_agents = len(quantity)
+    n_assets = summed_shares.sizes["asset"]
+    unassigned = (demand / (n_agents * n_assets)).where(
         logical_and(demand > 1e-12, total <= 1e-12), 0
     )
 

From 789b093f873f45ad50257e04f4a3cbbf693a372e Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Wed, 30 Oct 2024 11:47:26 +0000
Subject: [PATCH 17/18] Apply suggestions from code review

---
 src/muse/agents/agent.py  | 8 ++------
 src/muse/investments.py   | 2 +-
 tests/test_investments.py | 4 ++--
 3 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/src/muse/agents/agent.py b/src/muse/agents/agent.py
index fb8fef3b..8d7b503a 100644
--- a/src/muse/agents/agent.py
+++ b/src/muse/agents/agent.py
@@ -267,13 +267,9 @@ def __init__(
 
         super().__init__(*args, **kwargs)
 
-        if investment is None:
-            investment = ifactory()
-        self.invest = investment
+        self.invest = investment or ifactory()
         """Method to use when fulfilling demand from rated set of techs."""
-        if not callable(constraints):
-            constraints = csfactory()
-        self.constraints = constraints
+        self.constraints = constraints or csfactory()
         """Creates a set of constraints limiting investment."""
 
     def next(
diff --git a/src/muse/investments.py b/src/muse/investments.py
index 802e7f0b..0e2abb83 100644
--- a/src/muse/investments.py
+++ b/src/muse/investments.py
@@ -177,7 +177,7 @@ def compute_investment(
 
 def cliff_retirement_profile(
     technical_life: xr.DataArray,
-    investment_year: int = 0,
+    investment_year: int,
     interpolation: str = "linear",
     **kwargs,
 ) -> xr.DataArray:
diff --git a/tests/test_investments.py b/tests/test_investments.py
index a4fc8769..38f49a26 100644
--- a/tests/test_investments.py
+++ b/tests/test_investments.py
@@ -44,7 +44,7 @@ def test_cliff_retirement_known_profile():
         name="technical_life",
     )
 
-    profile = cliff_retirement_profile(technical_life=lifetime)
+    profile = cliff_retirement_profile(technical_life=lifetime, investment_year=2020)
     expected = array(
         [
             [True, False, False, False],
@@ -73,7 +73,7 @@ def test_cliff_retirement_random_profile(protected):
     )
     effective_lifetime = (protected // lifetime + 1) * lifetime
 
-    investment_year = 5
+    investment_year = 2020
     profile = cliff_retirement_profile(
         technical_life=lifetime.clip(min=protected), investment_year=investment_year
     )

From 557efd84c485154970e3f8f78fbc46a891258ca2 Mon Sep 17 00:00:00 2001
From: Tom Bland <tom_bland@hotmail.co.uk>
Date: Wed, 30 Oct 2024 12:16:58 +0000
Subject: [PATCH 18/18] Switch copy to deepcopy

---
 src/muse/utilities.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/muse/utilities.py b/src/muse/utilities.py
index b56427ba..78ef36b7 100644
--- a/src/muse/utilities.py
+++ b/src/muse/utilities.py
@@ -139,9 +139,9 @@ def reduce_assets(
             installed   (asset) int32 12B 1990 1991 1990
         Dimensions without coordinates: asset
     """
-    from copy import copy
+    from copy import deepcopy
 
-    assets = copy(assets)
+    assets = deepcopy(assets)
 
     if operation is None:
 
