Skip to main content

muse2/
input.rs

1//! Common routines for handling input data.
2use crate::graph::investment::solve_investment_order_for_model;
3use crate::graph::validate::validate_commodity_graphs_for_model;
4use crate::graph::{CommoditiesGraph, build_commodity_graphs_for_model};
5use crate::id::{HasID, ID};
6use crate::model::{Model, ModelParameters};
7use crate::region::RegionID;
8use crate::units::UnitType;
9use anyhow::{Context, Result, bail, ensure};
10use float_cmp::approx_eq;
11use indexmap::IndexMap;
12use itertools::Itertools;
13use serde::de::{Deserialize, DeserializeOwned, Deserializer};
14use std::collections::HashMap;
15use std::fmt::{self, Write};
16use std::fs;
17use std::hash::Hash;
18use std::path::Path;
19
20mod agent;
21use agent::read_agents;
22mod asset;
23use asset::read_user_assets;
24mod commodity;
25use commodity::read_commodities;
26mod process;
27use process::read_processes;
28mod region;
29use region::read_regions;
30mod time_slice;
31use time_slice::read_time_slice_info;
32mod range;
33use range::{parse_range, parse_range_parts};
34mod year;
35use year::parse_year_str;
36
37/// A trait which provides a method to insert a key and value into a map
38pub trait Insert<K, V> {
39    /// Insert a key and value into the map
40    fn insert(&mut self, key: K, value: V) -> Option<V>;
41}
42
43impl<K: Eq + Hash, V> Insert<K, V> for HashMap<K, V> {
44    fn insert(&mut self, key: K, value: V) -> Option<V> {
45        HashMap::insert(self, key, value)
46    }
47}
48
49impl<K: Eq + Hash, V> Insert<K, V> for IndexMap<K, V> {
50    fn insert(&mut self, key: K, value: V) -> Option<V> {
51        IndexMap::insert(self, key, value)
52    }
53}
54
55/// Read a series of type `T`s from a CSV file.
56///
57/// Returns an error if the file is empty.
58///
59/// # Arguments
60///
61/// * `file_path` - Path to the CSV file
62pub fn read_csv<'a, T: DeserializeOwned + 'a>(
63    file_path: &'a Path,
64) -> Result<impl Iterator<Item = T> + 'a> {
65    let vec = read_csv_internal(file_path)?;
66    if vec.is_empty() {
67        bail!("CSV file {} cannot be empty", file_path.display());
68    }
69    Ok(vec.into_iter())
70}
71
72/// Read a series of type `T`s from a CSV file.
73///
74/// # Arguments
75///
76/// * `file_path` - Path to the CSV file
77pub fn read_csv_optional<'a, T: DeserializeOwned + 'a>(
78    file_path: &'a Path,
79) -> Result<impl Iterator<Item = T> + 'a> {
80    if !file_path.exists() {
81        return Ok(Vec::new().into_iter());
82    }
83
84    let vec = read_csv_internal(file_path)?;
85    Ok(vec.into_iter())
86}
87
88fn read_csv_internal<'a, T: DeserializeOwned + 'a>(file_path: &'a Path) -> Result<Vec<T>> {
89    let vec = csv::ReaderBuilder::new()
90        .trim(csv::Trim::All)
91        .from_path(file_path)
92        .with_context(|| input_err_msg(file_path))?
93        .into_deserialize()
94        .process_results(|iter| iter.collect_vec())
95        .with_context(|| input_err_msg(file_path))?;
96
97    Ok(vec)
98}
99
100/// Parse a TOML file at the specified path.
101///
102/// # Arguments
103///
104/// * `file_path` - Path to the TOML file
105///
106/// # Returns
107///
108/// * The deserialised TOML data or an error if the file could not be read or parsed.
109pub fn read_toml<T: DeserializeOwned>(file_path: &Path) -> Result<T> {
110    let toml_str = fs::read_to_string(file_path).with_context(|| input_err_msg(file_path))?;
111    let toml_data = toml::from_str(&toml_str).with_context(|| input_err_msg(file_path))?;
112    Ok(toml_data)
113}
114
115/// Read a Dimensionless float, checking that it is between 0 and 1
116pub fn deserialise_proportion_nonzero<'de, D, T>(deserialiser: D) -> Result<T, D::Error>
117where
118    T: UnitType,
119    D: Deserializer<'de>,
120{
121    let value = f64::deserialize(deserialiser)?;
122    if !(value > 0.0 && value <= 1.0) {
123        Err(serde::de::Error::custom("Value must be > 0 and <= 1"))?;
124    }
125
126    Ok(T::new(value))
127}
128
129/// Read a value, checking that it is finite and greater than or equal to zero
130pub fn deserialise_finite_non_negative<'de, D, T>(deserialiser: D) -> Result<T, D::Error>
131where
132    T: UnitType,
133    D: Deserializer<'de>,
134{
135    let value = f64::deserialize(deserialiser)?;
136    if !(value.is_finite() && value >= 0.0) {
137        Err(serde::de::Error::custom(
138            "Value must be a finite number greater than or equal to zero",
139        ))?;
140    }
141
142    Ok(T::new(value))
143}
144
145/// Format an error message to include the file path. To be used with `anyhow::Context`.
146pub fn input_err_msg<P: AsRef<Path>>(file_path: P) -> String {
147    format!("Error reading {}", file_path.as_ref().display())
148}
149
150/// Read a CSV file of items with IDs.
151///
152/// As this function is only ever used for top-level CSV files (i.e. the ones which actually define
153/// the IDs for a given type), we use an ordered map to maintain the order in the input files.
154fn read_csv_id_file<T, I: ID>(file_path: &Path) -> Result<IndexMap<I, T>>
155where
156    T: HasID<I> + DeserializeOwned,
157{
158    fn fill_and_validate_map<T, I: ID>(file_path: &Path) -> Result<IndexMap<I, T>>
159    where
160        T: HasID<I> + DeserializeOwned,
161    {
162        let mut map = IndexMap::new();
163        for record in read_csv::<T>(file_path)? {
164            let id = record.get_id().clone();
165            let existing = map.insert(id.clone(), record).is_some();
166            ensure!(!existing, "Duplicate ID found: {id}");
167        }
168        ensure!(!map.is_empty(), "CSV file is empty");
169
170        Ok(map)
171    }
172
173    fill_and_validate_map(file_path).with_context(|| input_err_msg(file_path))
174}
175
176/// Check that fractions sum to (approximately) one
177fn check_values_sum_to_one_approx<I, T>(fractions: I) -> Result<()>
178where
179    T: UnitType,
180    I: Iterator<Item = T>,
181{
182    let sum = fractions.sum();
183    ensure!(
184        approx_eq!(T, sum, T::new(1.0), epsilon = 1e-5),
185        "Sum of fractions does not equal one (actual: {sum})"
186    );
187
188    Ok(())
189}
190
191/// Check whether an iterator contains values that are sorted and unique
192pub fn is_sorted_and_unique<T, I>(iter: I) -> bool
193where
194    T: PartialOrd + Clone,
195    I: IntoIterator<Item = T>,
196{
197    is_sorted_and_unique_with(iter, |a, b| a < b)
198}
199
200/// Check whether an iterator contains values that are sorted and unique, comparing with a custom
201/// function
202pub fn is_sorted_and_unique_with<T, I, F>(iter: I, mut less_than: F) -> bool
203where
204    T: Clone,
205    I: IntoIterator<Item = T>,
206    F: FnMut(&T, &T) -> bool,
207{
208    iter.into_iter()
209        .tuple_windows()
210        .all(|(a, b)| less_than(&a, &b))
211}
212
213/// Insert a key-value pair into a map implementing the `Insert` trait if the key does not
214/// already exist.
215///
216/// If the key already exists, this returns an error with a message indicating the key's
217/// existence.
218pub fn try_insert<M, K, V>(map: &mut M, key: &K, value: V) -> Result<()>
219where
220    M: Insert<K, V>,
221    K: Eq + Hash + Clone + std::fmt::Debug,
222{
223    let existing = map.insert(key.clone(), value).is_some();
224    ensure!(!existing, "Key {key:?} already exists in the map");
225    Ok(())
226}
227
228/// Format a list of items with a cap on display count for error messages
229pub fn format_items_with_cap<I, J, T>(items: I) -> String
230where
231    I: IntoIterator<Item = T, IntoIter = J>,
232    J: ExactSizeIterator<Item = T>,
233    T: fmt::Debug,
234{
235    const MAX_DISPLAY: usize = 10;
236
237    let items = items.into_iter();
238    let total_count = items.len();
239
240    // Format items with fmt::Debug::fmt() and separate with commas
241    let formatted_items = items
242        .take(MAX_DISPLAY)
243        .format_with(", ", |items, f| f(&format_args!("{items:?}")));
244    let mut out = format!("[{formatted_items}]");
245
246    // If there are remaining items, include the count
247    if total_count > MAX_DISPLAY {
248        write!(&mut out, " and {} more", total_count - MAX_DISPLAY).unwrap();
249    }
250
251    out
252}
253
254/// Read a model from the specified directory.
255///
256/// # Arguments
257///
258/// * `model_dir` - Folder containing model configuration files
259///
260/// # Returns
261///
262/// The static model data ([`Model`]) or an error.
263pub fn load_model<P: AsRef<Path>>(model_dir: P) -> Result<Model> {
264    let model_params = ModelParameters::from_path(&model_dir)?;
265
266    let time_slice_info = read_time_slice_info(model_dir.as_ref())?;
267    let regions = read_regions(model_dir.as_ref())?;
268    let region_ids = regions.keys().cloned().collect();
269    let years = &model_params.milestone_years;
270
271    let commodities = read_commodities(model_dir.as_ref(), &region_ids, &time_slice_info, years)?;
272    let processes = read_processes(
273        model_dir.as_ref(),
274        &commodities,
275        &region_ids,
276        &time_slice_info,
277        years,
278    )?;
279    let agents = read_agents(
280        model_dir.as_ref(),
281        &commodities,
282        &processes,
283        &region_ids,
284        years,
285    )?;
286    let agent_ids = agents.keys().cloned().collect();
287    let user_assets = read_user_assets(model_dir.as_ref(), &agent_ids, &processes, &region_ids)?;
288
289    // Build and validate commodity graphs for all regions and years
290    let commodity_graphs = build_commodity_graphs_for_model(&processes, &region_ids, years);
291    validate_commodity_graphs_for_model(
292        &commodity_graphs,
293        &processes,
294        &commodities,
295        &time_slice_info,
296    )?;
297
298    // Solve investment order for each region/year
299    let investment_order = solve_investment_order_for_model(&commodity_graphs, &commodities, years);
300
301    let model_path = model_dir
302        .as_ref()
303        .canonicalize()
304        .context("Could not parse path to model")?;
305    let model = Model {
306        model_path,
307        parameters: model_params,
308        agents,
309        commodities,
310        processes,
311        time_slice_info,
312        regions,
313        user_assets,
314        investment_order,
315    };
316    Ok(model)
317}
318
319/// Load commodity flow graphs for a model.
320///
321/// Loads necessary input data and creates a graph of commodity flows for each region and year,
322/// where nodes are commodities and edges are processes.
323///
324/// Graphs validation is NOT performed. This ensures that graphs can be generated even when
325/// validation would fail, which may be helpful for debugging.
326pub fn load_commodity_graphs<P: AsRef<Path>>(
327    model_dir: P,
328) -> Result<IndexMap<(RegionID, u32), CommoditiesGraph>> {
329    let model_params = ModelParameters::from_path(&model_dir)?;
330
331    let time_slice_info = read_time_slice_info(model_dir.as_ref())?;
332    let regions = read_regions(model_dir.as_ref())?;
333    let region_ids = regions.keys().cloned().collect();
334    let years = &model_params.milestone_years;
335
336    let commodities = read_commodities(model_dir.as_ref(), &region_ids, &time_slice_info, years)?;
337    let processes = read_processes(
338        model_dir.as_ref(),
339        &commodities,
340        &region_ids,
341        &time_slice_info,
342        years,
343    )?;
344
345    let commodity_graphs = build_commodity_graphs_for_model(&processes, &region_ids, years);
346    Ok(commodity_graphs)
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352    use crate::id::GenericID;
353    use crate::units::Dimensionless;
354    use rstest::rstest;
355    use serde::Deserialize;
356    use serde::de::IntoDeserializer;
357    use serde::de::value::{Error as ValueError, F64Deserializer};
358    use std::fs::File;
359    use std::io::Write;
360    use std::path::PathBuf;
361    use tempfile::tempdir;
362
363    #[derive(Debug, PartialEq, Deserialize)]
364    struct Record {
365        id: GenericID,
366        value: u32,
367    }
368
369    impl HasID<GenericID> for Record {
370        fn get_id(&self) -> &GenericID {
371            &self.id
372        }
373    }
374
375    /// Create an example CSV file in `dir_path`
376    fn create_csv_file(dir_path: &Path, contents: &str) -> PathBuf {
377        let file_path = dir_path.join("test.csv");
378        let mut file = File::create(&file_path).unwrap();
379        writeln!(file, "{contents}").unwrap();
380        file_path
381    }
382
383    /// Test a normal read
384    #[test]
385    fn read_csv_works() {
386        let dir = tempdir().unwrap();
387        let file_path = create_csv_file(dir.path(), "id,value\nhello,1\nworld,2\n");
388        let records: Vec<Record> = read_csv(&file_path).unwrap().collect();
389        assert_eq!(
390            records,
391            &[
392                Record {
393                    id: "hello".into(),
394                    value: 1,
395                },
396                Record {
397                    id: "world".into(),
398                    value: 2,
399                }
400            ]
401        );
402
403        // File with leading/trailing whitespace
404        let dir = tempdir().unwrap();
405        let file_path = create_csv_file(dir.path(), "id  , value\t\n  hello\t ,1\n world ,2\n");
406        let records: Vec<Record> = read_csv(&file_path).unwrap().collect();
407        assert_eq!(
408            records,
409            &[
410                Record {
411                    id: "hello".into(),
412                    value: 1,
413                },
414                Record {
415                    id: "world".into(),
416                    value: 2,
417                }
418            ]
419        );
420
421        // File with no data (only column headers)
422        let file_path = create_csv_file(dir.path(), "id,value\n");
423        assert!(read_csv::<Record>(&file_path).is_err());
424        assert!(
425            read_csv_optional::<Record>(&file_path)
426                .unwrap()
427                .next()
428                .is_none()
429        );
430
431        // Missing file
432        let dir = tempdir().unwrap();
433        let file_path = dir.path().join("a_missing_file.csv");
434        assert!(!file_path.exists());
435        assert!(read_csv::<Record>(&file_path).is_err());
436        // optional csv's should return empty iterator
437        assert!(
438            read_csv_optional::<Record>(&file_path)
439                .unwrap()
440                .next()
441                .is_none()
442        );
443    }
444
445    #[test]
446    fn read_toml_works() {
447        let dir = tempdir().unwrap();
448        let file_path = dir.path().join("test.toml");
449        {
450            let mut file = File::create(&file_path).unwrap();
451            writeln!(file, "id = \"hello\"\nvalue = 1").unwrap();
452        }
453
454        assert_eq!(
455            read_toml::<Record>(&file_path).unwrap(),
456            Record {
457                id: "hello".into(),
458                value: 1,
459            }
460        );
461
462        {
463            let mut file = File::create(&file_path).unwrap();
464            writeln!(file, "bad toml syntax").unwrap();
465        }
466
467        read_toml::<Record>(&file_path).unwrap_err();
468    }
469
470    /// Deserialise value with `deserialise_proportion_nonzero()`
471    fn deserialise_f64(value: f64) -> Result<Dimensionless, ValueError> {
472        let deserialiser: F64Deserializer<ValueError> = value.into_deserializer();
473        deserialise_proportion_nonzero(deserialiser)
474    }
475
476    #[test]
477    fn deserialise_proportion_nonzero_works() {
478        // Valid inputs
479        assert_eq!(deserialise_f64(0.01), Ok(Dimensionless(0.01)));
480        assert_eq!(deserialise_f64(0.5), Ok(Dimensionless(0.5)));
481        assert_eq!(deserialise_f64(1.0), Ok(Dimensionless(1.0)));
482
483        // Invalid inputs
484        deserialise_f64(0.0).unwrap_err();
485        deserialise_f64(-1.0).unwrap_err();
486        deserialise_f64(2.0).unwrap_err();
487        deserialise_f64(f64::NAN).unwrap_err();
488        deserialise_f64(f64::INFINITY).unwrap_err();
489    }
490
491    #[rstest]
492    #[case(0.0, true)] // Valid: zero
493    #[case(1e-10, true)] // Valid: small positive value
494    #[case(f64::MAX, true)] // Valid: maximum finite value
495    #[case(-1e-10, false)] // Invalid: negative value
496    #[case(-1.0, false)] // Invalid: negative value
497    #[case(f64::NAN, false)] // Invalid: NaN
498    #[case(f64::INFINITY, false)] // Invalid: positive infinity
499    #[case(f64::NEG_INFINITY, false)] // Invalid: negative infinity
500    fn deserialise_finite_non_negative_works(#[case] value: f64, #[case] expected_valid: bool) {
501        let deserialiser: F64Deserializer<ValueError> = value.into_deserializer();
502        let result = deserialise_finite_non_negative(deserialiser);
503        if expected_valid {
504            assert_eq!(result, Ok(Dimensionless(value)));
505        } else {
506            result.unwrap_err();
507        }
508    }
509
510    #[test]
511    fn check_values_sum_to_one_approx_works() {
512        // Single input, valid
513        check_values_sum_to_one_approx([Dimensionless(1.0)].into_iter()).unwrap();
514
515        // Multiple inputs, valid
516        check_values_sum_to_one_approx([Dimensionless(0.4), Dimensionless(0.6)].into_iter())
517            .unwrap();
518
519        // Single input, invalid
520        assert!(check_values_sum_to_one_approx([Dimensionless(0.5)].into_iter()).is_err());
521
522        // Multiple inputs, invalid
523        assert!(
524            check_values_sum_to_one_approx([Dimensionless(0.4), Dimensionless(0.3)].into_iter())
525                .is_err()
526        );
527
528        // Edge cases
529        assert!(
530            check_values_sum_to_one_approx([Dimensionless(f64::INFINITY)].into_iter()).is_err()
531        );
532        assert!(check_values_sum_to_one_approx([Dimensionless(f64::NAN)].into_iter()).is_err());
533    }
534
535    #[rstest]
536    #[case(&[], true)]
537    #[case(&[1], true)]
538    #[case(&[1,2], true)]
539    #[case(&[1,2,3,4], true)]
540    #[case(&[2,1],false)]
541    #[case(&[1,1],false)]
542    #[case(&[1,3,2,4], false)]
543    fn is_sorted_and_unique_works(#[case] values: &[u32], #[case] expected: bool) {
544        assert_eq!(is_sorted_and_unique(values), expected);
545    }
546
547    #[test]
548    fn format_items_with_cap_works() {
549        let items = vec!["a", "b", "c"];
550        assert_eq!(format_items_with_cap(&items), r#"["a", "b", "c"]"#);
551
552        // Test with more than 10 items to trigger the cap
553        let many_items = vec!["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"];
554        assert_eq!(
555            format_items_with_cap(&many_items),
556            r#"["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"] and 2 more"#
557        );
558    }
559}