New approach for partial dispatch

benches/assets.rs

@@ -9,7 +9,8 @@ use muse2::model::Model;
 use muse2::output::DataWriter;
 use muse2::process::{Process, ProcessID};
 use muse2::simulation::candidate_assets_for_next_year;
-use muse2::simulation::investment::{collect_preset_demands_for_year, select_best_assets};
+use muse2::simulation::demand::collect_preset_demands_for_year;
+use muse2::simulation::investment::select_best_assets;
 use muse2::simulation::market::{
     collect_agent_limits, get_asset_options, get_demand_portion_for_market, get_responsible_agents,
 };
@@ -95,13 +96,15 @@ fn calculate_seed_prices(
     candidates: &[AssetRef],
     writer: &mut DataWriter,
 ) -> Prices {
-    let solution_existing = DispatchRun::new(model, base_year_assets, BASE_YEAR)
+    let market_demands = collect_preset_demands_for_year(&model.commodities, BASE_YEAR);
+    let solution_existing = DispatchRun::new(model, base_year_assets, BASE_YEAR, &market_demands)
         .run("bench setup: without candidates", writer)
         .expect("Dispatch without candidates failed");
-    let solution_with_candidates = DispatchRun::new(model, base_year_assets, BASE_YEAR)
-        .with_candidates(candidates)
-        .run("bench setup: with candidates", writer)
-        .expect("Dispatch with candidates failed");
+    let solution_with_candidates =
+        DispatchRun::new(model, base_year_assets, BASE_YEAR, &market_demands)
+            .with_candidates(candidates)
+            .run("bench setup: with candidates", writer)
+            .expect("Dispatch with candidates failed");
 
     calculate_prices(
         model,

src/simulation.rs

@@ -16,6 +16,8 @@ pub mod optimisation;
 use optimisation::{DispatchRun, FlowMap};
 pub mod investment;
 use investment::perform_agent_investment;
+pub mod demand;
+use demand::collect_preset_demands_for_year;
 pub mod market;
 pub mod prices;
 pub use prices::PriceMap;
@@ -178,13 +180,15 @@ fn run_dispatch_for_year(
     debug_assert!(assets.iter().all(|asset| !asset.is_candidate()));
     debug_assert!(candidates.iter().all(|asset| asset.is_candidate()));
 
+    let market_demands = collect_preset_demands_for_year(&model.commodities, year);
+
     // Run dispatch optimisation with existing assets only, if there are any. If not, then assume no
     // flows (i.e. all are zero)
     let (solution_existing, flow_map) = if assets.is_empty() {
         (None, FlowMap::default())
     } else {
-        let solution =
-            DispatchRun::new(model, assets, year).run("final without candidates", writer)?;
+        let solution = DispatchRun::new(model, assets, year, &market_demands)
+            .run("final without candidates", writer)?;
         let flow_map = solution.create_flow_map();
         (Some(solution), flow_map)
     };
@@ -195,7 +199,7 @@ fn run_dispatch_for_year(
         None
     } else {
         Some(
-            DispatchRun::new(model, assets, year)
+            DispatchRun::new(model, assets, year, &market_demands)
                 .with_candidates(candidates)
                 .run("final with candidates", writer)?,
         )

src/simulation/demand.rs

@@ -0,0 +1,37 @@
+//! Code for collecting and storing commodity demand data
+use crate::commodity::{CommodityID, CommodityMap};
+use crate::region::RegionID;
+use crate::time_slice::TimeSliceSelection;
+use crate::units::Flow;
+use indexmap::IndexMap;
+
+/// A map of demand across time-slice selections for a specific market
+pub type DemandMap = IndexMap<TimeSliceSelection, Flow>;
+
+/// Demand for a given combination of commodity, region and time-slice selection
+pub type AllDemandMap = IndexMap<(CommodityID, RegionID, TimeSliceSelection), Flow>;
+
+/// Collect the preset commodity demands for a given year into a map of commodity, region and
+/// time slice selection to demand.
+///
+/// Demand for each commodity is stored at its natural time-slice selection level, matching the
+/// balance level at which the investment appraisal operates.
+pub fn collect_preset_demands_for_year(commodities: &CommodityMap, year: u32) -> AllDemandMap {
+    let mut demand_map = AllDemandMap::new();
+    for (commodity_id, commodity) in commodities {
+        for ((region_id, data_year, time_slice_selection), demand) in &commodity.demand {
+            if *data_year != year {
+                continue;
+            }
+            demand_map.insert(
+                (
+                    commodity_id.clone(),
+                    region_id.clone(),
+                    time_slice_selection.clone(),
+                ),
+                *demand,
+            );
+        }
+    }
+    demand_map
+}

src/simulation/investment.rs

@@ -1,8 +1,9 @@
 //! Code for performing agent investment.
+use super::demand::{AllDemandMap, DemandMap, collect_preset_demands_for_year};
 use super::optimisation::{DispatchRun, FlowMap};
 use crate::agent::{Agent, AgentID};
 use crate::asset::{Asset, AssetRef};
-use crate::commodity::{Commodity, CommodityID, CommodityMap};
+use crate::commodity::{Commodity, CommodityID};
 use crate::model::Model;
 use crate::output::DataWriter;
 use crate::process::ProcessID;
@@ -13,7 +14,6 @@ use crate::timeit::InvestmentTimer;
 use crate::units::{ActivityPerCapacity, Capacity, Flow, FlowPerCapacity};
 use anyhow::{Result, ensure};
 use context_manager;
-use indexmap::IndexMap;
 use itertools::Itertools;
 use log::{debug, warn};
 use rayon::prelude::*;
@@ -27,12 +27,6 @@ use appraisal::{
     sort_and_filter_appraisal_outputs,
 };
 
-/// A map of demand across time-slice selections for a specific market
-pub type DemandMap = IndexMap<TimeSliceSelection, Flow>;
-
-/// Demand for a given combination of commodity, region and time-slice selection
-pub type AllDemandMap = IndexMap<(CommodityID, RegionID, TimeSliceSelection), Flow>;
-
 /// Perform agent investment to determine capacity investment of new assets for next milestone year.
 ///
 /// # Arguments
@@ -68,29 +62,11 @@ pub fn perform_agent_investment(
         investment_order.iter().join(" -> ")
     );
 
-    // Keep track of the markets that have been seen so far. This will be used to apply
-    // balance constraints in the dispatch optimisation - we only apply balance constraints for
-    // markets that have been seen so far.
-    let mut seen_markets = Vec::new();
-
     // Iterate over market sets in the investment order for this year
     for market_set in investment_order {
         // Select assets for this market set
-        let selected_assets = market_set.select_assets(
-            model,
-            year,
-            &net_demand,
-            existing_assets,
-            prices,
-            &seen_markets,
-            &all_selected_assets,
-            writer,
-        )?;
-
-        // Update our list of seen markets
-        for market in market_set.iter_markets() {
-            seen_markets.push(market.clone());
-        }
+        let selected_assets =
+            market_set.select_assets(model, year, &net_demand, existing_assets, prices, writer)?;
 
         // If no assets have been selected, skip dispatch optimisation
         // **TODO**: this probably means there's no demand for the market, which we could
@@ -110,9 +86,10 @@ pub fn perform_agent_investment(
 
         // As upstream markets by definition will not yet have producers, we explicitly set
         // their prices using external values so that they don't appear free
-        let solution = DispatchRun::new(model, &all_selected_assets, year)
+        let current_markets: Vec<_> = market_set.iter_markets().cloned().collect();
+        let solution = DispatchRun::new(model, &selected_assets, year, &net_demand)
             .without_commodity_constraints()
-            .with_market_balance_subset(&seen_markets)
+            .with_market_balance_subset(&current_markets)
             .with_input_prices(&prices.shadow)
             .run(&format!("post {market_set} investment"), writer)?;
 
@@ -127,31 +104,6 @@ pub fn perform_agent_investment(
     Ok(all_selected_assets)
 }
 
-/// Collect the preset commodity demands for a given year into a map of commodity, region and
-/// time slice selection to demand.
-///
-/// Demand for each commodity is stored at its natural time-slice selection level, matching the
-/// balance level at which the investment appraisal operates.
-pub fn collect_preset_demands_for_year(commodities: &CommodityMap, year: u32) -> AllDemandMap {
-    let mut demand_map = AllDemandMap::new();
-    for (commodity_id, commodity) in commodities {
-        for ((region_id, data_year, time_slice_selection), demand) in &commodity.demand {
-            if *data_year != year {
-                continue;
-            }
-            demand_map.insert(
-                (
-                    commodity_id.clone(),
-                    region_id.clone(),
-                    time_slice_selection.clone(),
-                ),
-                *demand,
-            );
-        }
-    }
-    demand_map
-}
-
 /// Update net demand map with flows from a set of assets
 ///
 /// Non-primary output flows are ignored. This way, demand profiles aren't affected by production

src/simulation/market.rs

@@ -7,15 +7,14 @@ use crate::model::Model;
 use crate::output::DataWriter;
 use crate::process::{Process, ProcessID};
 use crate::region::RegionID;
+use crate::simulation::demand::{AllDemandMap, DemandMap};
 use crate::simulation::investment::{
-    AllDemandMap, DemandMap, calculate_candidate_asset_capacity_scale, select_best_assets,
-    update_net_demand_map,
+    calculate_candidate_asset_capacity_scale, select_best_assets, update_net_demand_map,
 };
 use crate::simulation::prices::Prices;
 use crate::time_slice::TimeSliceInfo;
 use crate::units::{Capacity, Dimensionless, Flow};
 use anyhow::{Context, Result};
-use indexmap::IndexMap;
 use itertools::{Itertools, chain};
 use log::debug;
 use std::collections::HashMap;
@@ -56,8 +55,6 @@ impl MarketSet {
     /// * `demand` – Net demand profiles available to all markets before selection.
     /// * `existing_assets` – Assets already commissioned in the system.
     /// * `prices` – Commodity price assumptions to use when valuing investments.
-    /// * `seen_markets` – Markets for which investments have already been settled.
-    /// * `previously_selected_assets` – Assets chosen in earlier market sets.
     /// * `writer` – Data sink used to log optimisation artefacts.
     #[allow(clippy::too_many_arguments)]
     pub fn select_assets(
@@ -67,8 +64,6 @@ impl MarketSet {
         demand: &AllDemandMap,
         existing_assets: &[AssetRef],
         prices: &Prices,
-        seen_markets: &[(CommodityID, RegionID)],
-        previously_selected_assets: &[AssetRef],
         writer: &mut DataWriter,
     ) -> Result<Vec<AssetRef>> {
         match self {
@@ -91,8 +86,6 @@ impl MarketSet {
                     demand,
                     existing_assets,
                     prices,
-                    seen_markets,
-                    previously_selected_assets,
                     writer,
                 )
                 .with_context(|| {
@@ -114,8 +107,6 @@ impl MarketSet {
                         demand,
                         existing_assets,
                         prices,
-                        seen_markets,
-                        previously_selected_assets,
                         writer,
                     )?;
                     all_assets.extend(assets);
@@ -247,64 +238,57 @@ pub fn select_assets_for_cycle(
     demand: &AllDemandMap,
     existing_assets: &[AssetRef],
     prices: &Prices,
-    seen_markets: &[(CommodityID, RegionID)],
-    previously_selected_assets: &[AssetRef],
     writer: &mut DataWriter,
 ) -> Result<Vec<AssetRef>> {
     // Precompute a joined string for logging
     let markets_str = markets.iter().map(|(c, r)| format!("{c}|{r}")).join(", ");
 
     // Iterate over the markets to select assets
-    let mut current_demand = demand.clone();
-    let mut assets_for_cycle = IndexMap::new();
-    for (idx, (commodity_id, region_id)) in markets.iter().enumerate() {
+    let mut net_demand = demand.clone();
+    let mut all_selected_assets = Vec::new();
+    for market in markets {
+        let (commodity_id, region_id) = market.clone();
+
         // Select assets for this market
-        let assets = select_assets_for_single_market(
+        let selected_assets = select_assets_for_single_market(
             model,
-            commodity_id,
-            region_id,
+            &commodity_id,
+            &region_id,
             year,
-            &current_demand,
+            &net_demand,
             existing_assets,
             prices,
             writer,
         )?;
-        assets_for_cycle.insert((commodity_id.clone(), region_id.clone()), assets);
 
-        // Assemble full list of assets for dispatch (previously selected + all chosen so far)
-        let mut all_assets = previously_selected_assets.to_vec();
-        let assets_for_cycle_flat: Vec<_> = assets_for_cycle
-            .values()
-            .flat_map(|v| v.iter().cloned())
-            .collect();
-        all_assets.extend_from_slice(&assets_for_cycle_flat);
+        // If no assets have been selected, skip dispatch optimisation
+        // **TODO**: this probably means there's no demand for the market, which we could
+        // presumably preempt
+        if selected_assets.is_empty() {
+            continue;
+        }
 
-        // We balance all previously seen markets plus all cycle markets up to and including this one
-        let mut markets_to_balance = seen_markets.to_vec();
-        markets_to_balance.extend_from_slice(&markets[0..=idx]);
+        all_selected_assets.extend(selected_assets.iter().cloned());
 
         // Run dispatch
-        let solution = DispatchRun::new(model, &all_assets, year)
+        let solution = DispatchRun::new(model, &selected_assets, year, &net_demand)
             .without_commodity_constraints()
-            .with_market_balance_subset(&markets_to_balance)
+            .with_market_balance_subset(std::slice::from_ref(market))
             .run(
                 &format!("cycle ({markets_str}) post {commodity_id}|{region_id} investment"),
                 writer,
             )
             .with_context(|| format!("Dispatch failed for cycle ({markets_str})"))?;
 
-        // Calculate new net demand map with all assets selected so far
-        current_demand.clone_from(demand);
+        // Update demand map with flows from newly selected assets
         update_net_demand_map(
-            &mut current_demand,
+            &mut net_demand,
             &solution.create_flow_map(),
-            &assets_for_cycle_flat,
+            &selected_assets,
         );
     }
 
-    // Collect assets
-    let all_cycle_assets: Vec<_> = assets_for_cycle.into_values().flatten().collect();
-    Ok(all_cycle_assets)
+    Ok(all_selected_assets)
 }
 
 /// Get a portion of the demand profile for this market

src/simulation/optimisation.rs

@@ -8,6 +8,7 @@ use crate::model::Model;
 use crate::output::DataWriter;
 use crate::region::RegionID;
 use crate::simulation::PriceMap;
+use crate::simulation::demand::AllDemandMap;
 use crate::time_slice::{TimeSliceID, TimeSliceInfo, TimeSliceSelection};
 use crate::units::{Activity, Flow, Money, MoneyPerActivity, MoneyPerFlow};
 use anyhow::{Context, Result, anyhow, bail};
@@ -401,19 +402,26 @@ pub struct DispatchRun<'model, 'run> {
     existing_assets: &'run [AssetRef],
     candidate_assets: &'run [AssetRef],
     markets_to_balance: &'run [(CommodityID, RegionID)],
+    market_demands: &'run AllDemandMap,
     input_prices: Option<&'run PriceMap>,
     include_commodity_constraints: bool,
     year: u32,
 }
 
 impl<'model, 'run> DispatchRun<'model, 'run> {
     /// Create a new [`DispatchRun`] for the specified model and assets for a given year
-    pub fn new(model: &'model Model, assets: &'run [AssetRef], year: u32) -> Self {
+    pub fn new(
+        model: &'model Model,
+        assets: &'run [AssetRef],
+        year: u32,
+        market_demands: &'run AllDemandMap,
+    ) -> Self {
         Self {
             model,
             existing_assets: assets,
             candidate_assets: &[],
             markets_to_balance: &[],
+            market_demands,
             input_prices: None,
             include_commodity_constraints: true,
             year,
@@ -691,6 +699,7 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
             self.model,
             &all_assets,
             markets_to_balance,
+            self.market_demands,
             self.year,
             self.candidate_assets,
             include_commodity_constraints,