muse2/
output.rs

1//! The module responsible for writing output data to disk.
2use crate::agent::AgentID;
3use crate::asset::{Asset, AssetID, AssetRef};
4use crate::commodity::CommodityID;
5use crate::process::ProcessID;
6use crate::region::RegionID;
7use crate::simulation::investment::appraisal::AppraisalOutput;
8use crate::simulation::optimisation::{FlowMap, Solution};
9use crate::simulation::prices::PriceMap;
10use crate::time_slice::TimeSliceID;
11use crate::units::{Activity, Capacity, Flow, Money, MoneyPerActivity, MoneyPerFlow};
12use anyhow::{Context, Result, ensure};
13use csv;
14use indexmap::IndexMap;
15use serde::{Deserialize, Serialize};
16use std::fs;
17use std::fs::File;
18use std::path::{Path, PathBuf};
19
20pub mod metadata;
21use metadata::write_metadata;
22
23/// The output file name for commodity flows
24const COMMODITY_FLOWS_FILE_NAME: &str = "commodity_flows.csv";
25
26/// The output file name for commodity prices
27const COMMODITY_PRICES_FILE_NAME: &str = "commodity_prices.csv";
28
29/// The output file name for assets
30const ASSETS_FILE_NAME: &str = "assets.csv";
31
32/// The output file name for asset capacities
33const ASSET_CAPACITIES_FILE_NAME: &str = "asset_capacities.csv";
34
35/// Debug output file for asset dispatch
36const ACTIVITY_ASSET_DISPATCH: &str = "debug_dispatch_assets.csv";
37
38/// The output file name for commodity balance duals
39const COMMODITY_BALANCE_DUALS_FILE_NAME: &str = "debug_commodity_balance_duals.csv";
40
41/// The output file name for unmet demand values
42const UNMET_DEMAND_FILE_NAME: &str = "debug_unmet_demand.csv";
43
44/// The output file name for extra solver output values
45const SOLVER_VALUES_FILE_NAME: &str = "debug_solver.csv";
46
47/// The output file name for appraisal results
48const APPRAISAL_RESULTS_FILE_NAME: &str = "debug_appraisal_results.csv";
49
50/// The output file name for appraisal time slice results
51const APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME: &str = "debug_appraisal_results_time_slices.csv";
52
53/// Get the default output directory for the model
54pub fn get_output_dir(model_dir: &Path, results_root: PathBuf) -> Result<PathBuf> {
55    // Get the model name from the dir path. This ends up being convoluted because we need to check
56    // for all possible errors. Ugh.
57    let model_dir = model_dir
58        .canonicalize() // canonicalise in case the user has specified "."
59        .context("Could not resolve path to model")?;
60
61    let model_name = model_dir
62        .file_name()
63        .context("Model cannot be in root folder")?
64        .to_str()
65        .context("Invalid chars in model dir name")?;
66
67    // Construct path
68    Ok([results_root, model_name.into()].iter().collect())
69}
70
71/// Get the default output directory for commodity flow graphs for the model
72pub fn get_graphs_dir(model_dir: &Path, graph_results_root: PathBuf) -> Result<PathBuf> {
73    let model_dir = model_dir
74        .canonicalize() // canonicalise in case the user has specified "."
75        .context("Could not resolve path to model")?;
76    let model_name = model_dir
77        .file_name()
78        .context("Model cannot be in root folder")?
79        .to_str()
80        .context("Invalid chars in model dir name")?;
81    Ok([graph_results_root, model_name.into()].iter().collect())
82}
83
84/// Create a new output directory for the model, optionally overwriting existing data
85///
86/// # Arguments
87///
88/// * `output_dir` - The output directory to create/overwrite
89/// * `allow_overwrite` - Whether to delete and recreate the folder if it is non-empty
90///
91/// # Returns
92///
93/// True if the output dir contained existing data that was deleted, false if not, or an error.
94pub fn create_output_directory(output_dir: &Path, allow_overwrite: bool) -> Result<bool> {
95    // If the folder already exists, then delete it
96    let overwrite = if let Ok(mut it) = fs::read_dir(output_dir) {
97        if it.next().is_none() {
98            // Folder exists and is empty: nothing to do
99            return Ok(false);
100        }
101
102        ensure!(
103            allow_overwrite,
104            "Output folder already exists and is not empty. \
105            Please delete the folder or pass the --overwrite command-line option."
106        );
107
108        fs::remove_dir_all(output_dir).context("Could not delete folder")?;
109        true
110    } else {
111        false
112    };
113
114    // Try to create the directory, with parents
115    fs::create_dir_all(output_dir)?;
116
117    Ok(overwrite)
118}
119
120/// Copy input files to output directory
121pub fn copy_input_files(model_dir: &Path, output_dir: &Path, model_name: &str) -> Result<()> {
122    // Get the model name from the dir path.
123    let mut input_copy_dir = output_dir.to_path_buf();
124    input_copy_dir.extend(["input", model_name]);
125
126    fs::create_dir_all(&input_copy_dir).context("Could not create input copy directory")?;
127
128    for entry in fs::read_dir(model_dir)? {
129        let entry = entry?;
130        let path = entry.path();
131        if path.is_file() {
132            let file_name = path.file_name().unwrap();
133            fs::copy(&path, input_copy_dir.join(file_name))?;
134        }
135    }
136    Ok(())
137}
138
139/// Represents a row in the assets output CSV file.
140#[derive(Serialize, Deserialize, Debug, PartialEq)]
141struct AssetRow {
142    asset_id: AssetID,
143    process_id: ProcessID,
144    region_id: RegionID,
145    agent_id: AgentID,
146    commission_year: u32,
147}
148
149impl AssetRow {
150    /// Create a new [`AssetRow`] for the given asset
151    fn new(asset: &Asset) -> Self {
152        Self {
153            asset_id: asset.id().unwrap(),
154            process_id: asset.process_id().clone(),
155            region_id: asset.region_id().clone(),
156            agent_id: asset.agent_id().unwrap().clone(),
157            commission_year: asset.commission_year(),
158        }
159    }
160}
161
162/// Represents a row in the asset capacities output CSV file.
163#[derive(Serialize, Deserialize, Debug, PartialEq)]
164struct AssetCapacityRow {
165    milestone_year: u32,
166    asset_id: AssetID,
167    capacity: Capacity,
168    num_units: Option<u32>,
169}
170
171/// Represents the flow-related data in a row of the commodity flows CSV file.
172#[derive(Serialize, Deserialize, Debug, PartialEq)]
173struct CommodityFlowRow {
174    milestone_year: u32,
175    asset_id: AssetID,
176    commodity_id: CommodityID,
177    time_slice: TimeSliceID,
178    flow: Flow,
179}
180
181/// Represents a row in the commodity prices CSV file
182#[derive(Serialize, Deserialize, Debug, PartialEq)]
183struct CommodityPriceRow {
184    milestone_year: u32,
185    commodity_id: CommodityID,
186    region_id: RegionID,
187    time_slice: TimeSliceID,
188    price: MoneyPerFlow,
189}
190
191/// Represents the activity in a row of the dispatch CSV file
192#[derive(Serialize, Deserialize, Debug, PartialEq)]
193struct DispatchRow {
194    milestone_year: u32,
195    run_description: String,
196    asset_id: Option<AssetID>,
197    process_id: ProcessID,
198    region_id: RegionID,
199    time_slice: TimeSliceID,
200    activity: Option<Activity>,
201    activity_dual: Option<MoneyPerActivity>,
202    column_dual: Option<MoneyPerActivity>,
203}
204
205/// Represents the commodity balance duals data in a row of the commodity balance duals CSV file
206#[derive(Serialize, Deserialize, Debug, PartialEq)]
207struct CommodityBalanceDualsRow {
208    milestone_year: u32,
209    run_description: String,
210    commodity_id: CommodityID,
211    region_id: RegionID,
212    time_slice: TimeSliceID,
213    value: MoneyPerFlow,
214}
215
216/// Represents the unmet demand data in a row of the unmet demand CSV file
217#[derive(Serialize, Deserialize, Debug, PartialEq)]
218struct UnmetDemandRow {
219    milestone_year: u32,
220    run_description: String,
221    commodity_id: CommodityID,
222    region_id: RegionID,
223    time_slice: TimeSliceID,
224    value: Flow,
225}
226
227/// Represents solver output values
228#[derive(Serialize, Deserialize, Debug, PartialEq)]
229struct SolverValuesRow {
230    milestone_year: u32,
231    run_description: String,
232    objective_value: Money,
233}
234
235/// Represents the appraisal results in a row of the appraisal results CSV file
236#[derive(Serialize, Deserialize, Debug, PartialEq)]
237struct AppraisalResultsRow {
238    milestone_year: u32,
239    run_description: String,
240    asset_id: Option<AssetID>,
241    process_id: ProcessID,
242    region_id: RegionID,
243    capacity: Capacity,
244    metric: Option<f64>,
245}
246
247/// Represents the appraisal results in a row of the appraisal results CSV file
248#[derive(Serialize, Deserialize, Debug, PartialEq)]
249struct AppraisalResultsTimeSliceRow {
250    milestone_year: u32,
251    run_description: String,
252    asset_id: Option<AssetID>,
253    process_id: ProcessID,
254    region_id: RegionID,
255    time_slice: TimeSliceID,
256    activity: Activity,
257    activity_coefficient: MoneyPerActivity,
258    demand: Flow,
259    unmet_demand: Flow,
260}
261
262/// For writing extra debug information about the model
263struct DebugDataWriter {
264    context: Option<String>,
265    unmet_demand_file_path: PathBuf,
266    commodity_balance_duals_writer: csv::Writer<File>,
267    unmet_demand_writer: Option<csv::Writer<File>>,
268    solver_values_writer: csv::Writer<File>,
269    appraisal_results_writer: csv::Writer<File>,
270    appraisal_results_time_slice_writer: csv::Writer<File>,
271    dispatch_asset_writer: csv::Writer<File>,
272}
273
274impl DebugDataWriter {
275    /// Open CSV files to write debug info to
276    ///
277    /// # Arguments
278    ///
279    /// * `output_path` - Folder where files will be saved
280    fn create(output_path: &Path) -> Result<Self> {
281        let new_writer = |file_name| {
282            let file_path = output_path.join(file_name);
283            csv::Writer::from_path(file_path)
284        };
285
286        Ok(Self {
287            context: None,
288            unmet_demand_file_path: output_path.join(UNMET_DEMAND_FILE_NAME),
289            commodity_balance_duals_writer: new_writer(COMMODITY_BALANCE_DUALS_FILE_NAME)?,
290            unmet_demand_writer: None,
291            solver_values_writer: new_writer(SOLVER_VALUES_FILE_NAME)?,
292            appraisal_results_writer: new_writer(APPRAISAL_RESULTS_FILE_NAME)?,
293            appraisal_results_time_slice_writer: new_writer(
294                APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME,
295            )?,
296            dispatch_asset_writer: new_writer(ACTIVITY_ASSET_DISPATCH)?,
297        })
298    }
299
300    /// Prepend the current context to the run description
301    fn with_context(&self, run_description: &str) -> String {
302        if let Some(context) = &self.context {
303            format!("{context}; {run_description}")
304        } else {
305            run_description.to_string()
306        }
307    }
308
309    /// Write debug info about the dispatch optimisation
310    fn write_dispatch_debug_info(
311        &mut self,
312        milestone_year: u32,
313        run_description: &str,
314        solution: &Solution,
315    ) -> Result<()> {
316        self.write_dispatch(
317            milestone_year,
318            run_description,
319            solution.iter_activity(),
320            solution.iter_activity_duals(),
321            solution.iter_column_duals(),
322        )?;
323        self.write_commodity_balance_duals(
324            milestone_year,
325            run_description,
326            solution.iter_commodity_balance_duals(),
327        )?;
328        self.write_unmet_demand(
329            milestone_year,
330            run_description,
331            solution.iter_unmet_demand(),
332        )?;
333        self.write_solver_values(milestone_year, run_description, solution.objective_value)?;
334        Ok(())
335    }
336
337    // Write activity to file
338    fn write_dispatch<'a, I, J, K>(
339        &mut self,
340        milestone_year: u32,
341        run_description: &str,
342        iter_activity: I,
343        iter_activity_duals: J,
344        iter_column_duals: K,
345    ) -> Result<()>
346    where
347        I: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, Activity)>,
348        J: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, MoneyPerActivity)>,
349        K: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, MoneyPerActivity)>,
350    {
351        // To account for different order of entries or missing ones, we first compile data in hash map
352        type CompiledActivityData = (
353            Option<Activity>,
354            Option<MoneyPerActivity>,
355            Option<MoneyPerActivity>,
356        );
357        let mut map: IndexMap<(&AssetRef, &TimeSliceID), CompiledActivityData> = IndexMap::new();
358
359        // For the activities
360        for (asset, time_slice, activity) in iter_activity {
361            map.entry((asset, time_slice)).or_default().0 = Some(activity);
362        }
363        // The activity duals
364        for (asset, time_slice, activity_dual) in iter_activity_duals {
365            map.entry((asset, time_slice)).or_default().1 = Some(activity_dual);
366        }
367        // And the column duals
368        for (asset, time_slice, column_dual) in iter_column_duals {
369            map.entry((asset, time_slice)).or_default().2 = Some(column_dual);
370        }
371
372        for ((asset, time_slice), (activity, activity_dual, column_dual)) in map {
373            let row = DispatchRow {
374                milestone_year,
375                run_description: self.with_context(run_description),
376                asset_id: asset.id(),
377                process_id: asset.process_id().clone(),
378                region_id: asset.region_id().clone(),
379                time_slice: time_slice.clone(),
380                activity,
381                activity_dual,
382                column_dual,
383            };
384            self.dispatch_asset_writer.serialize(row)?;
385        }
386
387        Ok(())
388    }
389
390    /// Write commodity balance duals to file
391    fn write_commodity_balance_duals<'a, I>(
392        &mut self,
393        milestone_year: u32,
394        run_description: &str,
395        iter: I,
396    ) -> Result<()>
397    where
398        I: Iterator<Item = (&'a CommodityID, &'a RegionID, &'a TimeSliceID, MoneyPerFlow)>,
399    {
400        for (commodity_id, region_id, time_slice, value) in iter {
401            let row = CommodityBalanceDualsRow {
402                milestone_year,
403                run_description: self.with_context(run_description),
404                commodity_id: commodity_id.clone(),
405                region_id: region_id.clone(),
406                time_slice: time_slice.clone(),
407                value,
408            };
409            self.commodity_balance_duals_writer.serialize(row)?;
410        }
411
412        Ok(())
413    }
414
415    /// Write unmet demand values to file
416    fn write_unmet_demand<'a, I>(
417        &mut self,
418        milestone_year: u32,
419        run_description: &str,
420        iter: I,
421    ) -> Result<()>
422    where
423        I: Iterator<Item = (&'a CommodityID, &'a RegionID, &'a TimeSliceID, Flow)>,
424    {
425        let mut rows = iter.peekable();
426
427        if rows.peek().is_none() {
428            return Ok(());
429        }
430
431        // If the unmet demand writer already exist, we panic, as it should not happen
432        assert!(
433            self.unmet_demand_writer.is_none(),
434            "Unmet demand file already exists!"
435        );
436
437        let run_description = self.with_context(run_description);
438        let writer = self
439            .unmet_demand_writer
440            .insert(csv::Writer::from_path(&self.unmet_demand_file_path)?);
441        for (commodity_id, region_id, time_slice, value) in rows {
442            let row = UnmetDemandRow {
443                milestone_year,
444                run_description: run_description.clone(),
445                commodity_id: commodity_id.clone(),
446                region_id: region_id.clone(),
447                time_slice: time_slice.clone(),
448                value,
449            };
450            writer.serialize(row)?;
451        }
452
453        Ok(())
454    }
455
456    /// Write additional solver output values to file
457    fn write_solver_values(
458        &mut self,
459        milestone_year: u32,
460        run_description: &str,
461        objective_value: Money,
462    ) -> Result<()> {
463        let row = SolverValuesRow {
464            milestone_year,
465            run_description: self.with_context(run_description),
466            objective_value,
467        };
468        self.solver_values_writer.serialize(row)?;
469        self.solver_values_writer.flush()?;
470
471        Ok(())
472    }
473
474    /// Write appraisal results to file
475    fn write_appraisal_results(
476        &mut self,
477        milestone_year: u32,
478        run_description: &str,
479        appraisal_results: &[AppraisalOutput],
480    ) -> Result<()> {
481        for result in appraisal_results {
482            let row = AppraisalResultsRow {
483                milestone_year,
484                run_description: self.with_context(run_description),
485                asset_id: result.asset.id(),
486                process_id: result.asset.process_id().clone(),
487                region_id: result.asset.region_id().clone(),
488                capacity: result.asset.capacity().total_capacity(),
489                metric: result.metric.as_ref().map(|m| m.value()),
490            };
491            self.appraisal_results_writer.serialize(row)?;
492        }
493
494        Ok(())
495    }
496
497    /// Write appraisal results to file
498    fn write_appraisal_time_slice_results(
499        &mut self,
500        milestone_year: u32,
501        run_description: &str,
502        appraisal_results: &[AppraisalOutput],
503        demand: &IndexMap<TimeSliceID, Flow>,
504    ) -> Result<()> {
505        for result in appraisal_results {
506            for (time_slice, activity) in &result.activity {
507                let activity_coefficient = result.coefficients.activity_coefficients[time_slice];
508                let demand = demand[time_slice];
509                let unmet_demand = result.unmet_demand[time_slice];
510                let row = AppraisalResultsTimeSliceRow {
511                    milestone_year,
512                    run_description: self.with_context(run_description),
513                    asset_id: result.asset.id(),
514                    process_id: result.asset.process_id().clone(),
515                    region_id: result.asset.region_id().clone(),
516                    time_slice: time_slice.clone(),
517                    activity: *activity,
518                    activity_coefficient,
519                    demand,
520                    unmet_demand,
521                };
522                self.appraisal_results_time_slice_writer.serialize(row)?;
523            }
524        }
525
526        Ok(())
527    }
528
529    /// Flush the underlying streams
530    fn flush(&mut self) -> Result<()> {
531        if let Some(wrt) = &mut self.unmet_demand_writer {
532            wrt.flush()?;
533        }
534        self.commodity_balance_duals_writer.flush()?;
535        self.solver_values_writer.flush()?;
536        self.appraisal_results_writer.flush()?;
537        self.appraisal_results_time_slice_writer.flush()?;
538        self.dispatch_asset_writer.flush()?;
539
540        Ok(())
541    }
542}
543
544/// An object for writing output data to file
545pub struct DataWriter {
546    assets: csv::Writer<File>,
547    asset_capacities: csv::Writer<File>,
548    flows: csv::Writer<File>,
549    prices: csv::Writer<File>,
550    debug: Option<DebugDataWriter>,
551}
552
553impl DataWriter {
554    /// Open CSV files to write output data to
555    ///
556    /// # Arguments
557    ///
558    /// * `output_path` - Folder where files will be saved
559    /// * `model_path` - Path to input model
560    /// * `save_debug_info` - Whether to include extra CSV files for debugging model
561    pub fn create(output_path: &Path, model_path: &Path, save_debug_info: bool) -> Result<Self> {
562        write_metadata(output_path, model_path).context("Failed to save metadata")?;
563
564        let new_writer = |file_name| {
565            let file_path = output_path.join(file_name);
566            csv::Writer::from_path(file_path)
567        };
568
569        let debug_writer = if save_debug_info {
570            // Create debug CSV files
571            Some(DebugDataWriter::create(output_path)?)
572        } else {
573            None
574        };
575
576        Ok(Self {
577            assets: new_writer(ASSETS_FILE_NAME)?,
578            asset_capacities: new_writer(ASSET_CAPACITIES_FILE_NAME)?,
579            flows: new_writer(COMMODITY_FLOWS_FILE_NAME)?,
580            prices: new_writer(COMMODITY_PRICES_FILE_NAME)?,
581            debug: debug_writer,
582        })
583    }
584
585    /// Write debug info about the dispatch optimisation
586    pub fn write_dispatch_debug_info(
587        &mut self,
588        milestone_year: u32,
589        run_description: &str,
590        solution: &Solution,
591    ) -> Result<()> {
592        if let Some(wtr) = &mut self.debug {
593            wtr.write_dispatch_debug_info(milestone_year, run_description, solution)?;
594        }
595
596        Ok(())
597    }
598
599    /// Write debug info about the investment appraisal
600    pub fn write_appraisal_debug_info(
601        &mut self,
602        milestone_year: u32,
603        run_description: &str,
604        appraisal_results: &[AppraisalOutput],
605        demand: &IndexMap<TimeSliceID, Flow>,
606    ) -> Result<()> {
607        if let Some(wtr) = &mut self.debug {
608            wtr.write_appraisal_results(milestone_year, run_description, appraisal_results)?;
609            wtr.write_appraisal_time_slice_results(
610                milestone_year,
611                run_description,
612                appraisal_results,
613                demand,
614            )?;
615        }
616
617        Ok(())
618    }
619
620    /// Append newly commissioned asset definitions to the assets CSV file
621    pub fn write_assets(&mut self, assets: &[AssetRef]) -> Result<()> {
622        for asset in assets {
623            self.assets.serialize(AssetRow::new(asset))?;
624        }
625
626        Ok(())
627    }
628
629    /// Write asset capacities for the current milestone year to a CSV file
630    pub fn write_asset_capacities(
631        &mut self,
632        milestone_year: u32,
633        assets: &[AssetRef],
634    ) -> Result<()> {
635        for asset in assets {
636            let row = AssetCapacityRow {
637                milestone_year,
638                asset_id: asset.id().unwrap(),
639                capacity: asset.total_capacity(),
640                num_units: asset.capacity().n_units(),
641            };
642            self.asset_capacities.serialize(row)?;
643        }
644
645        Ok(())
646    }
647
648    /// Write commodity flows to a CSV file
649    pub fn write_flows(&mut self, milestone_year: u32, flow_map: &FlowMap) -> Result<()> {
650        for ((asset, commodity_id, time_slice), flow) in flow_map {
651            let row = CommodityFlowRow {
652                milestone_year,
653                asset_id: asset.id().unwrap(),
654                commodity_id: commodity_id.clone(),
655                time_slice: time_slice.clone(),
656                flow: *flow,
657            };
658            self.flows.serialize(row)?;
659        }
660
661        Ok(())
662    }
663
664    /// Write commodity prices to a CSV file
665    pub fn write_prices(&mut self, milestone_year: u32, prices: &PriceMap) -> Result<()> {
666        for (commodity_id, region_id, time_slice, price) in prices.iter() {
667            let row = CommodityPriceRow {
668                milestone_year,
669                commodity_id: commodity_id.clone(),
670                region_id: region_id.clone(),
671                time_slice: time_slice.clone(),
672                price,
673            };
674            self.prices.serialize(row)?;
675        }
676
677        Ok(())
678    }
679
680    /// Flush the underlying streams
681    pub fn flush(&mut self) -> Result<()> {
682        self.assets.flush()?;
683        self.asset_capacities.flush()?;
684        self.flows.flush()?;
685        self.prices.flush()?;
686        if let Some(wtr) = &mut self.debug {
687            wtr.flush()?;
688        }
689
690        Ok(())
691    }
692
693    /// Add context to the debug writer
694    pub fn set_debug_context(&mut self, context: String) {
695        if let Some(wtr) = &mut self.debug {
696            wtr.context = Some(context);
697        }
698    }
699
700    /// Clear context from the debug writer
701    pub fn clear_debug_context(&mut self) {
702        if let Some(wtr) = &mut self.debug {
703            wtr.context = None;
704        }
705    }
706}
707
708#[cfg(test)]
709mod tests {
710    use super::*;
711    use crate::asset::AssetPool;
712    use crate::fixture::{appraisal_output, asset, assets, commodity_id, region_id, time_slice};
713    use crate::simulation::investment::appraisal::AppraisalOutput;
714    use crate::time_slice::TimeSliceID;
715    use indexmap::indexmap;
716    use itertools::{Itertools, assert_equal};
717    use rstest::rstest;
718    use std::iter;
719    use tempfile::tempdir;
720
721    #[rstest]
722    fn write_assets(assets: AssetPool) {
723        let dir = tempdir().unwrap();
724
725        // Write an asset
726        {
727            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
728            writer.write_assets(&assets).unwrap();
729            writer.flush().unwrap();
730        }
731
732        // Read back and compare
733        let asset = assets.iter().next().unwrap();
734        let expected = AssetRow::new(asset);
735        let records: Vec<AssetRow> = csv::Reader::from_path(dir.path().join(ASSETS_FILE_NAME))
736            .unwrap()
737            .into_deserialize()
738            .try_collect()
739            .unwrap();
740        assert_equal(records, iter::once(expected));
741    }
742
743    #[rstest]
744    fn write_asset_capacities(assets: AssetPool) {
745        let milestone_year = 2020;
746        let dir = tempdir().unwrap();
747
748        // Write asset capacities
749        {
750            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
751            writer
752                .write_asset_capacities(milestone_year, &assets)
753                .unwrap();
754            writer.flush().unwrap();
755        }
756
757        // Read back and compare
758        let asset = assets.iter().next().unwrap();
759        let expected = AssetCapacityRow {
760            milestone_year,
761            asset_id: asset.id().unwrap(),
762            capacity: asset.total_capacity(),
763            num_units: None,
764        };
765        let records: Vec<AssetCapacityRow> =
766            csv::Reader::from_path(dir.path().join(ASSET_CAPACITIES_FILE_NAME))
767                .unwrap()
768                .into_deserialize()
769                .try_collect()
770                .unwrap();
771        assert_equal(records, iter::once(expected));
772    }
773
774    #[rstest]
775    fn write_flows(assets: AssetPool, commodity_id: CommodityID, time_slice: TimeSliceID) {
776        let milestone_year = 2020;
777        let asset = assets.iter().next().unwrap();
778        let flow_map = indexmap! {
779            (asset.clone(), commodity_id.clone(), time_slice.clone()) => Flow(42.0)
780        };
781
782        // Write a flow
783        let dir = tempdir().unwrap();
784        {
785            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
786            writer.write_flows(milestone_year, &flow_map).unwrap();
787            writer.flush().unwrap();
788        }
789
790        // Read back and compare
791        let expected = CommodityFlowRow {
792            milestone_year,
793            asset_id: asset.id().unwrap(),
794            commodity_id,
795            time_slice,
796            flow: Flow(42.0),
797        };
798        let records: Vec<CommodityFlowRow> =
799            csv::Reader::from_path(dir.path().join(COMMODITY_FLOWS_FILE_NAME))
800                .unwrap()
801                .into_deserialize()
802                .try_collect()
803                .unwrap();
804        assert_equal(records, iter::once(expected));
805    }
806
807    #[rstest]
808    fn write_prices(commodity_id: CommodityID, region_id: RegionID, time_slice: TimeSliceID) {
809        let milestone_year = 2020;
810        let price = MoneyPerFlow(42.0);
811        let mut prices = PriceMap::default();
812        prices.insert(&commodity_id, &region_id, &time_slice, price);
813
814        let dir = tempdir().unwrap();
815
816        // Write a price
817        {
818            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
819            writer.write_prices(milestone_year, &prices).unwrap();
820            writer.flush().unwrap();
821        }
822
823        // Read back and compare
824        let expected = CommodityPriceRow {
825            milestone_year,
826            commodity_id,
827            region_id,
828            time_slice,
829            price,
830        };
831        let records: Vec<CommodityPriceRow> =
832            csv::Reader::from_path(dir.path().join(COMMODITY_PRICES_FILE_NAME))
833                .unwrap()
834                .into_deserialize()
835                .try_collect()
836                .unwrap();
837        assert_equal(records, iter::once(expected));
838    }
839
840    #[rstest]
841    fn write_commodity_balance_duals(
842        commodity_id: CommodityID,
843        region_id: RegionID,
844        time_slice: TimeSliceID,
845    ) {
846        let milestone_year = 2020;
847        let run_description = "test_run".to_string();
848        let value = MoneyPerFlow(0.5);
849        let dir = tempdir().unwrap();
850
851        // Write commodity balance dual
852        {
853            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
854            writer
855                .write_commodity_balance_duals(
856                    milestone_year,
857                    &run_description,
858                    iter::once((&commodity_id, &region_id, &time_slice, value)),
859                )
860                .unwrap();
861            writer.flush().unwrap();
862        }
863
864        // Read back and compare
865        let expected = CommodityBalanceDualsRow {
866            milestone_year,
867            run_description,
868            commodity_id,
869            region_id,
870            time_slice,
871            value,
872        };
873        let records: Vec<CommodityBalanceDualsRow> =
874            csv::Reader::from_path(dir.path().join(COMMODITY_BALANCE_DUALS_FILE_NAME))
875                .unwrap()
876                .into_deserialize()
877                .try_collect()
878                .unwrap();
879        assert_equal(records, iter::once(expected));
880    }
881
882    #[rstest]
883    fn write_unmet_demand(commodity_id: CommodityID, region_id: RegionID, time_slice: TimeSliceID) {
884        let milestone_year = 2020;
885        let run_description = "test_run".to_string();
886        let value = Flow(0.5);
887        let dir = tempdir().unwrap();
888
889        // Write unmet demand
890        {
891            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
892            writer
893                .write_unmet_demand(
894                    milestone_year,
895                    &run_description,
896                    iter::once((&commodity_id, &region_id, &time_slice, value)),
897                )
898                .unwrap();
899            writer.flush().unwrap();
900        }
901
902        // Read back and compare
903        let expected = UnmetDemandRow {
904            milestone_year,
905            run_description,
906            commodity_id,
907            region_id,
908            time_slice,
909            value,
910        };
911        let records: Vec<UnmetDemandRow> =
912            csv::Reader::from_path(dir.path().join(UNMET_DEMAND_FILE_NAME))
913                .unwrap()
914                .into_deserialize()
915                .try_collect()
916                .unwrap();
917        assert_equal(records, iter::once(expected));
918    }
919
920    #[rstest]
921    fn write_dispatch(assets: AssetPool, time_slice: TimeSliceID) {
922        let milestone_year = 2020;
923        let run_description = "test_run".to_string();
924        let activity = Activity(100.5);
925        let activity_dual = MoneyPerActivity(-1.5);
926        let column_dual = MoneyPerActivity(5.0);
927        let dir = tempdir().unwrap();
928        let asset = assets.iter().next().unwrap();
929
930        // Write activity
931        {
932            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
933            writer
934                .write_dispatch(
935                    milestone_year,
936                    &run_description,
937                    iter::once((asset, &time_slice, activity)),
938                    iter::once((asset, &time_slice, activity_dual)),
939                    iter::once((asset, &time_slice, column_dual)),
940                )
941                .unwrap();
942            writer.flush().unwrap();
943        }
944
945        // Read back and compare
946        let expected = DispatchRow {
947            milestone_year,
948            run_description,
949            asset_id: asset.id(),
950            process_id: asset.process_id().clone(),
951            region_id: asset.region_id().clone(),
952            time_slice,
953            activity: Some(activity),
954            activity_dual: Some(activity_dual),
955            column_dual: Some(column_dual),
956        };
957        let records: Vec<DispatchRow> =
958            csv::Reader::from_path(dir.path().join(ACTIVITY_ASSET_DISPATCH))
959                .unwrap()
960                .into_deserialize()
961                .try_collect()
962                .unwrap();
963        assert_equal(records, iter::once(expected));
964    }
965
966    #[rstest]
967    fn write_dispatch_with_missing_keys(assets: AssetPool, time_slice: TimeSliceID) {
968        let milestone_year = 2020;
969        let run_description = "test_run".to_string();
970        let activity = Activity(100.5);
971        let dir = tempdir().unwrap();
972        let asset = assets.iter().next().unwrap();
973
974        // Write activity
975        {
976            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
977            writer
978                .write_dispatch(
979                    milestone_year,
980                    &run_description,
981                    iter::once((asset, &time_slice, activity)),
982                    iter::empty::<(&AssetRef, &TimeSliceID, MoneyPerActivity)>(),
983                    iter::empty::<(&AssetRef, &TimeSliceID, MoneyPerActivity)>(),
984                )
985                .unwrap();
986            writer.flush().unwrap();
987        }
988
989        // Read back and compare
990        let expected = DispatchRow {
991            milestone_year,
992            run_description,
993            asset_id: asset.id(),
994            process_id: asset.process_id().clone(),
995            region_id: asset.region_id().clone(),
996            time_slice,
997            activity: Some(activity),
998            activity_dual: None,
999            column_dual: None,
1000        };
1001        let records: Vec<DispatchRow> =
1002            csv::Reader::from_path(dir.path().join(ACTIVITY_ASSET_DISPATCH))
1003                .unwrap()
1004                .into_deserialize()
1005                .try_collect()
1006                .unwrap();
1007        assert_equal(records, iter::once(expected));
1008    }
1009
1010    #[rstest]
1011    fn write_solver_values() {
1012        let milestone_year = 2020;
1013        let run_description = "test_run".to_string();
1014        let objective_value = Money(1234.56);
1015        let dir = tempdir().unwrap();
1016
1017        // Write solver values
1018        {
1019            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
1020            writer
1021                .write_solver_values(milestone_year, &run_description, objective_value)
1022                .unwrap();
1023            writer.flush().unwrap();
1024        }
1025
1026        // Read back and compare
1027        let expected = SolverValuesRow {
1028            milestone_year,
1029            run_description,
1030            objective_value,
1031        };
1032        let records: Vec<SolverValuesRow> =
1033            csv::Reader::from_path(dir.path().join(SOLVER_VALUES_FILE_NAME))
1034                .unwrap()
1035                .into_deserialize()
1036                .try_collect()
1037                .unwrap();
1038        assert_equal(records, iter::once(expected));
1039    }
1040
1041    #[rstest]
1042    fn write_appraisal_results(asset: Asset, appraisal_output: AppraisalOutput) {
1043        let milestone_year = 2020;
1044        let run_description = "test_run".to_string();
1045        let dir = tempdir().unwrap();
1046
1047        // Write appraisal results
1048        {
1049            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
1050            writer
1051                .write_appraisal_results(milestone_year, &run_description, &[appraisal_output])
1052                .unwrap();
1053            writer.flush().unwrap();
1054        }
1055
1056        // Read back and compare
1057        let expected = AppraisalResultsRow {
1058            milestone_year,
1059            run_description,
1060            asset_id: None,
1061            process_id: asset.process_id().clone(),
1062            region_id: asset.region_id().clone(),
1063            capacity: Capacity(2.0),
1064            metric: Some(4.14),
1065        };
1066        let records: Vec<AppraisalResultsRow> =
1067            csv::Reader::from_path(dir.path().join(APPRAISAL_RESULTS_FILE_NAME))
1068                .unwrap()
1069                .into_deserialize()
1070                .try_collect()
1071                .unwrap();
1072        assert_equal(records, iter::once(expected));
1073    }
1074
1075    #[rstest]
1076    fn write_appraisal_time_slice_results(
1077        asset: Asset,
1078        appraisal_output: AppraisalOutput,
1079        time_slice: TimeSliceID,
1080    ) {
1081        let milestone_year = 2020;
1082        let run_description = "test_run".to_string();
1083        let dir = tempdir().unwrap();
1084        let demand = indexmap! {time_slice.clone() => Flow(100.0) };
1085
1086        // Write appraisal time slice results
1087        {
1088            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
1089            writer
1090                .write_appraisal_time_slice_results(
1091                    milestone_year,
1092                    &run_description,
1093                    &[appraisal_output],
1094                    &demand,
1095                )
1096                .unwrap();
1097            writer.flush().unwrap();
1098        }
1099
1100        // Read back and compare
1101        let expected = AppraisalResultsTimeSliceRow {
1102            milestone_year,
1103            run_description,
1104            asset_id: None,
1105            process_id: asset.process_id().clone(),
1106            region_id: asset.region_id().clone(),
1107            time_slice: time_slice.clone(),
1108            activity: Activity(10.0),
1109            activity_coefficient: MoneyPerActivity(0.5),
1110            demand: Flow(100.0),
1111            unmet_demand: Flow(5.0),
1112        };
1113        let records: Vec<AppraisalResultsTimeSliceRow> =
1114            csv::Reader::from_path(dir.path().join(APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME))
1115                .unwrap()
1116                .into_deserialize()
1117                .try_collect()
1118                .unwrap();
1119        assert_equal(records, iter::once(expected));
1120    }
1121
1122    #[test]
1123    fn create_output_directory_new_directory() {
1124        let temp_dir = tempdir().unwrap();
1125        let output_dir = temp_dir.path().join("new_output");
1126
1127        // Create a new directory should succeed and return false (no overwrite)
1128        let result = create_output_directory(&output_dir, false).unwrap();
1129        assert!(!result);
1130        assert!(output_dir.exists());
1131        assert!(output_dir.is_dir());
1132    }
1133
1134    #[test]
1135    fn create_output_directory_existing_empty_directory() {
1136        let temp_dir = tempdir().unwrap();
1137        let output_dir = temp_dir.path().join("empty_output");
1138
1139        // Create the directory first
1140        fs::create_dir(&output_dir).unwrap();
1141
1142        // Creating again should succeed and return false (no overwrite needed)
1143        let result = create_output_directory(&output_dir, false).unwrap();
1144        assert!(!result);
1145        assert!(output_dir.exists());
1146        assert!(output_dir.is_dir());
1147    }
1148
1149    #[test]
1150    fn create_output_directory_existing_with_files_no_overwrite() {
1151        let temp_dir = tempdir().unwrap();
1152        let output_dir = temp_dir.path().join("output_with_files");
1153
1154        // Create directory with a file
1155        fs::create_dir(&output_dir).unwrap();
1156        fs::write(output_dir.join("existing_file.txt"), "some content").unwrap();
1157
1158        // Should fail when allow_overwrite is false
1159        let result = create_output_directory(&output_dir, false);
1160        assert!(result.is_err());
1161        assert!(
1162            result
1163                .unwrap_err()
1164                .to_string()
1165                .contains("Output folder already exists")
1166        );
1167    }
1168
1169    #[test]
1170    fn create_output_directory_existing_with_files_allow_overwrite() {
1171        let temp_dir = tempdir().unwrap();
1172        let output_dir = temp_dir.path().join("output_with_files");
1173
1174        // Create directory with a file
1175        fs::create_dir(&output_dir).unwrap();
1176        let file_path = output_dir.join("existing_file.txt");
1177        fs::write(&file_path, "some content").unwrap();
1178
1179        // Should succeed when allow_overwrite is true and return true (overwrite occurred)
1180        let result = create_output_directory(&output_dir, true).unwrap();
1181        assert!(result);
1182        assert!(output_dir.exists());
1183        assert!(output_dir.is_dir());
1184        assert!(!file_path.exists()); // File should be gone
1185    }
1186
1187    #[test]
1188    fn create_output_directory_nested_path() {
1189        let temp_dir = tempdir().unwrap();
1190        let output_dir = temp_dir.path().join("nested").join("path").join("output");
1191
1192        // Should create nested directories and return false (no overwrite)
1193        let result = create_output_directory(&output_dir, false).unwrap();
1194        assert!(!result);
1195        assert!(output_dir.exists());
1196        assert!(output_dir.is_dir());
1197    }
1198
1199    #[test]
1200    fn create_output_directory_existing_subdirs_with_files_allow_overwrite() {
1201        let temp_dir = tempdir().unwrap();
1202        let output_dir = temp_dir.path().join("output_with_subdirs");
1203
1204        // Create directory structure with files
1205        fs::create_dir_all(output_dir.join("subdir")).unwrap();
1206        fs::write(output_dir.join("file1.txt"), "content1").unwrap();
1207        fs::write(output_dir.join("subdir").join("file2.txt"), "content2").unwrap();
1208
1209        // Should succeed when allow_overwrite is true and return true (overwrite occurred)
1210        let result = create_output_directory(&output_dir, true).unwrap();
1211        assert!(result);
1212        assert!(output_dir.exists());
1213        assert!(output_dir.is_dir());
1214        // All previous content should be gone
1215        assert!(!output_dir.join("file1.txt").exists());
1216        assert!(!output_dir.join("subdir").exists());
1217    }
1218}
muse2/output.rs

muse2/
output.rs
