Merge pull request #49 from NREL/fix/more-data-in-train-history

Fix/more data in train history

applications/calibration/zanzeff_set_speed_train_cal_post.py

@@ -3,7 +3,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
-sns.set()
+sns.set_theme()
 
 from zanzeff_set_speed_train_cal import *
 

python/altrios/altrios_pyo3.pyi

@@ -613,6 +613,7 @@ class SpeedTrace(SerdeAPI):
     def __copy__(self) -> Any: ...
     def __len__(self) -> int: ...
     def from_csv_file(pathstr: str) -> Self: ...
+    def to_csv_file(self, pathstr: str): ...
 
 
 class TrainState:
@@ -621,6 +622,9 @@ class TrainState:
     offset_meters: float
     offset_back_meters: float
     total_dist_meters: float
+    link_idx_front: int
+    offset_in_link_meters: float
+    grade_front: float
     speed_meters_per_second: float
     speed_limit_meters_per_second: float
     speed_target_meters_per_second: float
@@ -668,6 +672,10 @@ class TrainState:
 class TrainStateHistoryVec(SerdeAPI):
     time_seconds: list[float]
     offset_meters: list[float]
+    offset_back_meters: list[float]
+    link_idx_front: list[int]
+    offset_in_link_meters: list[float]
+    grade_front: list[float]
     speed_meters_per_second: list[float]
     speed_limit_meters_per_second: list[float]
     speed_target_meters_per_second: list[float]

python/altrios/demos/bel_demo.py

@@ -11,7 +11,7 @@
 import os
 import seaborn as sns
 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 

python/altrios/demos/conv_demo.py

@@ -11,7 +11,7 @@
 import os
 import seaborn as sns
 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 

python/altrios/demos/rollout_demo.py

@@ -7,7 +7,7 @@
 import seaborn as sns
 from pathlib import Path
 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 

python/altrios/demos/set_speed_simple_corr_demo.py

@@ -10,7 +10,7 @@
 import seaborn as sns
 
 import altrios as alt 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 
@@ -78,62 +78,62 @@
 # pull out solved locomotive for plotting convenience
 loco0:alt.Locomotive = train_sim.loco_con.loco_vec.tolist()[0]
 
+fig, ax = plt.subplots(4, 1, sharex=True)
+ax[0].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.pwr_whl_out_watts,
+    label="tract pwr",
+)
+ax[0].set_ylabel('Power')
+ax[0].legend()
+
+ax[1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.res_aero_newtons,
+    label='aero',
+)
+ax[1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.res_rolling_newtons,
+    label='rolling',
+)
+ax[1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.res_curve_newtons,
+    label='curve',
+)
+ax[1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.res_bearing_newtons,
+    label='bearing',
+)
+ax[1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.history.res_grade_newtons,
+    label='grade',
+)
+ax[1].set_ylabel('Force [N]')
+ax[1].legend()
+
+ax[-1].plot(
+    np.array(train_sim.history.time_seconds) / 3_600,
+    train_sim.speed_trace.speed_meters_per_second,
+)
+ax[-1].set_xlabel('Time [hr]')
+ax[-1].set_ylabel('Speed [m/s]')
+
+ax[2].plot(
+    np.array(train_sim.history.time_seconds) / 3_600, 
+    np.array(loco0.res.history.soc)
+)
+
+ax[2].set_ylabel('SOC')
+ax[2].legend()
+
+plt.suptitle("Set Speed Train Sim Demo")
+plt.tight_layout()
+
 if SHOW_PLOTS:
-    fig, ax = plt.subplots(4, 1, sharex=True)
-    ax[0].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.pwr_whl_out_watts,
-        label="tract pwr",
-    )
-    ax[0].set_ylabel('Power')
-    ax[0].legend()
-
-    ax[1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.res_aero_newtons,
-        label='aero',
-    )
-    ax[1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.res_rolling_newtons,
-        label='rolling',
-    )
-    ax[1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.res_curve_newtons,
-        label='curve',
-    )
-    ax[1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.res_bearing_newtons,
-        label='bearing',
-    )
-    ax[1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.history.res_grade_newtons,
-        label='grade',
-    )
-    ax[1].set_ylabel('Force [N]')
-    ax[1].legend()
-
-    ax[-1].plot(
-        np.array(train_sim.history.time_seconds) / 3_600,
-        train_sim.speed_trace.speed_meters_per_second,
-    )
-    ax[-1].set_xlabel('Time [hr]')
-    ax[-1].set_ylabel('Speed [m/s]')
-
-    ax[2].plot(
-        np.array(train_sim.history.time_seconds) / 3_600, 
-        np.array(loco0.res.history.soc)
-    )
-
-    ax[2].set_ylabel('SOC')
-    ax[2].legend()
-
-    plt.suptitle("Set Speed Train Sim Demo")
-
-    plt.tight_layout()
-    plt.show()
+    fig.show()
 
 # %%

python/altrios/demos/set_speed_train_sim_demo.py

@@ -6,7 +6,7 @@
 import seaborn as sns
 
 import altrios as alt 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 
@@ -25,8 +25,7 @@
 # instantiate battery model
 # https://docs.rs/altrios-core/latest/altrios_core/consist/locomotive/powertrain/reversible_energy_storage/struct.ReversibleEnergyStorage.html#
 res = alt.ReversibleEnergyStorage.from_file(
-    alt.resources_root() / 
-        "powertrains/reversible_energy_storages/Kokam_NMC_75Ah_flx_drive.yaml"
+    alt.resources_root() / "powertrains/reversible_energy_storages/Kokam_NMC_75Ah_flx_drive.yaml"
 )
 # instantiate electric drivetrain (motors and any gearboxes)
 # https://docs.rs/altrios-core/latest/altrios_core/consist/locomotive/powertrain/electric_drivetrain/struct.ElectricDrivetrain.html
@@ -61,18 +60,15 @@
     loco_con=loco_con,
 )
 
-rail_vehicle_file = "rolling_stock/rail_vehicles.csv"
-rail_vehicle_map = alt.import_rail_vehicles(alt.resources_root() / rail_vehicle_file)
-rail_vehicle = rail_vehicle_map[train_config.rail_vehicle_type]
-
-network = alt.Network.from_file(alt.resources_root() / "networks/Taconite.yaml")
-# This data in this file were generated by running 
-# ```python
-# [lp.link_idx.idx for lp in sim0.path_tpc.link_points]
-# ``` 
-# in sim_manager_demo.py.
-link_path = alt.LinkPath.from_csv_file(alt.resources_root() / "demo_data/link_points_idx.csv")
+rail_vehicle_file = "rolling_stock/" + train_config.rail_vehicle_type + ".yaml"
+rail_vehicle = alt.RailVehicle.from_file(
+    alt.resources_root() / rail_vehicle_file)
 
+network = alt.Network.from_file(
+    alt.resources_root() / "networks/Taconite.yaml")
+link_path = alt.LinkPath.from_csv_file(
+    alt.resources_root() / "demo_data/link_points_idx.csv"
+)
 
 speed_trace = alt.SpeedTrace.from_csv_file(
     alt.resources_root() / "demo_data/speed_trace.csv"
@@ -86,7 +82,7 @@
     save_interval=SAVE_INTERVAL,
 )
 
-train_sim.set_save_interval(1)
+train_sim.set_save_interval(SAVE_INTERVAL)
 t0 = time.perf_counter()
 train_sim.walk()
 t1 = time.perf_counter()

python/altrios/demos/sim_manager_demo.py

@@ -8,7 +8,7 @@
 import seaborn as sns
 from pathlib import Path
 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()
 # %

python/altrios/demos/speed_limit_simple_corr_demo.py

@@ -10,7 +10,7 @@
 import seaborn as sns
 
 import altrios as alt 
-sns.set()
+sns.set_theme()
 
 SHOW_PLOTS = alt.utils.show_plots()