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::Dimensionless;
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: Dimensionless,
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 fractions = 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            fractions.insert(id.clone(), time_slice.fraction).is_none(),
64            "Duplicate time slice entry for {id}",
65        );
66    }
67
68    // Validate data
69    check_fractions_sum_to_one(fractions.values().cloned())
70        .context("Invalid time slice fractions")?;
71
72    Ok(TimeSliceInfo {
73        times_of_day,
74        seasons,
75        time_slices: fractions,
76    })
77}
78
79/// Read time slices from a CSV file.
80///
81/// # Arguments
82///
83/// * `model_dir` - Folder containing model configuration files
84///
85/// # Returns
86///
87/// This function returns a `TimeSliceInfo` struct or, if the file doesn't exist, a single time
88/// slice covering the whole year (see `TimeSliceInfo::default()`).
89pub fn read_time_slice_info(model_dir: &Path) -> Result<TimeSliceInfo> {
90    let file_path = model_dir.join(TIME_SLICES_FILE_NAME);
91    if !file_path.exists() {
92        return Ok(TimeSliceInfo::default());
93    }
94
95    let time_slices_csv = read_csv(&file_path)?;
96    read_time_slice_info_from_iter(time_slices_csv).with_context(|| input_err_msg(file_path))
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use std::fs::File;
103    use std::io::Write;
104    use std::path::Path;
105    use tempfile::tempdir;
106
107    /// Create an example time slices file in dir_path
108    fn create_time_slices_file(dir_path: &Path) {
109        let file_path = dir_path.join(TIME_SLICES_FILE_NAME);
110        let mut file = File::create(file_path).unwrap();
111        writeln!(
112            file,
113            "season,time_of_day,fraction
114winter,day,0.25
115peak,night,0.25
116summer,peak,0.25
117autumn,evening,0.25"
118        )
119        .unwrap();
120    }
121
122    #[test]
123    fn test_read_time_slice_info() {
124        let dir = tempdir().unwrap();
125        create_time_slices_file(dir.path());
126
127        let info = read_time_slice_info(dir.path()).unwrap();
128        assert_eq!(
129            info,
130            TimeSliceInfo {
131                seasons: [
132                    ("winter".into(), Dimensionless(0.25)),
133                    ("peak".into(), Dimensionless(0.25)),
134                    ("summer".into(), Dimensionless(0.25)),
135                    ("autumn".into(), Dimensionless(0.25))
136                ]
137                .into_iter()
138                .collect(),
139                times_of_day: [
140                    "day".into(),
141                    "night".into(),
142                    "peak".into(),
143                    "evening".into()
144                ]
145                .into_iter()
146                .collect(),
147                time_slices: [
148                    (
149                        TimeSliceID {
150                            season: "winter".into(),
151                            time_of_day: "day".into(),
152                        },
153                        Dimensionless(0.25),
154                    ),
155                    (
156                        TimeSliceID {
157                            season: "peak".into(),
158                            time_of_day: "night".into(),
159                        },
160                        Dimensionless(0.25),
161                    ),
162                    (
163                        TimeSliceID {
164                            season: "summer".into(),
165                            time_of_day: "peak".into(),
166                        },
167                        Dimensionless(0.25),
168                    ),
169                    (
170                        TimeSliceID {
171                            season: "autumn".into(),
172                            time_of_day: "evening".into(),
173                        },
174                        Dimensionless(0.25),
175                    ),
176                ]
177                .into_iter()
178                .collect()
179            }
180        );
181    }
182
183    #[test]
184    fn test_read_time_slice_info_non_existent() {
185        let actual = read_time_slice_info(tempdir().unwrap().path());
186        assert_eq!(actual.unwrap(), TimeSliceInfo::default());
187    }
188}
muse2/input/time_slice.rs

muse2/input/
time_slice.rs
