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