muse2/input/
time_slice.rs

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