muse2/
model.rs

1//! Code for simulation models.
2#![allow(missing_docs)]
3use crate::agent::AgentMap;
4use crate::commodity::CommodityMap;
5use crate::input::{input_err_msg, read_toml};
6use crate::process::ProcessMap;
7use crate::region::{RegionID, RegionMap};
8use crate::time_slice::TimeSliceInfo;
9use anyhow::{ensure, Context, Result};
10use serde::Deserialize;
11use std::path::Path;
12
13const MODEL_FILE_NAME: &str = "model.toml";
14
15/// Model definition
16pub struct Model {
17    pub milestone_years: Vec<u32>,
18    pub agents: AgentMap,
19    pub commodities: CommodityMap,
20    pub processes: ProcessMap,
21    pub time_slice_info: TimeSliceInfo,
22    pub regions: RegionMap,
23}
24
25/// Represents the contents of the entire model file.
26#[derive(Debug, Deserialize, PartialEq)]
27pub struct ModelFile {
28    pub milestone_years: MilestoneYears,
29}
30
31/// Represents the "milestone_years" section of the model file.
32#[derive(Debug, Deserialize, PartialEq)]
33pub struct MilestoneYears {
34    pub years: Vec<u32>,
35}
36
37/// Check that the milestone years parameter is valid
38///
39/// # Arguments
40///
41/// * `years` - Integer list of milestone years
42///
43/// # Returns
44///
45/// An error if the milestone years are invalid
46fn check_milestone_years(years: &[u32]) -> Result<()> {
47    ensure!(!years.is_empty(), "`milestone_years` is empty");
48
49    ensure!(
50        years[..years.len() - 1]
51            .iter()
52            .zip(years[1..].iter())
53            .all(|(y1, y2)| y1 < y2),
54        "`milestone_years` must be composed of unique values in order"
55    );
56
57    Ok(())
58}
59
60impl ModelFile {
61    /// Read a model file from the specified directory.
62    ///
63    /// # Arguments
64    ///
65    /// * `model_dir` - Folder containing model configuration files
66    ///
67    /// # Returns
68    ///
69    /// The model file contents as a `ModelFile` struct or an error if the file is invalid
70    pub fn from_path<P: AsRef<Path>>(model_dir: P) -> Result<ModelFile> {
71        let file_path = model_dir.as_ref().join(MODEL_FILE_NAME);
72        let model_file: ModelFile = read_toml(&file_path)?;
73        check_milestone_years(&model_file.milestone_years.years)
74            .with_context(|| input_err_msg(file_path))?;
75
76        Ok(model_file)
77    }
78}
79
80impl Model {
81    /// Iterate over the model's milestone years.
82    pub fn iter_years(&self) -> impl Iterator<Item = u32> + '_ {
83        self.milestone_years.iter().copied()
84    }
85
86    /// Iterate over the model's regions (region IDs).
87    pub fn iter_regions(&self) -> impl Iterator<Item = &RegionID> + '_ {
88        self.regions.keys()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use std::fs::File;
96    use std::io::Write;
97    use tempfile::tempdir;
98
99    #[test]
100    fn test_check_milestone_years() {
101        // Valid
102        assert!(check_milestone_years(&[1]).is_ok());
103        assert!(check_milestone_years(&[1, 2]).is_ok());
104
105        // Invalid
106        assert!(check_milestone_years(&[]).is_err());
107        assert!(check_milestone_years(&[1, 1]).is_err());
108        assert!(check_milestone_years(&[2, 1]).is_err());
109    }
110
111    #[test]
112    fn test_model_file_from_path() {
113        let dir = tempdir().unwrap();
114        {
115            let mut file = File::create(dir.path().join(MODEL_FILE_NAME)).unwrap();
116            writeln!(file, "[milestone_years]\nyears = [2020, 2100]").unwrap();
117        }
118
119        let model_file = ModelFile::from_path(dir.path()).unwrap();
120        assert_eq!(model_file.milestone_years.years, [2020, 2100]);
121    }
122}