Ability to specify markets with unmet demand vars

src/simulation/investment.rs

@@ -201,7 +201,7 @@ 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)
-            .with_market_subset(&seen_markets)
+            .with_market_balance_subset(&seen_markets)
             .with_input_prices(&external_prices)
             .run(&format!("post {investment_set} investment"), writer)?;
 

src/simulation/optimisation.rs

@@ -105,28 +105,30 @@ impl VariableMap {
     ///
     /// * `problem` - The optimisation problem
     /// * `model` - The model
-    /// * `markets` - The subset of markets the problem is being run for
+    /// * `markets_to_allow_unmet_demand` - The subset of markets to assign unmet demand variables to
     fn add_unmet_demand_variables(
         &mut self,
         problem: &mut Problem,
         model: &Model,
-        markets: &[(CommodityID, RegionID)],
+        markets_to_allow_unmet_demand: &[(CommodityID, RegionID)],
     ) {
-        assert!(!markets.is_empty());
+        assert!(!markets_to_allow_unmet_demand.is_empty());
 
         // This line **must** come before we add more variables
         let start = problem.num_cols();
 
         // Add variables
         let voll = model.parameters.value_of_lost_load;
         self.unmet_demand_vars.extend(
-            iproduct!(markets.iter(), model.time_slice_info.iter_ids()).map(
-                |((commodity_id, region_id), time_slice)| {
-                    let key = (commodity_id.clone(), region_id.clone(), time_slice.clone());
-                    let var = problem.add_column(voll.value(), 0.0..);
-                    (key, var)
-                },
-            ),
+            iproduct!(
+                markets_to_allow_unmet_demand.iter(),
+                model.time_slice_info.iter_ids()
+            )
+            .map(|((commodity_id, region_id), time_slice)| {
+                let key = (commodity_id.clone(), region_id.clone(), time_slice.clone());
+                let var = problem.add_column(voll.value(), 0.0..);
+                (key, var)
+            }),
         );
 
         self.unmet_demand_var_idx = start..problem.num_cols();
@@ -348,6 +350,7 @@ pub struct DispatchRun<'model, 'run> {
     existing_assets: &'run [AssetRef],
     candidate_assets: &'run [AssetRef],
     markets_to_balance: &'run [(CommodityID, RegionID)],
+    markets_to_allow_unmet_demand: &'run [(CommodityID, RegionID)],
     input_prices: Option<&'run CommodityPrices>,
     year: u32,
 }
@@ -360,6 +363,7 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
             existing_assets: assets,
             candidate_assets: &[],
             markets_to_balance: &[],
+            markets_to_allow_unmet_demand: &[],
             input_prices: None,
             year,
         }
@@ -374,7 +378,7 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
     }
 
     /// Only apply commodity balance constraints to the specified subset of markets
-    pub fn with_market_subset(self, markets: &'run [(CommodityID, RegionID)]) -> Self {
+    pub fn with_market_balance_subset(self, markets: &'run [(CommodityID, RegionID)]) -> Self {
         assert!(!markets.is_empty());
 
         Self {
@@ -391,6 +395,14 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
         }
     }
 
+    /// Allow unmet demand variables for the specified subset of markets
+    pub fn with_unmet_demand_vars(self, markets: &'run [(CommodityID, RegionID)]) -> Self {
+        Self {
+            markets_to_allow_unmet_demand: markets,
+            ..self
+        }
+    }
+
     /// Perform the dispatch optimisation.
     ///
     /// # Arguments
@@ -425,9 +437,9 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
         }
 
         // Try running dispatch. If it fails because the model is infeasible, it is likely that this
-        // is due to unmet demand, in this case, we rerun dispatch including extra variables to
-        // track the unmet demand so we can report the offending regions/commodities to users
-        match self.run_without_unmet_demand(markets_to_balance) {
+        // is due to unmet demand, in this case, we rerun dispatch including with unmet demand
+        // variables for all markets so we can report the offending markets to users
+        match self.run_internal(markets_to_balance, self.markets_to_allow_unmet_demand) {
             Ok(solution) => Ok(solution),
             Err(ModelError::NonOptimal(HighsModelStatus::Infeasible)) => {
                 let markets = self
@@ -442,40 +454,35 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
                 bail!(
                     "The solver has indicated that the problem is infeasible, probably because \
                     the supplied assets could not meet the required demand. Demand was not met \
-                    for the following region and commodity pairs: {}",
+                    for the following markets: {}",
                     format_items_with_cap(markets)
                 )
             }
             Err(err) => Err(err)?,
         }
     }
 
-    /// Run dispatch without unmet demand variables
-    fn run_without_unmet_demand(
-        &self,
-        markets: &[(CommodityID, RegionID)],
-    ) -> Result<Solution<'model>, ModelError> {
-        self.run_internal(markets, /*allow_unmet_demand=*/ false)
-    }
-
     /// Run dispatch to diagnose which markets have unmet demand
     fn get_markets_with_unmet_demand(
         &self,
-        markets: &[(CommodityID, RegionID)],
+        markets_to_balance: &[(CommodityID, RegionID)],
     ) -> Result<IndexSet<(CommodityID, RegionID)>> {
-        let solution = self.run_internal(markets, /*allow_unmet_demand=*/ true)?;
+        // Run dispatch including unmet demand variables for all markets being balanced
+        let solution = self.run_internal(markets_to_balance, markets_to_balance)?;
+
+        // Collect markets with unmet demand
         Ok(solution
             .iter_unmet_demand()
             .filter(|(_, _, _, flow)| *flow > Flow(0.0))
             .map(|(commodity_id, region_id, _, _)| (commodity_id.clone(), region_id.clone()))
             .collect())
     }
 
-    /// Run dispatch for specified commodities, optionally including unmet demand variables
+    /// Run dispatch to balance the specified markets, allowing unmet demand for a subset of these
     fn run_internal(
         &self,
-        markets: &[(CommodityID, RegionID)],
-        allow_unmet_demand: bool,
+        markets_to_balance: &[(CommodityID, RegionID)],
+        markets_to_allow_unmet_demand: &[(CommodityID, RegionID)],
     ) -> Result<Solution<'model>, ModelError> {
         // Set up problem
         let mut problem = Problem::default();
@@ -488,10 +495,22 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
             self.year,
         );
 
-        // If unmet demand is enabled for this dispatch run (and is allowed by the model param) then
-        // we add variables representing unmet demand
-        if allow_unmet_demand {
-            variables.add_unmet_demand_variables(&mut problem, self.model, markets);
+        // Check that markets_to_allow_unmet_demand is a subset of markets_to_balance
+        for market in markets_to_allow_unmet_demand {
+            assert!(
+                markets_to_balance.contains(market),
+                "markets_to_allow_unmet_demand must be a subset of markets_to_balance. \
+                 Offending market: {market:?}"
+            );
+        }
+
+        // Add variables representing unmet demand
+        if !markets_to_allow_unmet_demand.is_empty() {
+            variables.add_unmet_demand_variables(
+                &mut problem,
+                self.model,
+                markets_to_allow_unmet_demand,
+            );
         }
 
         // Add constraints
@@ -501,7 +520,8 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
             &variables,
             self.model,
             &all_assets,
-            markets,
+            markets_to_balance,
+            markets_to_allow_unmet_demand,
             self.year,
         );
 

src/simulation/optimisation/constraints.rs

@@ -56,7 +56,8 @@ pub struct ConstraintKeys {
 /// * `variables` - The variables in the problem
 /// * `model` - The model
 /// * `assets` - The asset pool
-/// * `markets` - The subset of markets to apply constraints to
+/// * `markets_to_balance` - The subset of markets to apply balance constraints to
+/// * `markets_to_allow_unmet_demand` - The subset of markets to assign unmet demand variables to
 /// * `year` - Current milestone year
 ///
 /// # Returns
@@ -67,14 +68,22 @@ pub fn add_asset_constraints<'a, I>(
     variables: &VariableMap,
     model: &'a Model,
     assets: &I,
-    markets: &'a [(CommodityID, RegionID)],
+    markets_to_balance: &'a [(CommodityID, RegionID)],
+    markets_to_allow_unmet_demand: &'a [(CommodityID, RegionID)],
     year: u32,
 ) -> ConstraintKeys
 where
     I: Iterator<Item = &'a AssetRef> + Clone + 'a,
 {
-    let commodity_balance_keys =
-        add_commodity_balance_constraints(problem, variables, model, assets, markets, year);
+    let commodity_balance_keys = add_commodity_balance_constraints(
+        problem,
+        variables,
+        model,
+        assets,
+        markets_to_balance,
+        markets_to_allow_unmet_demand,
+        year,
+    );
 
     let activity_keys = add_activity_constraints(problem, variables);
 
@@ -97,7 +106,8 @@ fn add_commodity_balance_constraints<'a, I>(
     variables: &VariableMap,
     model: &'a Model,
     assets: &I,
-    markets: &'a [(CommodityID, RegionID)],
+    markets_to_balance: &'a [(CommodityID, RegionID)],
+    markets_to_allow_unmet_demand: &'a [(CommodityID, RegionID)],
     year: u32,
 ) -> CommodityBalanceKeys
 where
@@ -108,7 +118,7 @@ where
 
     let mut keys = Vec::new();
     let mut terms = Vec::new();
-    for (commodity_id, region_id) in markets {
+    for (commodity_id, region_id) in markets_to_balance {
         let commodity = &model.commodities[commodity_id];
         if !matches!(
             commodity.kind,
@@ -142,7 +152,7 @@ where
             }
 
             // Also include unmet demand variables if required
-            if !variables.unmet_demand_var_idx.is_empty() {
+            if markets_to_allow_unmet_demand.contains(&(commodity_id.clone(), region_id.clone())) {
                 for (time_slice, _) in ts_selection.iter(&model.time_slice_info) {
                     let var = variables.get_unmet_demand_var(commodity_id, region_id, time_slice);
                     terms.push((var, 1.0));