Enforce total capacity limit constraint in investment

src/input/process/investment_constraints.rs

@@ -199,7 +199,10 @@ where
                 .addition_limit
                 .map(|limit| limit * Year(years_since_prev as f64));
 
-            let constraint = Arc::new(ProcessInvestmentConstraint { addition_limit });
+            let constraint = Arc::new(ProcessInvestmentConstraint {
+                addition_limit,
+                total_capacity_limit: record.total_capacity_limit,
+            });
 
             try_insert(process_map, &(region.clone(), year), constraint.clone())?;
         }

src/process.rs

@@ -77,13 +77,14 @@ impl Process {
         self.years.contains(&year)
     }
 
-    /// Calculate an agent's share of the addition limit for this process in a region and year
-    /// based on commodity portion.
-    pub fn agent_addition_limit(
+    /// Calculate an agent's share of any limit for this process in a region and year based on
+    /// commodity portion.
+    fn agent_limit(
         &self,
         region_id: &RegionID,
         commission_year: u32,
         commodity_portion: Dimensionless,
+        get_limit: fn(&ProcessInvestmentConstraint) -> Option<Capacity>,
     ) -> Option<Capacity> {
         assert!(
             commodity_portion >= Dimensionless(0.0) && commodity_portion <= Dimensionless(1.0),
@@ -92,11 +93,39 @@ impl Process {
 
         self.investment_constraints
             .get(&(region_id.clone(), commission_year))
-            .and_then(|constraint| {
-                constraint
-                    .get_addition_limit()
-                    .map(|limit| limit * commodity_portion)
-            })
+            .and_then(|constraint| get_limit(constraint).map(|limit| limit * commodity_portion))
+    }
+
+    /// Calculate an agent's share of the addition limit for this process in a region and year
+    /// based on commodity portion.
+    pub fn agent_addition_limit(
+        &self,
+        region_id: &RegionID,
+        commission_year: u32,
+        commodity_portion: Dimensionless,
+    ) -> Option<Capacity> {
+        self.agent_limit(
+            region_id,
+            commission_year,
+            commodity_portion,
+            ProcessInvestmentConstraint::get_addition_limit,
+        )
+    }
+
+    /// Calculate an agent's share of the total limit for this process in a region and year
+    /// based on commodity portion.
+    pub fn agent_total_limit(
+        &self,
+        region_id: &RegionID,
+        commission_year: u32,
+        commodity_portion: Dimensionless,
+    ) -> Option<Capacity> {
+        self.agent_limit(
+            region_id,
+            commission_year,
+            commodity_portion,
+            ProcessInvestmentConstraint::get_total_limit,
+        )
     }
 }
 
@@ -527,24 +556,27 @@ pub struct ProcessInvestmentConstraint {
     /// Addition constraint: Limit an agent can invest in the process, shared according to the
     /// agent's proportion of the process's primary commodity demand
     pub addition_limit: Option<Capacity>,
+    /// Total capacity limit for the process
+    pub total_capacity_limit: Option<Capacity>,
 }
 
 impl ProcessInvestmentConstraint {
-    /// Calculate the effective addition limit
-    ///
-    /// For now, this just returns `addition_limit`, but in the future when we add growth
-    /// limits and total capacity limits, this will have more complex logic which will depend on the
-    /// current total capacity.
+    /// Get the addition limit
     pub fn get_addition_limit(&self) -> Option<Capacity> {
         self.addition_limit
     }
+
+    /// Get the total capacity limit allowed
+    pub fn get_total_limit(&self) -> Option<Capacity> {
+        self.total_capacity_limit
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::commodity::{CommodityLevyMap, CommodityType, DemandMap, PricingStrategy};
-    use crate::fixture::{assert_error, region_id, time_slice, time_slice_info2};
+    use crate::fixture::{assert_error, process, region_id, time_slice, time_slice_info2};
     use crate::time_slice::TimeSliceLevel;
     use crate::time_slice::TimeSliceSelection;
     use float_cmp::assert_approx_eq;
@@ -1156,4 +1188,71 @@ mod tests {
             "Availability limit for season winter clashes with time slice limits"
         );
     }
+
+    #[rstest]
+    #[case(Dimensionless(1.1))]
+    #[case(Dimensionless(-0.1))]
+    #[should_panic(expected = "commodity_portion must be between 0 and 1 inclusive")]
+    fn agent_limit_invalid_commodity_portion(
+        process: Process,
+        region_id: RegionID,
+        #[case] commodity_portion: Dimensionless,
+    ) {
+        process.agent_limit(
+            &region_id,
+            2015,
+            commodity_portion,
+            ProcessInvestmentConstraint::get_addition_limit,
+        );
+    }
+
+    #[rstest]
+    fn agent_addition_limit_no_constraint(process: Process, region_id: RegionID) {
+        assert!(
+            process
+                .agent_addition_limit(&region_id, 2015, Dimensionless(1.0))
+                .is_none()
+        );
+    }
+
+    #[rstest]
+    fn agent_addition_limit_scaled_by_portion(mut process: Process, region_id: RegionID) {
+        process.investment_constraints.insert(
+            (region_id.clone(), 2015),
+            Arc::new(ProcessInvestmentConstraint {
+                addition_limit: Some(crate::units::Capacity(10.0)),
+                total_capacity_limit: Some(crate::units::Capacity(100.0)),
+            }),
+        );
+
+        let result = process
+            .agent_addition_limit(&region_id, 2015, Dimensionless(0.5))
+            .unwrap();
+        assert_eq!(result, Capacity(5.0));
+    }
+
+    #[rstest]
+    fn agent_total_limit_no_constraint(process: Process, region_id: RegionID) {
+        assert!(
+            process
+                .agent_total_limit(&region_id, 2015, Dimensionless(1.0))
+                .is_none()
+        );
+    }
+
+    #[rstest]
+    fn agent_total_limit_scaled_by_portion(mut process: Process, region_id: RegionID) {
+        process.investment_constraints.insert(
+            (region_id.clone(), 2015),
+            Arc::new(ProcessInvestmentConstraint {
+                addition_limit: Some(crate::units::Capacity(10.0)),
+                total_capacity_limit: Some(crate::units::Capacity(100.0)),
+            }),
+        );
+
+        let result = process
+            .agent_total_limit(&region_id, 2015, Dimensionless(0.5))
+            .unwrap();
+        assert_eq!(result, Capacity(50.0));
+    }
 }

src/simulation/investment.rs

@@ -354,6 +354,7 @@ pub fn select_best_assets(
     model: &Model,
     mut opt_assets: Vec<AssetRef>,
     agent_addition_limits: HashMap<ProcessID, Capacity>,
+    agent_total_limits: HashMap<ProcessID, Capacity>,
     commodity: &Commodity,
     agent: &Agent,
     region_id: &RegionID,
@@ -372,6 +373,9 @@ pub fn select_best_assets(
         &remaining_agent_addition_limits,
     );
 
+    // Remaining total capacity for all assets
+    let mut remaining_agent_total_limit = agent_total_limits;
+
     // Store commissioned units available for retention and replace assets with single units
     let mut available_retention_units = prepare_commissioned_assets_for_retention(&mut opt_assets);
 
@@ -462,6 +466,7 @@ pub fn select_best_assets(
             best_output.asset,
             &mut opt_assets,
             &mut remaining_agent_addition_limits,
+            &mut remaining_agent_total_limit,
             &mut available_retention_units,
             &mut best_assets,
         );
@@ -533,12 +538,14 @@ fn remove_candidates_exceeding_agent_addition_limits(
 /// * `best_asset` - The asset that has been selected as the best option in this round
 /// * `opt_assets` - The list of remaining asset options to be considered in future rounds
 /// * `remaining_agent_addition_limits` - The remaining agent addition limits for processes
+/// * `remaining_agent_total_limit` - The remaining capacity for processes
 /// * `available_retention_units` - The commissioned units available for retention
 /// * `best_assets` - The list of assets that have been selected so far
 fn record_asset_selection(
     best_asset: AssetRef,
     opt_assets: &mut Vec<AssetRef>,
     remaining_agent_addition_limits: &mut HashMap<ProcessID, Capacity>,
+    remaining_agent_total_limit: &mut HashMap<ProcessID, Capacity>,
     available_retention_units: &mut HashMap<AssetRef, u32>,
     best_assets: &mut Vec<AssetRef>,
 ) {
@@ -547,6 +554,22 @@ fn record_asset_selection(
         "Invalid asset type"
     );
 
+    // Remove capacity from the total capacity limit, if applicable.
+    if let Some(remaining_capacity) = remaining_agent_total_limit.get_mut(best_asset.process_id()) {
+        *remaining_capacity -= best_asset.total_capacity();
+
+        // If there's not enough capacity remaining to install any more units, remove the
+        // asset from the investment options.
+        if *remaining_capacity < best_asset.total_capacity() {
+            let old_idx = opt_assets
+                .iter()
+                .position(|asset| *asset == best_asset)
+                .unwrap();
+            opt_assets.swap_remove(old_idx);
+            remaining_agent_total_limit.remove(best_asset.process_id());
+        }
+    }
+
     // Update the remaining agent addition limit for the selected asset, if applicable, and remove it
     // from the options if the limit is exhausted.
     if best_asset.is_candidate() {

src/simulation/market.rs

@@ -5,7 +5,7 @@ use crate::asset::{Asset, AssetCapacity, AssetIterator, AssetRef, AssetState};
 use crate::commodity::{Commodity, CommodityID};
 use crate::model::Model;
 use crate::output::DataWriter;
-use crate::process::ProcessID;
+use crate::process::{Process, ProcessID};
 use crate::region::RegionID;
 use crate::simulation::investment::{
     AllDemandMap, DemandMap, calculate_candidate_asset_capacity_scale, select_best_assets,
@@ -194,14 +194,31 @@ pub fn select_assets_for_single_market(
         .collect::<Vec<_>>();
 
         // Calculate the agent's share of addition limits for candidate processes
-        let agent_addition_limits =
-            collect_agent_addition_limits(agent, region_id, &commodity.id, year, commodity_portion);
+        let agent_addition_limits = collect_agent_limits(
+            agent,
+            region_id,
+            commodity_id,
+            year,
+            commodity_portion,
+            Process::agent_addition_limit,
+        );
+
+        // Calculate the agent's share of total capacity limits for all processes
+        let agent_total_limits = collect_agent_limits(
+            agent,
+            region_id,
+            commodity_id,
+            year,
+            commodity_portion,
+            Process::agent_total_limit,
+        );
 
         // Choose assets from among existing pool and candidates
         let best_assets = select_best_assets(
             model,
             opt_assets,
             agent_addition_limits,
+            agent_total_limits,
             commodity,
             agent,
             region_id,
@@ -468,22 +485,34 @@ fn get_candidate_assets<'a>(
         })
 }
 
-/// Agent addition limits are based on process addition limits that have already been
-/// scaled from annual addition limits to the interval since the previous milestone year.
-/// The resulting limit is then scaled according to the agent's portion of commodity demand.
-pub fn collect_agent_addition_limits(
+/// Collects capacity limits for all processes in the agent's search space for a given market.
+///
+/// Processes without a defined limit are excluded from the returned map. The limit type is
+/// determined by `get_agent_limit`, which should be one of [`Process::agent_addition_limit`] (the
+/// agent's share of the annual addition limit) or [`Process::agent_total_limit`] (the agent's share
+/// of the maximum total installed capacity).
+///
+/// # Arguments
+///
+/// * `agent` – Agent whose search space is queried.
+/// * `region_id` – Region for which limits are calculated.
+/// * `commodity_id` – Commodity for which limits are calculated.
+/// * `year` – Milestone year being solved.
+/// * `commodity_portion` – Agent's fractional share of commodity demand, used to scale limits.
+/// * `get_agent_limit` – Method on [`Process`] that returns the limit value.
+fn collect_agent_limits(
     agent: &Agent,
     region_id: &RegionID,
     commodity_id: &CommodityID,
     year: u32,
     commodity_portion: Dimensionless,
+    get_agent_limit: fn(&Process, &RegionID, u32, Dimensionless) -> Option<Capacity>,
 ) -> HashMap<ProcessID, Capacity> {
     agent
         .iter_search_space(region_id, commodity_id, year)
         .filter_map(|process| {
-            process
-                .agent_addition_limit(region_id, year, commodity_portion)
-                .map(|agent_limit| (process.id.clone(), agent_limit))
+            get_agent_limit(process, region_id, year, commodity_portion)
+                .map(|limit| (process.id.clone(), limit))
         })
         .collect()
 }
@@ -522,43 +551,57 @@ mod tests {
     }
 
     #[rstest]
-    fn collect_agent_addition_limits_uses_search_space(mut process: Process, region_id: RegionID) {
+    fn collect_agent_limits_uses_search_space(mut process: Process, region_id: RegionID) {
         process.investment_constraints.insert(
             (region_id.clone(), 2015),
             Arc::new(ProcessInvestmentConstraint {
                 addition_limit: Some(crate::units::Capacity(10.0)),
+                total_capacity_limit: Some(crate::units::Capacity(100.0)),
             }),
         );
         let commodity_id = "commodity".into();
         let process_id = process.id.clone();
         let agent = agent_with_process(process, &region_id, &commodity_id);
 
-        let result = collect_agent_addition_limits(
+        let result = collect_agent_limits(
             &agent,
             &region_id,
             &commodity_id,
             2015,
             Dimensionless(0.5),
+            Process::agent_addition_limit,
         );
 
         assert_eq!(result.get(&process_id), Some(&crate::units::Capacity(5.0)));
+
+        let result = collect_agent_limits(
+            &agent,
+            &region_id,
+            &commodity_id,
+            2015,
+            Dimensionless(0.5),
+            Process::agent_total_limit,
+        );
+
+        assert_eq!(result.get(&process_id), Some(&crate::units::Capacity(50.0)));
     }
 
     #[rstest]
-    fn collect_agent_addition_limits_excludes_processes_without_limits(
+    fn collect_agent_limits_excludes_processes_without_limits(
         process: Process,
         region_id: RegionID,
     ) {
         let commodity_id = "commodity".into();
         let process_id = process.id.clone();
         let agent = agent_with_process(process, &region_id, &commodity_id);
 
-        let result = collect_agent_addition_limits(
+        let result = collect_agent_limits(
             &agent,
             &region_id,
             &commodity_id,
             2015,
             Dimensionless(1.0),
+            Process::agent_addition_limit,
         );
 
         assert!(!result.contains_key(&process_id));