muse2/input/
time_slice.rs

1//! Code for reading in time slice info from a CSV file.
2use super::*;
3use crate::time_slice::{TimeSliceID, TimeSliceInfo};
4use anyhow::{ensure, Context, Result};
5use indexmap::{IndexMap, IndexSet};
6use serde::Deserialize;
7use std::path::Path;
8use std::rc::Rc;
9
10const TIME_SLICES_FILE_NAME: &str = "time_slices.csv";
11
12/// A time slice record retrieved from a CSV file
13#[derive(PartialEq, Debug, Deserialize)]
14struct TimeSliceRaw {
15    season: String,
16    time_of_day: String,
17    #[serde(deserialize_with = "deserialise_proportion_nonzero")]
18    fraction: f64,
19}
20
21/// Get the specified `String` from `set` or insert if it doesn't exist
22fn get_or_insert(value: String, set: &mut IndexSet<Rc<str>>) -> Rc<str> {
23    // Sadly there's no entry API for HashSets: https://github.com/rust-lang/rfcs/issues/1490
24    match set.get(value.as_str()) {
25        Some(value) => Rc::clone(value),
26        None => {
27            let value = Rc::from(value);
28            set.insert(Rc::clone(&value));
29            value
30        }
31    }
32}
33
34/// Read time slice information from an iterator of raw time slice records
35fn read_time_slice_info_from_iter<I>(iter: I) -> Result<TimeSliceInfo>
36where
37    I: Iterator<Item = TimeSliceRaw>,
38{
39    let mut seasons = IndexSet::new();
40    let mut times_of_day = IndexSet::new();
41    let mut fractions = IndexMap::new();
42    for time_slice in iter {
43        let season = get_or_insert(time_slice.season, &mut seasons);
44        let time_of_day = get_or_insert(time_slice.time_of_day, &mut times_of_day);
45        let id = TimeSliceID {
46            season,
47            time_of_day,
48        };
49
50        ensure!(
51            fractions.insert(id.clone(), time_slice.fraction).is_none(),
52            "Duplicate time slice entry for {id}",
53        );
54    }
55
56    // Validate data
57    check_fractions_sum_to_one(fractions.values().cloned())
58        .context("Invalid time slice fractions")?;
59
60    Ok(TimeSliceInfo {
61        seasons,
62        times_of_day,
63        fractions,
64    })
65}
66
67/// Read time slices from a CSV file.
68///
69/// # Arguments
70///
71/// * `model_dir` - Folder containing model configuration files
72///
73/// # Returns
74///
75/// This function returns a `TimeSliceInfo` struct or, if the file doesn't exist, a single time
76/// slice covering the whole year (see `TimeSliceInfo::default()`).
77pub fn read_time_slice_info(model_dir: &Path) -> Result<TimeSliceInfo> {
78    let file_path = model_dir.join(TIME_SLICES_FILE_NAME);
79    if !file_path.exists() {
80        return Ok(TimeSliceInfo::default());
81    }
82
83    let time_slices_csv = read_csv(&file_path)?;
84    read_time_slice_info_from_iter(time_slices_csv).with_context(|| input_err_msg(file_path))
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use std::fs::File;
91    use std::io::Write;
92    use std::path::Path;
93    use tempfile::tempdir;
94
95    /// Create an example time slices file in dir_path
96    fn create_time_slices_file(dir_path: &Path) {
97        let file_path = dir_path.join(TIME_SLICES_FILE_NAME);
98        let mut file = File::create(file_path).unwrap();
99        writeln!(
100            file,
101            "season,time_of_day,fraction
102winter,day,0.25
103peak,night,0.25
104summer,peak,0.25
105autumn,evening,0.25"
106        )
107        .unwrap();
108    }
109
110    #[test]
111    fn test_read_time_slice_info() {
112        let dir = tempdir().unwrap();
113        create_time_slices_file(dir.path());
114
115        let info = read_time_slice_info(dir.path()).unwrap();
116        assert_eq!(
117            info,
118            TimeSliceInfo {
119                seasons: [
120                    "winter".into(),
121                    "peak".into(),
122                    "summer".into(),
123                    "autumn".into()
124                ]
125                .into_iter()
126                .collect(),
127                times_of_day: [
128                    "day".into(),
129                    "night".into(),
130                    "peak".into(),
131                    "evening".into()
132                ]
133                .into_iter()
134                .collect(),
135                fractions: [
136                    (
137                        TimeSliceID {
138                            season: "winter".into(),
139                            time_of_day: "day".into(),
140                        },
141                        0.25,
142                    ),
143                    (
144                        TimeSliceID {
145                            season: "peak".into(),
146                            time_of_day: "night".into(),
147                        },
148                        0.25,
149                    ),
150                    (
151                        TimeSliceID {
152                            season: "summer".into(),
153                            time_of_day: "peak".into(),
154                        },
155                        0.25,
156                    ),
157                    (
158                        TimeSliceID {
159                            season: "autumn".into(),
160                            time_of_day: "evening".into(),
161                        },
162                        0.25,
163                    ),
164                ]
165                .into_iter()
166                .collect()
167            }
168        );
169    }
170
171    #[test]
172    fn test_read_time_slice_info_non_existent() {
173        let actual = read_time_slice_info(tempdir().unwrap().path());
174        assert_eq!(actual.unwrap(), TimeSliceInfo::default());
175    }
176}