muse2/input/agent/
cost_limit.rs1use 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
31pub 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 let agent_id = agent_ids.get_id(&agent_cost_limits_raw.agent_id)?;
66
67 let entry = map.entry(agent_id.clone()).or_default();
69
70 for year in years {
72 entry.insert(year, cost_limits.clone());
73 }
74 }
75
76 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}