muse2/input/
time_slice.rs

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

muse2/input/
time_slice.rs
