muse2/input/agent/
search_space.rs

1//! Code for reading the agent search space CSV file.
2use super::super::*;
3use crate::agent::{AgentID, AgentMap, AgentSearchSpace};
4use crate::commodity::CommodityMap;
5use crate::id::IDCollection;
6use crate::process::ProcessID;
7use anyhow::{Context, Result};
8use indexmap::IndexSet;
9use serde::Deserialize;
10use std::collections::HashMap;
11use std::path::Path;
12use std::rc::Rc;
13
14const AGENT_SEARCH_SPACE_FILE_NAME: &str = "agent_search_space.csv";
15
16#[derive(PartialEq, Debug, Deserialize)]
17struct AgentSearchSpaceRaw {
18    /// The agent to apply the search space to.
19    agent_id: String,
20    /// The commodity to apply the search space to.
21    commodity_id: String,
22    /// The year to apply the search space to.
23    year: u32,
24    /// The processes that the agent will consider investing in. Expressed as process IDs separated
25    /// by semicolons or `None`, meaning all processes.
26    search_space: String,
27}
28
29impl AgentSearchSpaceRaw {
30    fn to_agent_search_space(
31        &self,
32        process_ids: &IndexSet<ProcessID>,
33        commodities: &CommodityMap,
34        milestone_years: &[u32],
35    ) -> Result<AgentSearchSpace> {
36        // Parse search_space string
37        let search_space = parse_search_space_str(&self.search_space, process_ids)?;
38
39        // Get commodity
40        let commodity = commodities
41            .get(self.commodity_id.as_str())
42            .context("Invalid commodity ID")?;
43
44        // Check that the year is a valid milestone year
45        ensure!(
46            milestone_years.binary_search(&self.year).is_ok(),
47            "Invalid milestone year {}",
48            self.year
49        );
50
51        // Create AgentSearchSpace
52        Ok(AgentSearchSpace {
53            year: self.year,
54            commodity: Rc::clone(commodity),
55            search_space,
56        })
57    }
58}
59
60/// Parse a string representing the processes the agent will invest in.
61///
62/// This string can either be:
63///  * Empty, meaning all processes
64///  * "all", meaning the same
65///  * A list of process IDs separated by semicolons
66fn parse_search_space_str(
67    search_space: &str,
68    process_ids: &IndexSet<ProcessID>,
69) -> Result<Vec<ProcessID>> {
70    let search_space = search_space.trim();
71    if search_space.is_empty() || search_space.eq_ignore_ascii_case("all") {
72        Ok(process_ids.iter().cloned().collect())
73    } else {
74        search_space
75            .split(';')
76            .map(|id| process_ids.get_id_by_str(id.trim()))
77            .try_collect()
78    }
79}
80
81/// Read agent search space info from the agent_search_space.csv file.
82///
83/// # Arguments
84///
85/// * `model_dir` - Folder containing model configuration files
86///
87/// # Returns
88///
89/// A map of Agents, with the agent ID as the key
90pub fn read_agent_search_space(
91    model_dir: &Path,
92    agents: &AgentMap,
93    process_ids: &IndexSet<ProcessID>,
94    commodities: &CommodityMap,
95    milestone_years: &[u32],
96) -> Result<HashMap<AgentID, Vec<AgentSearchSpace>>> {
97    let file_path = model_dir.join(AGENT_SEARCH_SPACE_FILE_NAME);
98    let iter = read_csv_optional::<AgentSearchSpaceRaw>(&file_path)?;
99    read_agent_search_space_from_iter(iter, agents, process_ids, commodities, milestone_years)
100        .with_context(|| input_err_msg(&file_path))
101}
102
103fn read_agent_search_space_from_iter<I>(
104    iter: I,
105    agents: &AgentMap,
106    process_ids: &IndexSet<ProcessID>,
107    commodities: &CommodityMap,
108    milestone_years: &[u32],
109) -> Result<HashMap<AgentID, Vec<AgentSearchSpace>>>
110where
111    I: Iterator<Item = AgentSearchSpaceRaw>,
112{
113    let mut search_spaces = HashMap::new();
114    for search_space_raw in iter {
115        let search_space =
116            search_space_raw.to_agent_search_space(process_ids, commodities, milestone_years)?;
117
118        let (id, _agent) = agents
119            .get_key_value(search_space_raw.agent_id.as_str())
120            .context("Invalid agent ID")?;
121
122        // Append to Vec with the corresponding key or create
123        search_spaces
124            .entry(id.clone())
125            .or_insert_with(|| Vec::with_capacity(1))
126            .push(search_space);
127    }
128
129    Ok(search_spaces)
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::commodity::{Commodity, CommodityCostMap, CommodityType, DemandMap};
136    use crate::time_slice::TimeSliceLevel;
137    use std::iter;
138
139    #[test]
140    fn test_search_space_raw_into_search_space() {
141        let process_ids = ["A".into(), "B".into(), "C".into()].into_iter().collect();
142        let commodity = Rc::new(Commodity {
143            id: "commodity1".into(),
144            description: "A commodity".into(),
145            kind: CommodityType::SupplyEqualsDemand,
146            time_slice_level: TimeSliceLevel::Annual,
147            costs: CommodityCostMap::new(),
148            demand: DemandMap::new(),
149        });
150        let commodities = iter::once(("commodity1".into(), Rc::clone(&commodity))).collect();
151
152        // Valid search space
153        let raw = AgentSearchSpaceRaw {
154            agent_id: "agent1".into(),
155            commodity_id: "commodity1".into(),
156            year: 2020,
157            search_space: "A;B".into(),
158        };
159        assert!(raw
160            .to_agent_search_space(&process_ids, &commodities, &[2020])
161            .is_ok());
162
163        // Invalid commodity ID
164        let raw = AgentSearchSpaceRaw {
165            agent_id: "agent1".into(),
166            commodity_id: "invalid_commodity".into(),
167            year: 2020,
168            search_space: "A;B".into(),
169        };
170        assert!(raw
171            .to_agent_search_space(&process_ids, &commodities, &[2020])
172            .is_err());
173
174        // Invalid process ID
175        let raw = AgentSearchSpaceRaw {
176            agent_id: "agent1".into(),
177            commodity_id: "commodity1".into(),
178            year: 2020,
179            search_space: "A;D".into(),
180        };
181        assert!(raw
182            .to_agent_search_space(&process_ids, &commodities, &[2020])
183            .is_err());
184    }
185}
muse2/input/agent/search_space.rs

muse2/input/agent/
search_space.rs
