altrios_core/train/
train_state.rs

1use crate::imports::*;
2use crate::track::PathTpc;
3
4#[serde_api]
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, HistoryVec)]
6#[cfg_attr(feature = "pyo3", pyclass(module = "altrios", subclass, eq))]
7/// For `SetSpeedTrainSim`, it is typically best to use the default for this.
8pub struct InitTrainState {
9    pub time: TrackedState<si::Time>,
10    pub offset: TrackedState<si::Length>,
11    pub speed: TrackedState<si::Velocity>,
12}
13
14#[pyo3_api]
15impl InitTrainState {
16    #[new]
17    #[pyo3(signature = (
18        time_seconds=None,
19        offset_meters=None,
20        speed_meters_per_second=None,
21    ))]
22    fn __new__(
23        time_seconds: Option<f64>,
24        offset_meters: Option<f64>,
25        speed_meters_per_second: Option<f64>,
26    ) -> Self {
27        Self::new(
28            time_seconds.map(|x| x * uc::S),
29            offset_meters.map(|x| x * uc::M),
30            speed_meters_per_second.map(|x| x * uc::MPS),
31        )
32    }
33}
34
35impl Init for InitTrainState {}
36impl SerdeAPI for InitTrainState {}
37
38impl Default for InitTrainState {
39    fn default() -> Self {
40        Self {
41            time: TrackedState::new(si::Time::ZERO),
42            offset: TrackedState::new(f64::NAN * uc::M),
43            speed: TrackedState::new(si::Velocity::ZERO),
44        }
45    }
46}
47
48impl InitTrainState {
49    pub fn new(
50        time: Option<si::Time>,
51        offset: Option<si::Length>,
52        speed: Option<si::Velocity>,
53    ) -> Self {
54        let base = InitTrainState::default();
55        Self {
56            time: TrackedState::new(
57                time.unwrap_or(*base.time.get_fresh(|| format_dbg!()).unwrap()),
58            ),
59            offset: TrackedState::new(
60                offset.unwrap_or(*base.offset.get_fresh(|| format_dbg!()).unwrap()),
61            ),
62            speed: TrackedState::new(
63                speed.unwrap_or(*base.speed.get_fresh(|| format_dbg!()).unwrap()),
64            ),
65        }
66    }
67}
68
69#[serde_api]
70#[derive(
71    Debug, Clone, Serialize, Deserialize, HistoryVec, PartialEq, StateMethods, SetCumulative,
72)]
73#[cfg_attr(feature = "pyo3", pyclass(module = "altrios", subclass, eq))]
74pub struct TrainState {
75    /// time since user-defined datum
76    pub time: TrackedState<si::Time>,
77    /// index for time steps
78    pub i: TrackedState<usize>,
79    /// Linear-along-track, directional distance of front of train from original
80    /// starting position of back of train.
81    ///
82    /// If this is provided in [InitTrainState::new], it gets set as the train length or the value,
83    /// whichever is larger, and if it is not provided, then it defaults to the train length.
84    pub offset: TrackedState<si::Length>,
85    /// Linear-along-track, directional distance of back of train from original
86    /// starting position of back of train.
87    pub offset_back: TrackedState<si::Length>,
88    /// Linear-along-track, cumulative, absolute distance from initial starting position.
89    pub total_dist: TrackedState<si::Length>,
90    /// Current link containing head end (i.e. pulling locomotives) of train
91    pub link_idx_front: TrackedState<u32>,
92    /// Current link containing tail/back end of train
93    pub link_idx_back: TrackedState<u32>,
94    /// Offset from start of current link
95    pub offset_in_link: TrackedState<si::Length>,
96    /// Achieved speed based on consist capabilities and train resistance
97    pub speed: TrackedState<si::Velocity>,
98    /// Speed limit
99    pub speed_limit: TrackedState<si::Velocity>,
100    /// Speed target from meet-pass planner
101    pub speed_target: TrackedState<si::Velocity>,
102    /// Time step size
103    pub dt: TrackedState<si::Time>,
104    /// Train length
105    pub length: TrackedState<si::Length>,
106    /// Static mass of train, including freight
107    pub mass_static: TrackedState<si::Mass>,
108    /// Effective additional mass of train due to rotational inertia
109    pub mass_rot: TrackedState<si::Mass>,
110    /// Mass of freight being hauled by the train (not including railcar empty weight)
111    pub mass_freight: TrackedState<si::Mass>,
112    /// Static weight of train
113    pub weight_static: TrackedState<si::Force>,
114    /// Rolling resistance force
115    pub res_rolling: TrackedState<si::Force>,
116    /// Bearing resistance force
117    pub res_bearing: TrackedState<si::Force>,
118    /// Davis B term resistance force
119    pub res_davis_b: TrackedState<si::Force>,
120    /// Aerodynamic resistance force
121    pub res_aero: TrackedState<si::Force>,
122    /// Grade resistance force
123    pub res_grade: TrackedState<si::Force>,
124    /// Curvature resistance force
125    pub res_curve: TrackedState<si::Force>,
126
127    /// Grade at front of train
128    pub grade_front: TrackedState<si::Ratio>,
129    /// Grade at back of train of train if strap method is used
130    pub grade_back: TrackedState<si::Ratio>,
131    /// Elevation at front of train
132    pub elev_front: TrackedState<si::Length>,
133    /// Elevation at back of train
134    pub elev_back: TrackedState<si::Length>,
135
136    /// Power to overcome train resistance forces
137    pub pwr_res: TrackedState<si::Power>,
138    /// Power to overcome inertial forces
139    pub pwr_accel: TrackedState<si::Power>,
140    /// Total tractive power exerted by locomotive consist
141    pub pwr_whl_out: TrackedState<si::Power>,
142    /// Integral of [Self::pwr_whl_out]
143    pub energy_whl_out: TrackedState<si::Energy>,
144    /// Energy out during positive or zero traction
145    pub energy_whl_out_pos: TrackedState<si::Energy>,
146    /// Energy out during negative traction (positive value means negative traction)
147    pub energy_whl_out_neg: TrackedState<si::Energy>,
148}
149
150impl Init for TrainState {}
151impl SerdeAPI for TrainState {}
152
153impl Default for TrainState {
154    fn default() -> Self {
155        Self {
156            time: Default::default(),
157            i: Default::default(),
158            offset: Default::default(),
159            offset_back: Default::default(),
160            total_dist: Default::default(),
161            link_idx_front: Default::default(),
162            link_idx_back: Default::default(),
163            offset_in_link: Default::default(),
164            speed: Default::default(),
165            speed_limit: Default::default(),
166            dt: TrackedState::new(uc::S),
167            length: Default::default(),
168            mass_static: Default::default(),
169            mass_rot: Default::default(),
170            mass_freight: Default::default(),
171            elev_front: Default::default(),
172            elev_back: Default::default(),
173            energy_whl_out: Default::default(),
174            grade_front: Default::default(),
175            grade_back: Default::default(),
176            speed_target: Default::default(),
177            weight_static: Default::default(),
178            res_rolling: Default::default(),
179            res_bearing: Default::default(),
180            res_davis_b: Default::default(),
181            res_aero: Default::default(),
182            res_grade: Default::default(),
183            res_curve: Default::default(),
184            pwr_res: Default::default(),
185            pwr_accel: Default::default(),
186            pwr_whl_out: Default::default(),
187            energy_whl_out_pos: Default::default(),
188            energy_whl_out_neg: Default::default(),
189        }
190    }
191}
192
193impl Mass for TrainState {
194    /// Static mass of train, not including effective rotational mass
195    fn mass(&self) -> anyhow::Result<Option<si::Mass>> {
196        self.derived_mass()
197    }
198
199    fn set_mass(
200        &mut self,
201        _new_mass: Option<si::Mass>,
202        _side_effect: MassSideEffect,
203    ) -> anyhow::Result<()> {
204        bail!("`set_mass` is not enabled for `TrainState`")
205    }
206
207    fn derived_mass(&self) -> anyhow::Result<Option<si::Mass>> {
208        // NOTE: if we ever dynamically change mass, this needs attention!
209        Ok(Some(*self.mass_static.get_unchecked(|| format_dbg!())?))
210    }
211
212    fn expunge_mass_fields(&mut self) {}
213}
214
215impl TrainState {
216    #[allow(clippy::too_many_arguments)]
217    pub fn new(
218        length: si::Length,
219        mass_static: si::Mass,
220        mass_rot: si::Mass,
221        mass_freight: si::Mass,
222        init_train_state: Option<InitTrainState>,
223    ) -> Self {
224        let init_train_state = init_train_state.unwrap_or_default();
225        let offset = init_train_state
226            .offset
227            .get_fresh(|| format_dbg!())
228            .unwrap()
229            .max(length);
230        Self {
231            time: init_train_state.time,
232            i: Default::default(),
233            offset: TrackedState::new(offset),
234            offset_back: TrackedState::new(offset - length),
235            total_dist: TrackedState::new(si::Length::ZERO),
236            speed: init_train_state.speed.clone(),
237            // this needs to be set to something greater than or equal to actual speed and will be
238            // updated after the first time step anyway
239            speed_limit: init_train_state.speed,
240            length: TrackedState::new(length),
241            mass_static: TrackedState::new(mass_static),
242            mass_rot: TrackedState::new(mass_rot),
243            mass_freight: TrackedState::new(mass_freight),
244            ..Self::default()
245        }
246    }
247
248    pub fn res_net(&self) -> anyhow::Result<si::Force> {
249        Ok(*self.res_rolling.get_fresh(|| format_dbg!())?
250            + *self.res_bearing.get_fresh(|| format_dbg!())?
251            + *self.res_davis_b.get_fresh(|| format_dbg!())?
252            + *self.res_aero.get_fresh(|| format_dbg!())?
253            + *self.res_grade.get_fresh(|| format_dbg!())?
254            + *self.res_curve.get_fresh(|| format_dbg!())?)
255    }
256
257    /// All base, freight, and rotational mass
258    pub fn mass_compound(&self) -> anyhow::Result<si::Mass> {
259        Ok(self
260            .mass() 
261            .with_context(|| format_dbg!())? // extract result
262            .with_context(|| format!("{}\nExpected `Some`", format_dbg!()))? // extract option
263            + *self.mass_rot.get_unchecked(|| format_dbg!())?)
264    }
265}
266
267impl Valid for TrainState {
268    fn valid() -> Self {
269        Self {
270            length: TrackedState::new(2000.0 * uc::M),
271            offset: TrackedState::new(2000.0 * uc::M),
272            offset_back: TrackedState::new(si::Length::ZERO),
273            mass_static: TrackedState::new(6000.0 * uc::TON),
274            mass_rot: TrackedState::new(200.0 * uc::TON),
275
276            dt: TrackedState::new(uc::S),
277            ..Self::default()
278        }
279    }
280}
281
282// TODO: Add new values!
283impl ObjState for TrainState {
284    fn validate(&self) -> ValidationResults {
285        let mut errors = ValidationErrors::new();
286        if let Err(err) = self.mass_static.get_fresh(|| format_dbg!()) {
287            errors.push(err);
288            return errors.make_err();
289        }
290        if let Err(err) = self.length.get_fresh(|| format_dbg!()) {
291            errors.push(err);
292            return errors.make_err();
293        }
294        si_chk_num_gtz_fin(
295            &mut errors,
296            self.mass_static.get_fresh(|| format_dbg!()).unwrap(),
297            "Mass static",
298        );
299        si_chk_num_gtz_fin(
300            &mut errors,
301            self.length.get_fresh(|| format_dbg!()).unwrap(),
302            "Length",
303        );
304        // si_chk_num_gtz_fin(&mut errors, &self.res_bearing, "Resistance bearing");
305        // si_chk_num_fin(&mut errors, &self.res_davis_b, "Resistance Davis B");
306        // si_chk_num_gtz_fin(&mut errors, &self.cd_area, "cd area");
307        errors.make_err()
308    }
309}
310
311/// Sets `link_idx_front` and `offset_in_link` based on `state` and `path_tpc`
312///
313/// Assumes that `offset` in `link_points()` is monotically increasing, which may not always be true.
314pub fn set_link_and_offset(state: &mut TrainState, path_tpc: &PathTpc) -> anyhow::Result<()> {
315    // index of current link within `path_tpc`
316    // if the link_point.offset is greater than the train `state` offset, then
317    // the train is in the previous link
318    let offset = *state.offset.get_stale(|| format_dbg!())?;
319    let idx_curr_link = path_tpc
320        .link_points()
321        .iter()
322        .position(|&lp| lp.offset > offset)
323        // if None, assume that it's the last element
324        .unwrap_or_else(|| path_tpc.link_points().len())
325        - 1;
326    let link_point = path_tpc
327        .link_points()
328        .get(idx_curr_link)
329        .with_context(|| format_dbg!())?;
330    state
331        .link_idx_front
332        .update(link_point.link_idx.idx() as u32, || format_dbg!())?;
333    state.offset_in_link.update(
334        *state.offset.get_stale(|| format_dbg!())? - link_point.offset,
335        || format_dbg!(),
336    )?;
337
338    // link index of back of train at current time step
339    let offset_back = *state.offset_back.get_fresh(|| format_dbg!())?;
340    let idx_back_link = path_tpc
341        .link_points()
342        .iter()
343        .position(|&lp| lp.offset > offset_back)
344        // if None, assume that it's the last element
345        .unwrap_or_else(|| path_tpc.link_points().len())
346        - 1;
347    state.link_idx_back.update(
348        path_tpc
349            .link_points()
350            .get(idx_back_link)
351            .with_context(|| format_dbg!())?
352            .link_idx
353            .idx() as u32,
354        || format_dbg!(),
355    )?;
356
357    Ok(())
358}