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

muse2/
input.rs
