muse2/
time_slice.rs

1//! Code for working with time slices.
2//!
3//! Time slices provide a mechanism for users to indicate that production (or other quantities)
4//! varies with time of day and season.
5use crate::id::{IDCollection, define_id_type};
6use crate::units::{Dimensionless, Year};
7use anyhow::{Context, Result};
8use indexmap::{IndexMap, IndexSet};
9use itertools::Itertools;
10use serde::de::Error;
11use serde::{Deserialize, Serialize};
12use serde_string_enum::DeserializeLabeledStringEnum;
13use std::fmt::Display;
14use std::iter;
15
16define_id_type! {Season}
17define_id_type! {TimeOfDay}
18
19/// An ID describing season and time of day
20#[derive(Hash, Eq, PartialEq, Ord, PartialOrd, Clone, Debug)]
21pub struct TimeSliceID {
22    /// The name of each season.
23    pub season: Season,
24    /// The name of each time slice within a day.
25    pub time_of_day: TimeOfDay,
26}
27
28/// Only implement for tests as this is a bit of a footgun
29#[cfg(test)]
30impl From<&str> for TimeSliceID {
31    fn from(value: &str) -> Self {
32        let (season, time_of_day) = value
33            .split('.')
34            .collect_tuple()
35            .expect("Time slice not in form season.time_of_day");
36        TimeSliceID {
37            season: season.into(),
38            time_of_day: time_of_day.into(),
39        }
40    }
41}
42
43impl Display for TimeSliceID {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        write!(f, "{}.{}", self.season, self.time_of_day)
46    }
47}
48
49impl<'de> Deserialize<'de> for TimeSliceID {
50    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
51    where
52        D: serde::Deserializer<'de>,
53    {
54        let s: &str = Deserialize::deserialize(deserializer)?;
55        let (season, time_of_day) = s.split('.').collect_tuple().ok_or_else(|| {
56            D::Error::custom(format!(
57                "Invalid input '{s}': Should be in form season.time_of_day"
58            ))
59        })?;
60        Ok(Self {
61            season: season.into(),
62            time_of_day: time_of_day.into(),
63        })
64    }
65}
66
67impl Serialize for TimeSliceID {
68    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
69    where
70        S: serde::Serializer,
71    {
72        serializer.collect_str(self)
73    }
74}
75
76/// Represents a time slice read from an input file, which can be all
77#[derive(PartialEq, Eq, Hash, Clone, Debug)]
78pub enum TimeSliceSelection {
79    /// All year and all day
80    Annual,
81    /// Only applies to one season
82    Season(Season),
83    /// Only applies to a single time slice
84    Single(TimeSliceID),
85}
86
87impl TimeSliceSelection {
88    /// The [`TimeSliceLevel`] to which this [`TimeSliceSelection`] corresponds
89    pub fn level(&self) -> TimeSliceLevel {
90        match self {
91            Self::Annual => TimeSliceLevel::Annual,
92            Self::Season(_) => TimeSliceLevel::Season,
93            Self::Single(_) => TimeSliceLevel::DayNight,
94        }
95    }
96
97    /// Iterate over the subset of time slices in this selection
98    pub fn iter<'a>(
99        &'a self,
100        time_slice_info: &'a TimeSliceInfo,
101    ) -> Box<dyn Iterator<Item = (&'a TimeSliceID, Year)> + 'a> {
102        let ts_info = time_slice_info;
103        match self {
104            Self::Annual => Box::new(ts_info.iter()),
105            Self::Season(season) => {
106                Box::new(ts_info.iter().filter(move |(ts, _)| ts.season == *season))
107            }
108            Self::Single(ts) => Box::new(iter::once((ts, ts_info.time_slices[ts]))),
109        }
110    }
111
112    /// Iterate over this [`TimeSliceSelection`] at the specified level.
113    ///
114    /// For example, this allows you to iterate over a [`TimeSliceSelection::Season`] at the level
115    /// of either seasons (in which case, the iterator will just contain the season) or time slices
116    /// (in which case it will contain all time slices for that season).
117    ///
118    /// Note that you cannot iterate over a [`TimeSliceSelection`] with coarser temporal granularity
119    /// than the [`TimeSliceSelection`] itself (for example, you cannot iterate over a
120    /// [`TimeSliceSelection::Season`] at the [`TimeSliceLevel::Annual`] level). In this case, the
121    /// function will return `None`.
122    pub fn iter_at_level<'a>(
123        &'a self,
124        time_slice_info: &'a TimeSliceInfo,
125        level: TimeSliceLevel,
126    ) -> Option<Box<dyn Iterator<Item = (Self, Year)> + 'a>> {
127        if level > self.level() {
128            return None;
129        }
130
131        let ts_info = time_slice_info;
132        let iter: Box<dyn Iterator<Item = _>> = match self {
133            Self::Annual => match level {
134                TimeSliceLevel::Annual => Box::new(iter::once((Self::Annual, Year(1.0)))),
135                TimeSliceLevel::Season => Box::new(
136                    ts_info
137                        .seasons
138                        .iter()
139                        .map(|(season, duration)| (season.clone().into(), *duration)),
140                ),
141                TimeSliceLevel::DayNight => Box::new(
142                    ts_info
143                        .time_slices
144                        .iter()
145                        .map(|(ts, duration)| (ts.clone().into(), *duration)),
146                ),
147            },
148            Self::Season(season) => match level {
149                TimeSliceLevel::Season => {
150                    Box::new(iter::once((self.clone(), ts_info.seasons[season])))
151                }
152                TimeSliceLevel::DayNight => Box::new(
153                    ts_info
154                        .time_slices
155                        .iter()
156                        .filter(move |(ts, _)| &ts.season == season)
157                        .map(|(ts, duration)| (ts.clone().into(), *duration)),
158                ),
159                TimeSliceLevel::Annual => unreachable!(),
160            },
161            Self::Single(time_slice) => Box::new(iter::once((
162                time_slice.clone().into(),
163                ts_info.time_slices[time_slice],
164            ))),
165        };
166
167        Some(iter)
168    }
169}
170
171impl From<TimeSliceID> for TimeSliceSelection {
172    fn from(value: TimeSliceID) -> Self {
173        Self::Single(value)
174    }
175}
176
177impl From<Season> for TimeSliceSelection {
178    fn from(value: Season) -> Self {
179        Self::Season(value)
180    }
181}
182
183impl Display for TimeSliceSelection {
184    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
185        match self {
186            Self::Annual => write!(f, "annual"),
187            Self::Season(season) => write!(f, "{season}"),
188            Self::Single(ts) => write!(f, "{ts}"),
189        }
190    }
191}
192
193/// The time granularity for a particular operation
194#[derive(
195    PartialEq, PartialOrd, Copy, Clone, Debug, DeserializeLabeledStringEnum, strum::EnumIter,
196)]
197pub enum TimeSliceLevel {
198    /// Treat individual time slices separately
199    #[string = "daynight"]
200    DayNight,
201    /// Whole seasons
202    #[string = "season"]
203    Season,
204    /// The whole year
205    #[string = "annual"]
206    Annual,
207}
208
209impl TimeSliceLevel {
210    /// Get the [`TimeSliceSelection`] containing the given [`TimeSliceID`] at this level.
211    pub fn containing_selection(&self, ts: &TimeSliceID) -> TimeSliceSelection {
212        match self {
213            Self::Annual => TimeSliceSelection::Annual,
214            Self::Season => TimeSliceSelection::Season(ts.season.clone()),
215            Self::DayNight => TimeSliceSelection::Single(ts.clone()),
216        }
217    }
218}
219
220/// Information about the time slices in the simulation, including names and durations
221#[derive(PartialEq, Debug)]
222pub struct TimeSliceInfo {
223    /// Names of times of day (e.g. "evening")
224    pub times_of_day: IndexSet<TimeOfDay>,
225    /// Names and fraction of year occupied by each season
226    pub seasons: IndexMap<Season, Year>,
227    /// The fraction of the year that this combination of season and time of day occupies
228    pub time_slices: IndexMap<TimeSliceID, Year>,
229}
230
231impl Default for TimeSliceInfo {
232    /// The default `TimeSliceInfo` is a single time slice covering the whole year
233    fn default() -> Self {
234        let id = TimeSliceID {
235            season: "all-year".into(),
236            time_of_day: "all-day".into(),
237        };
238        let time_slices = [(id.clone(), Year(1.0))].into_iter().collect();
239
240        Self {
241            seasons: iter::once((id.season, Year(1.0))).collect(),
242            times_of_day: iter::once(id.time_of_day).collect(),
243            time_slices,
244        }
245    }
246}
247
248impl TimeSliceInfo {
249    #[allow(clippy::doc_markdown)]
250    /// Get the `TimeSliceID` corresponding to the `time_slice`.
251    ///
252    /// `time_slice` must be in the form "season.time_of_day".
253    pub fn get_time_slice_id_from_str(&self, time_slice: &str) -> Result<TimeSliceID> {
254        let (season, time_of_day) = time_slice
255            .split('.')
256            .collect_tuple()
257            .context("Time slice must be in the form season.time_of_day")?;
258        let season = self
259            .seasons
260            .get_id(season)
261            .with_context(|| format!("{season} is not a known season"))?;
262        let time_of_day = self
263            .times_of_day
264            .get_id(time_of_day)
265            .with_context(|| format!("{time_of_day} is not a known time of day"))?;
266
267        Ok(TimeSliceID {
268            season: season.clone(),
269            time_of_day: time_of_day.clone(),
270        })
271    }
272
273    /// Get a `TimeSliceSelection` from the specified string.
274    ///
275    /// If the string is empty, the default value is `TimeSliceSelection::Annual`.
276    pub fn get_selection(&self, time_slice: &str) -> Result<TimeSliceSelection> {
277        if time_slice.eq_ignore_ascii_case("annual") {
278            Ok(TimeSliceSelection::Annual)
279        } else if time_slice.contains('.') {
280            let time_slice = self.get_time_slice_id_from_str(time_slice)?;
281            Ok(TimeSliceSelection::Single(time_slice))
282        } else {
283            let season = self
284                .seasons
285                .get_id(time_slice)
286                .with_context(|| format!("'{time_slice}' is not a valid season"))?
287                .clone();
288            Ok(TimeSliceSelection::Season(season))
289        }
290    }
291
292    /// Iterate over all [`TimeSliceID`]s
293    pub fn iter_ids(&self) -> indexmap::map::Keys<'_, TimeSliceID, Year> {
294        self.time_slices.keys()
295    }
296
297    /// Iterate over all seasons
298    pub fn iter_seasons(&self) -> indexmap::map::Keys<'_, Season, Year> {
299        self.seasons.keys()
300    }
301
302    /// Iterate over all time slices
303    pub fn iter(&self) -> impl Iterator<Item = (&TimeSliceID, Year)> {
304        self.time_slices
305            .iter()
306            .map(|(ts, duration)| (ts, *duration))
307    }
308
309    /// Iterate over the different time slice selections for a given time slice level.
310    ///
311    /// For example, if [`TimeSliceLevel::Season`] is specified, this function will return an
312    /// iterator of [`TimeSliceSelection`]s covering each season.
313    pub fn iter_selections_at_level(
314        &self,
315        level: TimeSliceLevel,
316    ) -> Box<dyn Iterator<Item = TimeSliceSelection> + '_> {
317        match level {
318            TimeSliceLevel::Annual => Box::new(iter::once(TimeSliceSelection::Annual)),
319            TimeSliceLevel::Season => {
320                Box::new(self.seasons.keys().cloned().map(TimeSliceSelection::Season))
321            }
322            TimeSliceLevel::DayNight => {
323                Box::new(self.iter_ids().cloned().map(TimeSliceSelection::Single))
324            }
325        }
326    }
327
328    /// Iterate over a subset of time slices calculating the relative duration of each.
329    ///
330    /// The relative duration is specified as a fraction of the total time covered by `selection`.
331    ///
332    /// # Arguments
333    ///
334    /// * `selection` - A subset of time slices
335    ///
336    /// # Returns
337    ///
338    /// An iterator of time slices along with the fraction of the total selection.
339    pub fn iter_selection_share<'a>(
340        &'a self,
341        selection: &'a TimeSliceSelection,
342        level: TimeSliceLevel,
343    ) -> Option<impl Iterator<Item = (TimeSliceSelection, Dimensionless)> + use<>> {
344        // Store selections as we have to iterate twice
345        let selections = selection.iter_at_level(self, level)?.collect_vec();
346
347        // Total duration covered by selection
348        let total_duration: Year = selections.iter().map(|(_, duration)| *duration).sum();
349
350        // Calculate share
351        let iter = selections
352            .into_iter()
353            .map(move |(selection, duration)| (selection, duration / total_duration));
354        Some(iter)
355    }
356
357    /// Calculate the total length of a selection of time slices.
358    pub fn length_for_selection(&self, selection: &TimeSliceSelection) -> Result<Year> {
359        let length: Year = selection.iter(self).map(|(_, duration)| duration).sum();
360        Ok(length)
361    }
362
363    /// Share a value between a subset of time slices in proportion to their lengths.
364    ///
365    /// For instance, you could use this function to compute how demand is distributed between the
366    /// different time slices of winter.
367    ///
368    /// # Arguments
369    ///
370    /// * `selection` - A subset of time slices
371    /// * `value` - The value to be shared between the time slices
372    ///
373    /// # Returns
374    ///
375    /// An iterator of time slices along with a fraction of `value`.
376    pub fn calculate_share<'a>(
377        &'a self,
378        selection: &'a TimeSliceSelection,
379        level: TimeSliceLevel,
380        value: Dimensionless,
381    ) -> Option<impl Iterator<Item = (TimeSliceSelection, Dimensionless)> + use<>> {
382        let iter = self
383            .iter_selection_share(selection, level)?
384            .map(move |(selection, share)| (selection, value * share));
385        Some(iter)
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392    use crate::units::UnitType;
393    use itertools::assert_equal;
394    use rstest::{fixture, rstest};
395
396    #[fixture]
397    fn time_slices1() -> [TimeSliceID; 2] {
398        [
399            TimeSliceID {
400                season: "winter".into(),
401                time_of_day: "day".into(),
402            },
403            TimeSliceID {
404                season: "summer".into(),
405                time_of_day: "night".into(),
406            },
407        ]
408    }
409
410    #[fixture]
411    fn time_slice_info1(time_slices1: [TimeSliceID; 2]) -> TimeSliceInfo {
412        TimeSliceInfo {
413            seasons: [("winter".into(), Year(0.5)), ("summer".into(), Year(0.5))]
414                .into_iter()
415                .collect(),
416            times_of_day: ["day".into(), "night".into()].into_iter().collect(),
417            time_slices: time_slices1.map(|ts| (ts, Year(0.5))).into_iter().collect(),
418        }
419    }
420
421    #[fixture]
422    fn time_slice_info2() -> TimeSliceInfo {
423        let time_slices = [
424            TimeSliceID {
425                season: "winter".into(),
426                time_of_day: "day".into(),
427            },
428            TimeSliceID {
429                season: "winter".into(),
430                time_of_day: "night".into(),
431            },
432            TimeSliceID {
433                season: "summer".into(),
434                time_of_day: "day".into(),
435            },
436            TimeSliceID {
437                season: "summer".into(),
438                time_of_day: "night".into(),
439            },
440        ];
441        TimeSliceInfo {
442            times_of_day: ["day".into(), "night".into()].into_iter().collect(),
443            seasons: [("winter".into(), Year(0.5)), ("summer".into(), Year(0.5))]
444                .into_iter()
445                .collect(),
446            time_slices: time_slices
447                .iter()
448                .map(|ts| (ts.clone(), Year(0.25)))
449                .collect(),
450        }
451    }
452
453    #[rstest]
454    fn ts_selection_iter_annual(time_slice_info1: TimeSliceInfo, time_slices1: [TimeSliceID; 2]) {
455        assert_equal(
456            TimeSliceSelection::Annual.iter(&time_slice_info1),
457            time_slices1.iter().map(|ts| (ts, Year(0.5))),
458        );
459    }
460
461    #[rstest]
462    fn ts_selection_iter_season(time_slice_info1: TimeSliceInfo, time_slices1: [TimeSliceID; 2]) {
463        assert_equal(
464            TimeSliceSelection::Season("winter".into()).iter(&time_slice_info1),
465            iter::once((&time_slices1[0], Year(0.5))),
466        );
467    }
468
469    #[rstest]
470    fn ts_selection_iter_single(time_slice_info1: TimeSliceInfo, time_slices1: [TimeSliceID; 2]) {
471        let ts = time_slice_info1
472            .get_time_slice_id_from_str("summer.night")
473            .unwrap();
474        assert_equal(
475            TimeSliceSelection::Single(ts).iter(&time_slice_info1),
476            iter::once((&time_slices1[1], Year(0.5))),
477        );
478    }
479
480    fn assert_selection_equal<I, T>(actual: Option<I>, expected: Option<Vec<(&str, T)>>)
481    where
482        T: UnitType,
483        I: Iterator<Item = (TimeSliceSelection, T)>,
484    {
485        let Some(actual) = actual else {
486            assert!(expected.is_none());
487            return;
488        };
489
490        let ts_info = time_slice_info2();
491        let expected = expected
492            .unwrap()
493            .into_iter()
494            .map(move |(sel, frac)| (ts_info.get_selection(sel).unwrap(), frac));
495        assert_equal(actual, expected);
496    }
497
498    #[rstest]
499    #[case(TimeSliceSelection::Annual, TimeSliceLevel::Annual, Some(vec![("annual", Year(1.0))]))]
500    #[case(TimeSliceSelection::Annual, TimeSliceLevel::Season, Some(vec![("winter", Year(0.5)), ("summer", Year(0.5))]))]
501    #[case(TimeSliceSelection::Annual, TimeSliceLevel::DayNight,
502           Some(vec![("winter.day", Year(0.25)), ("winter.night", Year(0.25)), ("summer.day", Year(0.25)), ("summer.night", Year(0.25))]))]
503    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::Annual, None)]
504    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::Season, Some(vec![("winter", Year(0.5))]))]
505    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::DayNight,
506           Some(vec![("winter.day", Year(0.25)), ("winter.night", Year(0.25))]))]
507    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::Annual, None)]
508    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::Season, None)]
509    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::DayNight, Some(vec![("winter.day", Year(0.25))]))]
510    fn ts_selection_iter_at_level(
511        time_slice_info2: TimeSliceInfo,
512        #[case] selection: TimeSliceSelection,
513        #[case] level: TimeSliceLevel,
514        #[case] expected: Option<Vec<(&str, Year)>>,
515    ) {
516        let actual = selection.iter_at_level(&time_slice_info2, level);
517        assert_selection_equal(actual, expected);
518    }
519
520    #[rstest]
521    #[case(TimeSliceSelection::Annual, TimeSliceLevel::Annual, Some(vec![("annual", Dimensionless(8.0))]))]
522    #[case(TimeSliceSelection::Annual, TimeSliceLevel::Season, Some(vec![("winter", Dimensionless(4.0)), ("summer", Dimensionless(4.0))]))]
523    #[case(TimeSliceSelection::Annual, TimeSliceLevel::DayNight,
524           Some(vec![("winter.day", Dimensionless(2.0)), ("winter.night", Dimensionless(2.0)), ("summer.day", Dimensionless(2.0)), ("summer.night", Dimensionless(2.0))]))]
525    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::Annual, None)]
526    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::Season, Some(vec![("winter", Dimensionless(8.0))]))]
527    #[case(TimeSliceSelection::Season("winter".into()), TimeSliceLevel::DayNight,
528           Some(vec![("winter.day", Dimensionless(4.0)), ("winter.night", Dimensionless(4.0))]))]
529    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::Annual, None)]
530    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::Season, None)]
531    #[case(TimeSliceSelection::Single("winter.day".into()), TimeSliceLevel::DayNight, Some(vec![("winter.day", Dimensionless(8.0))]))]
532    fn calculate_share(
533        time_slice_info2: TimeSliceInfo,
534        #[case] selection: TimeSliceSelection,
535        #[case] level: TimeSliceLevel,
536        #[case] expected: Option<Vec<(&str, Dimensionless)>>,
537    ) {
538        let actual = time_slice_info2.calculate_share(&selection, level, Dimensionless(8.0));
539        assert_selection_equal(actual, expected);
540    }
541}