muse2/input/agent/
search_space.rs

1//! Code for reading the agent search space CSV file.
2use 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    /// 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(s) to apply the search space to.
23    years: String,
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
29/// Search space for an agent
30#[derive(Debug)]
31struct AgentSearchSpace {
32    /// The agent to which this search space applies
33    agent_id: AgentID,
34    /// The commodity to apply the search space to
35    commodity_id: CommodityID,
36    /// The year(s) the objective is relevant for
37    years: Vec<u32>,
38    /// The agent's search space
39    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        // Parse search_space string
51        let search_space = Rc::new(parse_search_space_str(&self.search_space, processes)?);
52
53        // Get commodity
54        let commodity_id = commodity_ids.get_id(&self.commodity_id)?;
55
56        // Check that the year is a valid milestone year
57        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
70/// Parse a string representing the processes the agent will invest in.
71///
72/// This string can either be:
73///  * Empty, meaning all processes
74///  * "all", meaning the same
75///  * A list of process IDs separated by semicolons
76fn 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
93/// Read agent search space info from the agent_search_space.csv file.
94///
95/// # Arguments
96///
97/// * `model_dir` - Folder containing model configuration files
98/// * `agents` - Map of agents
99/// * `processes` - Map of processes
100/// * `commodity_ids` - All possible valid commodity IDs
101/// * `milestone_years` - The milestone years of the simulation
102///
103/// # Returns
104///
105/// A map of Agents, with the agent ID as the key
106pub 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        // Get or create search space map
139        let map = search_spaces
140            .entry(search_space.agent_id)
141            .or_insert_with(AgentSearchSpaceMap::new);
142
143        // Store process IDs
144        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        // Valid search space
200        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        // Invalid commodity ID
218        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        // Invalid process ID
237        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}