|
| 1 | +import time |
| 2 | +import altrios as alt |
| 3 | + |
| 4 | +SAVE_INTERVAL = 100 |
| 5 | +def get_solved_speed_limit_train_sim(): |
| 6 | + # Build the train config |
| 7 | + rail_vehicle_loaded = alt.RailVehicle.from_file( |
| 8 | + alt.resources_root() / "rolling_stock/Manifest_Loaded.yaml") |
| 9 | + rail_vehicle_empty = alt.RailVehicle.from_file( |
| 10 | + alt.resources_root() / "rolling_stock/Manifest_Empty.yaml") |
| 11 | + |
| 12 | + # https://docs.rs/altrios-core/latest/altrios_core/train/struct.TrainConfig.html |
| 13 | + train_config = alt.TrainConfig( |
| 14 | + rail_vehicles=[rail_vehicle_loaded, rail_vehicle_empty], |
| 15 | + n_cars_by_type={ |
| 16 | + "Manifest_Loaded": 50, |
| 17 | + "Manifest_Empty": 50, |
| 18 | + }, |
| 19 | + train_length_meters=None, |
| 20 | + train_mass_kilograms=None, |
| 21 | + ) |
| 22 | + |
| 23 | + # Build the locomotive consist model |
| 24 | + # instantiate battery model |
| 25 | + # https://docs.rs/altrios-core/latest/altrios_core/consist/locomotive/powertrain/reversible_energy_storage/struct.ReversibleEnergyStorage.html# |
| 26 | + res = alt.ReversibleEnergyStorage.from_file( |
| 27 | + alt.resources_root() / "powertrains/reversible_energy_storages/Kokam_NMC_75Ah_flx_drive.yaml" |
| 28 | + ) |
| 29 | + |
| 30 | + edrv = alt.ElectricDrivetrain( |
| 31 | + pwr_out_frac_interp=[0., 1.], |
| 32 | + eta_interp=[0.98, 0.98], |
| 33 | + pwr_out_max_watts=5e9, |
| 34 | + save_interval=SAVE_INTERVAL, |
| 35 | + ) |
| 36 | + |
| 37 | + bel: alt.Locomotive = alt.Locomotive.build_battery_electric_loco( |
| 38 | + reversible_energy_storage=res, |
| 39 | + drivetrain=edrv, |
| 40 | + loco_params=alt.LocoParams.from_dict(dict( |
| 41 | + pwr_aux_offset_watts=8.55e3, |
| 42 | + pwr_aux_traction_coeff_ratio=540.e-6, |
| 43 | + force_max_newtons=667.2e3, |
| 44 | + ))) |
| 45 | + |
| 46 | + # construct a vector of one BEL and several conventional locomotives |
| 47 | + loco_vec = [bel.clone()] + [alt.Locomotive.default()] * 7 |
| 48 | + # instantiate consist |
| 49 | + loco_con = alt.Consist( |
| 50 | + loco_vec |
| 51 | + ) |
| 52 | + |
| 53 | + # Instantiate the intermediate `TrainSimBuilder` |
| 54 | + tsb = alt.TrainSimBuilder( |
| 55 | + train_id="0", |
| 56 | + origin_id="A", |
| 57 | + destination_id="B", |
| 58 | + train_config=train_config, |
| 59 | + loco_con=loco_con, |
| 60 | + ) |
| 61 | + |
| 62 | + # Load the network and construct the timed link path through the network. |
| 63 | + network = alt.Network.from_file( |
| 64 | + alt.resources_root() / 'networks/simple_corridor_network.yaml') |
| 65 | + |
| 66 | + location_map = alt.import_locations( |
| 67 | + alt.resources_root() / "networks/simple_corridor_locations.csv") |
| 68 | + train_sim: alt.SetSpeedTrainSim = tsb.make_speed_limit_train_sim( |
| 69 | + location_map=location_map, |
| 70 | + save_interval=1, |
| 71 | + ) |
| 72 | + train_sim.set_save_interval(SAVE_INTERVAL) |
| 73 | + est_time_net, _consist = alt.make_est_times(train_sim, network) |
| 74 | + |
| 75 | + timed_link_path = alt.run_dispatch( |
| 76 | + network, |
| 77 | + alt.SpeedLimitTrainSimVec([train_sim]), |
| 78 | + [est_time_net], |
| 79 | + False, |
| 80 | + False, |
| 81 | + )[0] |
| 82 | + |
| 83 | + train_sim.walk_timed_path( |
| 84 | + network=network, |
| 85 | + timed_path=timed_link_path, |
| 86 | + ) |
| 87 | + assert len(train_sim.history) > 1 |
| 88 | + |
| 89 | + return train_sim |
| 90 | + |
| 91 | + |
| 92 | +def test_pydict(): |
| 93 | + ts = get_solved_speed_limit_train_sim() |
| 94 | + |
| 95 | + t0 = time.perf_counter_ns() |
| 96 | + ts_dict_msg = ts.to_pydict(flatten=False, data_fmt="msg_pack") |
| 97 | + ts_msg = alt.SpeedLimitTrainSim.from_pydict( |
| 98 | + ts_dict_msg, data_fmt="msg_pack") |
| 99 | + t1 = time.perf_counter_ns() |
| 100 | + t_msg = t1 - t0 |
| 101 | + print(f"\nElapsed time for MessagePack: {t_msg:.3e} ns ") |
| 102 | + |
| 103 | + t0 = time.perf_counter_ns() |
| 104 | + ts_dict_yaml = ts.to_pydict(flatten=False, data_fmt="yaml") |
| 105 | + ts_yaml = alt.SpeedLimitTrainSim.from_pydict(ts_dict_yaml, data_fmt="yaml") |
| 106 | + t1 = time.perf_counter_ns() |
| 107 | + t_yaml = t1 - t0 |
| 108 | + print(f"Elapsed time for YAML: {t_yaml:.3e} ns ") |
| 109 | + print(f"YAML time per MessagePack time: {(t_yaml / t_msg):.3e} ") |
| 110 | + |
| 111 | + t0 = time.perf_counter_ns() |
| 112 | + ts_dict_json = ts.to_pydict(flatten=False, data_fmt="json") |
| 113 | + _ts_json = alt.SpeedLimitTrainSim.from_pydict( |
| 114 | + ts_dict_json, data_fmt="json") |
| 115 | + t1 = time.perf_counter_ns() |
| 116 | + t_json = t1 - t0 |
| 117 | + print(f"Elapsed time for json: {t_json:.3e} ns ") |
| 118 | + print(f"JSON time per MessagePack time: {(t_json / t_msg):.3e} ") |
| 119 | + |
| 120 | + # `to_pydict` is necessary because of some funkiness with direct equality comparison |
| 121 | + assert ts_msg.to_pydict() == ts.to_pydict() |
| 122 | + assert ts_yaml.to_pydict() == ts.to_pydict() |
| 123 | + |
| 124 | +if __name__ == "__main__": |
| 125 | + test_pydict() |
0 commit comments