1use super::super::*;
3use crate::agent::{AgentID, AgentMap, AgentSearchSpaceMap};
4use crate::commodity::CommodityID;
5use crate::id::IDCollection;
6use crate::process::{Process, ProcessMap};
7use crate::year::parse_year_str;
8use anyhow::{Context, Result};
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 agent_id: String,
20 commodity_id: String,
22 years: String,
24 search_space: String,
27}
28
29#[derive(Debug)]
31struct AgentSearchSpace {
32 agent_id: AgentID,
34 commodity_id: CommodityID,
36 years: Vec<u32>,
38 search_space: Rc<Vec<Rc<Process>>>,
40}
41
42impl AgentSearchSpaceRaw {
43 fn into_agent_search_space(
44 self,
45 agents: &AgentMap,
46 processes: &ProcessMap,
47 commodity_ids: &HashSet<CommodityID>,
48 milestone_years: &[u32],
49 ) -> Result<AgentSearchSpace> {
50 let search_space = Rc::new(parse_search_space_str(&self.search_space, processes)?);
52
53 let commodity_id = commodity_ids.get_id(&self.commodity_id)?;
55
56 let year = parse_year_str(&self.years, milestone_years)?;
58
59 let agent_id = agents.get_id(&self.agent_id)?;
60
61 Ok(AgentSearchSpace {
62 agent_id: agent_id.clone(),
63 commodity_id: commodity_id.clone(),
64 years: year,
65 search_space,
66 })
67 }
68}
69
70fn parse_search_space_str(search_space: &str, processes: &ProcessMap) -> Result<Vec<Rc<Process>>> {
77 let search_space = search_space.trim();
78 if search_space.is_empty() || search_space.eq_ignore_ascii_case("all") {
79 Ok(processes.values().cloned().collect())
80 } else {
81 search_space
82 .split(';')
83 .map(|id| {
84 let process = processes
85 .get(id.trim())
86 .with_context(|| format!("Invalid process '{id}'"))?;
87 Ok(process.clone())
88 })
89 .try_collect()
90 }
91}
92
93pub fn read_agent_search_space(
107 model_dir: &Path,
108 agents: &AgentMap,
109 processes: &ProcessMap,
110 commodity_ids: &HashSet<CommodityID>,
111 milestone_years: &[u32],
112) -> Result<HashMap<AgentID, AgentSearchSpaceMap>> {
113 let file_path = model_dir.join(AGENT_SEARCH_SPACE_FILE_NAME);
114 let iter = read_csv_optional::<AgentSearchSpaceRaw>(&file_path)?;
115 read_agent_search_space_from_iter(iter, agents, processes, commodity_ids, milestone_years)
116 .with_context(|| input_err_msg(&file_path))
117}
118
119fn read_agent_search_space_from_iter<I>(
120 iter: I,
121 agents: &AgentMap,
122 processes: &ProcessMap,
123 commodity_ids: &HashSet<CommodityID>,
124 milestone_years: &[u32],
125) -> Result<HashMap<AgentID, AgentSearchSpaceMap>>
126where
127 I: Iterator<Item = AgentSearchSpaceRaw>,
128{
129 let mut search_spaces = HashMap::new();
130 for search_space_raw in iter {
131 let search_space = search_space_raw.into_agent_search_space(
132 agents,
133 processes,
134 commodity_ids,
135 milestone_years,
136 )?;
137
138 let map = search_spaces
140 .entry(search_space.agent_id)
141 .or_insert_with(AgentSearchSpaceMap::new);
142
143 for year in search_space.years {
145 try_insert(
146 map,
147 (search_space.commodity_id.clone(), year),
148 search_space.search_space.clone(),
149 )?;
150 }
151 }
152
153 Ok(search_spaces)
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159 use crate::fixture::{agents, assert_error, region_ids};
160 use crate::process::{
161 ProcessActivityLimitsMap, ProcessFlowsMap, ProcessID, ProcessParameterMap,
162 };
163 use crate::region::RegionID;
164 use indexmap::IndexSet;
165 use rstest::{fixture, rstest};
166 use std::iter;
167
168 #[fixture]
169 pub fn processes(region_ids: IndexSet<RegionID>) -> ProcessMap {
170 ["A", "B", "C"]
171 .map(|id| {
172 let id: ProcessID = id.into();
173 let process = Process {
174 id: id.clone(),
175 description: "Description".into(),
176 years: vec![2010, 2020],
177 activity_limits: ProcessActivityLimitsMap::new(),
178 flows: ProcessFlowsMap::new(),
179 parameters: ProcessParameterMap::new(),
180 regions: region_ids.clone(),
181 };
182 (id, process.into())
183 })
184 .into_iter()
185 .collect()
186 }
187
188 #[fixture]
189 fn commodity_ids() -> HashSet<CommodityID> {
190 iter::once("commodity1".into()).collect()
191 }
192
193 #[rstest]
194 fn test_search_space_raw_into_search_space_valid(
195 agents: AgentMap,
196 processes: ProcessMap,
197 commodity_ids: HashSet<CommodityID>,
198 ) {
199 let raw = AgentSearchSpaceRaw {
201 agent_id: "agent1".into(),
202 commodity_id: "commodity1".into(),
203 years: "2020".into(),
204 search_space: "A;B".into(),
205 };
206 assert!(raw
207 .into_agent_search_space(&agents, &processes, &commodity_ids, &[2020])
208 .is_ok());
209 }
210
211 #[rstest]
212 fn test_search_space_raw_into_search_space_invalid_commodity_id(
213 agents: AgentMap,
214 processes: ProcessMap,
215 commodity_ids: HashSet<CommodityID>,
216 ) {
217 let raw = AgentSearchSpaceRaw {
219 agent_id: "agent1".into(),
220 commodity_id: "invalid_commodity".into(),
221 years: "2020".into(),
222 search_space: "A;B".into(),
223 };
224 assert_error!(
225 raw.into_agent_search_space(&agents, &processes, &commodity_ids, &[2020]),
226 "Unknown ID invalid_commodity found"
227 );
228 }
229
230 #[rstest]
231 fn test_search_space_raw_into_search_space_invalid_process_id(
232 agents: AgentMap,
233 processes: ProcessMap,
234 commodity_ids: HashSet<CommodityID>,
235 ) {
236 let raw = AgentSearchSpaceRaw {
238 agent_id: "agent1".into(),
239 commodity_id: "commodity1".into(),
240 years: "2020".into(),
241 search_space: "A;D".into(),
242 };
243 assert_error!(
244 raw.into_agent_search_space(&agents, &processes, &commodity_ids, &[2020]),
245 "Invalid process 'D'"
246 );
247 }
248}