muse2/input/agent/
cost_limit.rs

1//! Code for reading the agent cost limits CSV file.
2use super::super::*;
3use crate::agent::{AgentCostLimits, AgentCostLimitsMap, AgentID};
4use crate::id::IDCollection;
5use crate::year::parse_year_str;
6use anyhow::{Context, Result};
7use serde::Deserialize;
8use std::collections::HashMap;
9use std::path::Path;
10
11const AGENT_COST_LIMITS_FILE_NAME: &str = "agent_cost_limits.csv";
12
13#[derive(PartialEq, Debug, Deserialize)]
14struct AgentCostLimitsRaw {
15    agent_id: String,
16    year: String,
17    capex_limit: Option<f64>,
18    annual_cost_limit: Option<f64>,
19}
20
21impl AgentCostLimitsRaw {
22    fn to_agent_cost_limits(&self) -> AgentCostLimits {
23        AgentCostLimits {
24            capex_limit: self.capex_limit,
25            annual_cost_limit: self.annual_cost_limit,
26        }
27    }
28}
29
30/// Read agent cost limits info from the agent_cost_limits.csv file.
31///
32/// # Arguments
33///
34/// * `model_dir` - Folder containing model configuration files
35///
36/// # Returns
37///
38/// A map of Agents, with the agent ID as the key
39pub fn read_agent_cost_limits(
40    model_dir: &Path,
41    agent_ids: &HashSet<AgentID>,
42    milestone_years: &[u32],
43) -> Result<HashMap<AgentID, AgentCostLimitsMap>> {
44    let file_path = model_dir.join(AGENT_COST_LIMITS_FILE_NAME);
45    let agent_cost_limits_csv = read_csv_optional(&file_path)?;
46    read_agent_cost_limits_from_iter(agent_cost_limits_csv, agent_ids, milestone_years)
47        .with_context(|| input_err_msg(&file_path))
48}
49
50fn read_agent_cost_limits_from_iter<I>(
51    iter: I,
52    agent_ids: &HashSet<AgentID>,
53    milestone_years: &[u32],
54) -> Result<HashMap<AgentID, AgentCostLimitsMap>>
55where
56    I: Iterator<Item = AgentCostLimitsRaw>,
57{
58    let mut map: HashMap<AgentID, AgentCostLimitsMap> = HashMap::new();
59    for agent_cost_limits_raw in iter {
60        let cost_limits = agent_cost_limits_raw.to_agent_cost_limits();
61        let years = parse_year_str(&agent_cost_limits_raw.year, milestone_years)?;
62
63        // Get agent ID
64        let agent_id = agent_ids.get_id_by_str(&agent_cost_limits_raw.agent_id)?;
65
66        // Get or create entry in the map
67        let entry = map.entry(agent_id.clone()).or_default();
68
69        // Insert cost limits for the specified year(s)
70        for year in years {
71            entry.insert(year, cost_limits.clone());
72        }
73    }
74
75    // Validation: if cost limits are specified for an agent, they must be present for all years.
76    for (id, cost_limits) in map.iter() {
77        for year in milestone_years {
78            ensure!(
79                cost_limits.contains_key(year),
80                "Agent {id} is missing cost limits for year {year}"
81            );
82        }
83    }
84
85    Ok(map)
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::agent::AgentCostLimits;
92    use std::collections::HashSet;
93
94    fn create_agent_cost_limits_raw(
95        agent_id: &str,
96        year: &str,
97        capex_limit: Option<f64>,
98        annual_cost_limit: Option<f64>,
99    ) -> AgentCostLimitsRaw {
100        AgentCostLimitsRaw {
101            agent_id: agent_id.to_string(),
102            year: year.to_string(),
103            capex_limit,
104            annual_cost_limit,
105        }
106    }
107
108    #[test]
109    fn test_read_agent_cost_limits_from_iter_all_years() {
110        let agent_ids: HashSet<AgentID> = ["Agent1", "Agent2"]
111            .iter()
112            .map(|&id| AgentID::from(id))
113            .collect();
114        let milestone_years = [2020, 2025];
115
116        let iter = [
117            create_agent_cost_limits_raw("Agent1", "all", Some(100.0), Some(200.0)),
118            create_agent_cost_limits_raw("Agent2", "all", Some(150.0), Some(250.0)),
119        ]
120        .into_iter();
121
122        let result = read_agent_cost_limits_from_iter(iter, &agent_ids, &milestone_years).unwrap();
123
124        assert_eq!(result.len(), 2);
125        for year in milestone_years {
126            assert_eq!(
127                result[&AgentID::from("Agent1")][&year],
128                AgentCostLimits {
129                    capex_limit: Some(100.0),
130                    annual_cost_limit: Some(200.0),
131                }
132            );
133            assert_eq!(
134                result[&AgentID::from("Agent2")][&year],
135                AgentCostLimits {
136                    capex_limit: Some(150.0),
137                    annual_cost_limit: Some(250.0),
138                }
139            );
140        }
141    }
142
143    #[test]
144    fn test_read_agent_cost_limits_from_iter_some_years() {
145        let agent_ids: HashSet<AgentID> = ["Agent1"].iter().map(|&id| AgentID::from(id)).collect();
146        let milestone_years = [2020, 2025];
147
148        let iter = [create_agent_cost_limits_raw(
149            "Agent1",
150            "2020;2025",
151            Some(100.0),
152            Some(200.0),
153        )]
154        .into_iter();
155
156        let result = read_agent_cost_limits_from_iter(iter, &agent_ids, &milestone_years).unwrap();
157
158        assert_eq!(result.len(), 1);
159        assert_eq!(
160            result[&AgentID::from("Agent1")][&2020],
161            AgentCostLimits {
162                capex_limit: Some(100.0),
163                annual_cost_limit: Some(200.0),
164            }
165        );
166        assert_eq!(
167            result[&AgentID::from("Agent1")][&2025],
168            AgentCostLimits {
169                capex_limit: Some(100.0),
170                annual_cost_limit: Some(200.0),
171            }
172        );
173    }
174
175    #[test]
176    fn test_read_agent_cost_limits_from_iter_missing_years() {
177        let agent_ids: HashSet<AgentID> = ["Agent1"].iter().map(|&id| AgentID::from(id)).collect();
178        let milestone_years = [2020, 2025];
179
180        let iter = [create_agent_cost_limits_raw(
181            "Agent1",
182            "2020",
183            Some(100.0),
184            Some(200.0),
185        )]
186        .into_iter();
187
188        let result = read_agent_cost_limits_from_iter(iter, &agent_ids, &milestone_years);
189
190        assert!(result.is_err());
191        assert_eq!(
192            result.unwrap_err().to_string(),
193            "Agent Agent1 is missing cost limits for year 2025"
194        );
195    }
196}
muse2/input/agent/cost_limit.rs

muse2/input/agent/
cost_limit.rs
