1use super::*;
2
3pub trait LocoTrait {
5 fn set_curr_pwr_max_out(
15 &mut self,
16 pwr_aux: Option<si::Power>,
17 elev_and_temp: Option<(si::Length, si::ThermodynamicTemperature)>,
18 train_mass: Option<si::Mass>,
19 train_speed: Option<si::Velocity>,
20 dt: si::Time,
21 ) -> anyhow::Result<()>;
22 fn get_energy_loss(&self) -> anyhow::Result<si::Energy>;
24}
25
26#[serde_api]
27#[derive(Debug, Default, Serialize, Deserialize, Clone, PartialEq)]
28#[cfg_attr(feature = "pyo3", pyclass(module = "altrios", subclass, eq))]
29pub struct Pyo3VecLocoWrapper(pub Vec<Locomotive>);
31
32#[pyo3_api]
33impl Pyo3VecLocoWrapper {
34 #[new]
35 fn __new__(v: Vec<Locomotive>) -> Self {
37 Self(v)
38 }
39}
40
41impl Pyo3VecLocoWrapper {
42 pub fn new(value: Vec<Locomotive>) -> Self {
43 Self(value)
44 }
45}
46
47impl Init for Pyo3VecLocoWrapper {
48 fn init(&mut self) -> Result<(), Error> {
49 self.0.iter_mut().try_for_each(|l| l.init())?;
50 Ok(())
51 }
52}
53impl SerdeAPI for Pyo3VecLocoWrapper {}
54
55pub trait SolvePower {
56 fn solve_positive_traction(
58 &mut self,
59 loco_vec: &[Locomotive],
60 state: &ConsistState,
61 train_mass: Option<si::Mass>,
62 train_speed: Option<si::Velocity>,
63 ) -> anyhow::Result<Vec<si::Power>>;
64 fn solve_negative_traction(
65 &mut self,
66 loco_vec: &[Locomotive],
67 state: &ConsistState,
68 train_mass: Option<si::Mass>,
69 train_speed: Option<si::Velocity>,
70 ) -> anyhow::Result<Vec<si::Power>>;
71}
72
73#[derive(PartialEq, Eq, Clone, Deserialize, Serialize, Debug)]
74pub struct RESGreedy;
78impl SolvePower for RESGreedy {
79 fn solve_positive_traction(
80 &mut self,
81 loco_vec: &[Locomotive],
82 state: &ConsistState,
83 _train_mass: Option<si::Mass>,
84 _train_speed: Option<si::Velocity>,
85 ) -> anyhow::Result<Vec<si::Power>> {
86 let loco_pwr_out_vec: Vec<si::Power> = if *state
87 .pwr_out_deficit
88 .get_fresh(|| format_dbg!())?
89 == si::Power::ZERO
90 {
91 let mut loco_pwr_out_vec: Vec<si::Power> = vec![];
93 for loco in loco_vec {
94 loco_pwr_out_vec.push(match &loco.loco_type {
95 PowertrainType::ConventionalLoco(_) => si::Power::ZERO,
96 PowertrainType::HybridLoco(_) => {
97 *loco.state.pwr_out_max.get_fresh(|| format_dbg!())?
98 / *state.pwr_out_max_reves.get_fresh(|| format_dbg!())?
99 * *state.pwr_out_req.get_fresh(|| format_dbg!())?
100 }
101 PowertrainType::BatteryElectricLoco(_) => {
102 *loco.state.pwr_out_max.get_fresh(|| format_dbg!())?
103 / *state.pwr_out_max_reves.get_fresh(|| format_dbg!())?
104 * *state.pwr_out_req.get_fresh(|| format_dbg!())?
105 }
106 PowertrainType::DummyLoco(_) => {
109 *state.pwr_out_req.get_fresh(|| format_dbg!())?
110 }
111 })
112 }
113 loco_pwr_out_vec
114 } else {
115 let mut loco_pwr_out_vec: Vec<si::Power> = vec![];
117 for loco in loco_vec {
118 loco_pwr_out_vec.push( match &loco.loco_type {
119 PowertrainType::ConventionalLoco(_) => {
120* loco.state.pwr_out_max.get_fresh(|| format_dbg!())? / *state.pwr_out_max_non_reves.get_fresh(|| format_dbg!())?
121 * *state.pwr_out_deficit.get_fresh(|| format_dbg!())?
122 }
123 PowertrainType::HybridLoco(_) => *loco.state.pwr_out_max.get_fresh(|| format_dbg!())?,
124 PowertrainType::BatteryElectricLoco(_) => *loco.state.pwr_out_max.get_fresh(|| format_dbg!())?,
125 PowertrainType::DummyLoco(_) => {
126 si::Power::ZERO }
128 })
129 }
130 loco_pwr_out_vec
131 };
132 let loco_pwr_out_vec_sum: si::Power = loco_pwr_out_vec.iter().copied().sum();
133 ensure!(
134 utils::almost_eq_uom(
135 &loco_pwr_out_vec.iter().copied().sum(),
136 state.pwr_out_req.get_fresh(|| format_dbg!())?,
137 None,
138 ),
139 format!(
140 "{}\n{}",
141 format_dbg!(loco_pwr_out_vec_sum.get::<si::kilowatt>()),
142 format_dbg!(state
143 .pwr_out_req
144 .get_fresh(|| format_dbg!())?
145 .get::<si::kilowatt>())
146 )
147 );
148 Ok(loco_pwr_out_vec)
149 }
150
151 fn solve_negative_traction(
152 &mut self,
153 loco_vec: &[Locomotive],
154 state: &ConsistState,
155 train_mass: Option<si::Mass>,
156 train_speed: Option<si::Velocity>,
157 ) -> anyhow::Result<Vec<si::Power>> {
158 solve_negative_traction(loco_vec, state, train_mass, train_speed)
159 }
160}
161
162fn get_pwr_regen_vec(
163 loco_vec: &[Locomotive],
164 regen_frac: si::Ratio,
165) -> anyhow::Result<Vec<si::Power>> {
166 let mut pwr_regen_vec: Vec<si::Power> = vec![];
167 for loco in loco_vec {
168 pwr_regen_vec.push(match &loco.loco_type {
169 PowertrainType::ConventionalLoco(_) => si::Power::ZERO,
171 PowertrainType::HybridLoco(_) => {
172 *loco.state.pwr_regen_max.get_fresh(|| format_dbg!())? * regen_frac
173 }
174 PowertrainType::BatteryElectricLoco(_) => {
175 *loco.state.pwr_regen_max.get_fresh(|| format_dbg!())? * regen_frac
176 }
177 PowertrainType::DummyLoco(_) => si::Power::ZERO,
180 })
181 }
182 Ok(pwr_regen_vec)
183}
184
185fn solve_negative_traction(
188 loco_vec: &[Locomotive],
189 consist_state: &ConsistState,
190 _train_mass: Option<si::Mass>,
191 _train_speed: Option<si::Velocity>,
192) -> anyhow::Result<Vec<si::Power>> {
193 let pwr_brake_req = -*consist_state.pwr_out_req.get_fresh(|| format_dbg!())?;
195
196 let regen_frac = if *consist_state.pwr_regen_max.get_fresh(|| format_dbg!())? == si::Power::ZERO
198 {
199 si::Ratio::ZERO
201 } else {
202 (pwr_brake_req / *consist_state.pwr_regen_max.get_fresh(|| format_dbg!())?).min(uc::R * 1.)
203 };
204 let pwr_out_vec: Vec<si::Power> = if *consist_state
205 .pwr_regen_deficit
206 .get_fresh(|| format_dbg!())?
207 == si::Power::ZERO
208 {
209 get_pwr_regen_vec(loco_vec, regen_frac)?
210 } else {
211 let pwr_regen_vec = get_pwr_regen_vec(loco_vec, regen_frac)?;
215 let pwr_surplus_vec: Vec<si::Power> = loco_vec
217 .iter()
218 .zip(&pwr_regen_vec)
219 .map(|(loco, pwr_regen)| {
220 loco.electric_drivetrain()
221 .expect("this `expect` might cause problems for DummyLoco")
222 .pwr_out_max
223 - *pwr_regen
224 })
225 .collect();
226 let pwr_surplus_sum = pwr_surplus_vec
227 .iter()
228 .fold(0.0 * uc::W, |acc, &curr| acc + curr);
229
230 let surplus_frac = *consist_state
232 .pwr_regen_deficit
233 .get_fresh(|| format_dbg!())?
234 / pwr_surplus_sum;
235 ensure!(
236 surplus_frac >= si::Ratio::ZERO && surplus_frac <= uc::R,
237 format_dbg!(surplus_frac),
238 );
239 let pwr_dyn_brake_vec: Vec<si::Power> = pwr_surplus_vec
241 .iter()
242 .zip(pwr_regen_vec)
243 .map(|(pwr_surplus, pwr_regen)| *pwr_surplus * surplus_frac + pwr_regen)
244 .collect();
245 pwr_dyn_brake_vec
246 };
247 let pwr_out_vec: Vec<si::Power> = pwr_out_vec.iter().map(|x| -*x).collect();
249 Ok(pwr_out_vec)
250}
251
252#[derive(PartialEq, Eq, Clone, Deserialize, Serialize, Debug)]
253pub struct Proportional;
258#[allow(unused_variables)]
259#[allow(unreachable_code)]
260impl SolvePower for Proportional {
261 fn solve_positive_traction(
262 &mut self,
263 loco_vec: &[Locomotive],
264 state: &ConsistState,
265 _train_mass: Option<si::Mass>,
266 _train_speed: Option<si::Velocity>,
267 ) -> anyhow::Result<Vec<si::Power>> {
268 todo!("this need some attention to make sure it handles the hybrid correctly");
269 let mut loco_pwr_vec: Vec<si::Power> = vec![];
270 for loco in loco_vec {
271 loco_pwr_vec.push(
272 *loco.state.pwr_out_max.get_fresh(|| format_dbg!())?
274 / *state.pwr_out_max.get_fresh(|| format_dbg!())?
275 * *state.pwr_out_req.get_fresh(|| format_dbg!())?,
276 )
277 }
278 Ok(loco_pwr_vec)
279 }
280
281 fn solve_negative_traction(
282 &mut self,
283 loco_vec: &[Locomotive],
284 state: &ConsistState,
285 train_mass: Option<si::Mass>,
286 train_speed: Option<si::Velocity>,
287 ) -> anyhow::Result<Vec<si::Power>> {
288 todo!("this need some attention to make sure it handles the hybrid correctly");
289 solve_negative_traction(loco_vec, state, train_mass, train_speed)
290 }
291}
292
293#[derive(PartialEq, Eq, Clone, Deserialize, Serialize, Debug)]
294pub struct FrontAndBack;
296impl SerdeAPI for FrontAndBack {}
297impl Init for FrontAndBack {}
298impl SolvePower for FrontAndBack {
299 fn solve_positive_traction(
300 &mut self,
301 _loco_vec: &[Locomotive],
302 _state: &ConsistState,
303 _train_mass: Option<si::Mass>,
304 _train_speed: Option<si::Velocity>,
305 ) -> anyhow::Result<Vec<si::Power>> {
306 todo!() }
308
309 fn solve_negative_traction(
310 &mut self,
311 _loco_vec: &[Locomotive],
312 _state: &ConsistState,
313 _train_mass: Option<si::Mass>,
314 _train_speed: Option<si::Velocity>,
315 ) -> anyhow::Result<Vec<si::Power>> {
316 todo!() }
318}
319
320#[derive(PartialEq, Clone, Deserialize, Serialize, Debug)]
323pub enum PowerDistributionControlType {
324 RESGreedy(RESGreedy),
325 Proportional(Proportional),
326 FrontAndBack(FrontAndBack),
327}
328
329impl SolvePower for PowerDistributionControlType {
330 fn solve_negative_traction(
331 &mut self,
332 loco_vec: &[Locomotive],
333 state: &ConsistState,
334 train_mass: Option<si::Mass>,
335 train_speed: Option<si::Velocity>,
336 ) -> anyhow::Result<Vec<si::Power>> {
337 match self {
338 Self::RESGreedy(res_greedy) => {
339 res_greedy.solve_negative_traction(loco_vec, state, train_mass, train_speed)
340 }
341 Self::Proportional(prop) => {
342 prop.solve_negative_traction(loco_vec, state, train_mass, train_speed)
343 }
344 Self::FrontAndBack(fab) => {
345 fab.solve_negative_traction(loco_vec, state, train_mass, train_speed)
346 }
347 }
348 }
349
350 fn solve_positive_traction(
351 &mut self,
352 loco_vec: &[Locomotive],
353 state: &ConsistState,
354 train_mass: Option<si::Mass>,
355 train_speed: Option<si::Velocity>,
356 ) -> anyhow::Result<Vec<si::Power>> {
357 match self {
358 Self::RESGreedy(res_greedy) => {
359 res_greedy.solve_positive_traction(loco_vec, state, train_mass, train_speed)
360 }
361 Self::Proportional(prop) => {
362 prop.solve_positive_traction(loco_vec, state, train_mass, train_speed)
363 }
364 Self::FrontAndBack(fab) => {
365 fab.solve_positive_traction(loco_vec, state, train_mass, train_speed)
366 }
367 }
368 }
369}
370
371impl Default for PowerDistributionControlType {
372 fn default() -> Self {
373 Self::RESGreedy(RESGreedy)
374 }
375}
376
377impl Init for PowerDistributionControlType {}
378impl SerdeAPI for PowerDistributionControlType {}