Allow debug files in unmet demand scenario

src/simulation/optimisation.rs

@@ -455,16 +455,7 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
     ///
     /// A solution containing new commodity flows for assets and prices for (some) commodities or an
     /// error.
-    pub fn run(self, run_description: &str, writer: &mut DataWriter) -> Result<Solution<'model>> {
-        let solution = self.run_no_save()?;
-        writer.write_dispatch_debug_info(self.year, run_description, &solution)?;
-        Ok(solution)
-    }
-
-    /// Run dispatch without saving the results.
-    ///
-    /// This is an internal function as callers always want to save results.
-    fn run_no_save(&self) -> Result<Solution<'model>> {
+    pub fn run(&self, run_description: &str, writer: &mut DataWriter) -> Result<Solution<'model>> {
         // If the user provided no markets to balance, we use all of them
         let all_markets: Vec<_>;
         let markets_to_balance = if self.markets_to_balance.is_empty() {
@@ -484,11 +475,31 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
         // 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 markets to users
         match self.run_without_unmet_demand_variables(markets_to_balance, input_prices) {
-            Ok(solution) => Ok(solution),
+            Ok(solution) => {
+                // Normal successful run: write debug info and return
+                writer.write_dispatch_debug_info(self.year, run_description, &solution)?;
+                Ok(solution)
+            }
             Err(ModelError::NonOptimal(HighsModelStatus::Infeasible)) => {
-                let markets = self
-                    .get_markets_with_unmet_demand(markets_to_balance, input_prices)
-                    .expect("Failed to run dispatch to calculate unmet demand");
+                // Re-run including unmet demand variables so we can record detailed unmet-demand
+                // debug output before returning an error to the caller.
+                let solution = self.run_internal(
+                    markets_to_balance,
+                    /*allow_unmet_demand=*/ true,
+                    input_prices,
+                )?;
+
+                // Write debug CSVs to help diagnosis
+                writer.write_dispatch_debug_info(self.year, run_description, &solution)?;
+
+                // Collect markets with unmet demand from the solution
+                let markets: IndexSet<_> = solution
+                    .iter_unmet_demand()
+                    .filter(|(_, _, _, flow)| *flow > Flow(0.0))
+                    .map(|(commodity_id, region_id, _, _)| {
+                        (commodity_id.clone(), region_id.clone())
+                    })
+                    .collect();
 
                 ensure!(
                     !markets.is_empty(),
@@ -519,27 +530,6 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
         )
     }
 
-    /// Run dispatch to diagnose which markets have unmet demand
-    fn get_markets_with_unmet_demand(
-        &self,
-        markets_to_balance: &[(CommodityID, RegionID)],
-        input_prices: Option<&CommodityPrices>,
-    ) -> Result<IndexSet<(CommodityID, RegionID)>> {
-        // Run dispatch including unmet demand variables for all markets being balanced
-        let solution = self.run_internal(
-            markets_to_balance,
-            /*allow_unmet_demand=*/ true,
-            input_prices,
-        )?;
-
-        // 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 to balance the specified markets, optionally including unmet demand variables
     fn run_internal(
         &self,