editoast: stdcm fix train simulation

editoast/src/core/simulation.rs

@@ -82,7 +82,12 @@ pub struct PhysicsConsistParameters {
 }
 
 impl PhysicsConsistParameters {
-    pub fn with_traction_engine(traction_engine: RollingStock) -> Self {
+    pub fn set_traction_engine(mut self, traction_engine: RollingStock) -> Self {
+        self.traction_engine = traction_engine;
+        self
+    }
+
+    pub fn from_traction_engine(traction_engine: RollingStock) -> Self {
         PhysicsConsistParameters {
             max_speed: None,
             total_length: None,

editoast/src/views/timetable/stdcm.rs

@@ -42,8 +42,7 @@ use crate::models::train_schedule::TrainSchedule;
 use crate::models::Infra;
 use crate::models::RollingStockModel;
 use crate::views::path::pathfinding::PathfindingResult;
-use crate::views::train_schedule::train_simulation;
-use crate::views::train_schedule::train_simulation_batch;
+use crate::views::train_schedule::consist_train_simulation_batch;
 use crate::views::AuthenticationExt;
 use crate::views::AuthorizationError;
 use crate::AppState;
@@ -91,6 +90,8 @@ enum StdcmError {
     RollingStockNotFound { rolling_stock_id: i64 },
     #[error("Towed rolling stock {towed_rolling_stock_id} does not exist")]
     TowedRollingStockNotFound { towed_rolling_stock_id: i64 },
+    #[error("Train simulation fail")]
+    TrainSimulationFail,
     #[error("Path items are invalid")]
     InvalidPathItems { items: Vec<InvalidPathItem> },
 }
@@ -163,6 +164,17 @@ async fn stdcm(
         })
         .await?;
 
+    let physics_consist_parameters = PhysicsConsistParameters {
+        max_speed: stdcm_request.max_speed,
+        total_length: stdcm_request.total_length,
+        total_mass: stdcm_request.total_mass,
+        towed_rolling_stock: stdcm_request
+            .get_towed_rolling_stock(conn)
+            .await?
+            .map(From::from),
+        traction_engine: rolling_stock.clone().into(),
+    };
+
     // 2. Compute the earliest start time and maximum departure delay
     let virtual_train_run = VirtualTrainRun::simulate(
         db_pool.clone(),
@@ -171,6 +183,7 @@ async fn stdcm(
         &stdcm_request,
         &infra,
         &rolling_stock,
+        &physics_consist_parameters,
         timetable_id,
     )
     .await?;
@@ -186,13 +199,14 @@ async fn stdcm(
     let latest_simulation_end = stdcm_request.get_latest_simulation_end(simulation_run_time);
 
     // 3. Get scheduled train requirements
-    let simulations: Vec<_> = train_simulation_batch(
+    let simulations: Vec<_> = consist_train_simulation_batch(
         conn,
         valkey_client.clone(),
         core_client.clone(),
         &train_schedules,
         &infra,
         stdcm_request.electrical_profile_set_id,
+        physics_consist_parameters.clone(),
     )
     .await?
     .into_iter()
@@ -217,17 +231,7 @@ async fn stdcm(
         rolling_stock_supported_signaling_systems: rolling_stock
             .supported_signaling_systems
             .clone(),
-        rolling_stock: PhysicsConsistParameters {
-            max_speed: stdcm_request.max_speed,
-            total_length: stdcm_request.total_length,
-            total_mass: stdcm_request.total_mass,
-            towed_rolling_stock: stdcm_request
-                .get_towed_rolling_stock(conn)
-                .await?
-                .map(From::from),
-            traction_engine: rolling_stock.into(),
-        }
-        .into(),
+        rolling_stock: physics_consist_parameters.into(),
         temporary_speed_limits: stdcm_request
             .get_temporary_speed_limits(conn, simulation_run_time)
             .await?,
@@ -382,6 +386,7 @@ impl VirtualTrainRun {
         stdcm_request: &Request,
         infra: &Infra,
         rolling_stock: &RollingStockModel,
+        consist_parameters: &PhysicsConsistParameters,
         timetable_id: i64,
     ) -> Result<Self> {
         // Doesn't matter for now, but eventually it will affect tmp speed limits
@@ -415,15 +420,19 @@ impl VirtualTrainRun {
             options: Default::default(),
         };
 
-        let (simulation, pathfinding) = train_simulation(
+        // Compute simulation of a train schedule
+        let (simulation, pathfinding) = consist_train_simulation_batch(
             &mut db_pool.get().await?,
             valkey_client,
             core_client,
-            train_schedule.clone(),
+            &[train_schedule.clone()],
             infra,
             None,
+            consist_parameters.clone(),
         )
-        .await?;
+        .await?
+        .pop()
+        .ok_or(StdcmError::TrainSimulationFail)?;
 
         Ok(Self {
             train_schedule,
@@ -543,7 +552,7 @@ mod tests {
     #[test]
     fn simulation_without_parameters() {
         let rolling_stock = create_simple_rolling_stock();
-        let simulation_parameters = PhysicsConsistParameters::with_traction_engine(rolling_stock);
+        let simulation_parameters = PhysicsConsistParameters::from_traction_engine(rolling_stock);
 
         let physics_consist: PhysicsConsist = simulation_parameters.into();
 

editoast/src/views/train_schedule.rs

@@ -55,6 +55,7 @@ use crate::views::AuthorizationError;
 use crate::AppState;
 use crate::RollingStockModel;
 use crate::ValkeyClient;
+use crate::ValkeyConnection;
 use editoast_derive::EditoastError;
 use editoast_models::DbConnection;
 use editoast_models::DbConnectionPoolV2;
@@ -358,66 +359,40 @@ async fn simulation(
         })
         .await?;
 
-    Ok(Json(
-        train_simulation(
-            &mut db_pool.get().await?,
-            valkey_client,
-            core_client,
-            train_schedule,
-            &infra,
-            electrical_profile_set_id,
-        )
-        .await?
-        .0,
-    ))
-}
-
-/// Compute simulation of a train schedule
-pub async fn train_simulation(
-    conn: &mut DbConnection,
-    valkey_client: Arc<ValkeyClient>,
-    core: Arc<CoreClient>,
-    train_schedule: TrainSchedule,
-    infra: &Infra,
-    electrical_profile_set_id: Option<i64>,
-) -> Result<(SimulationResponse, PathfindingResult)> {
-    Ok(train_simulation_batch(
-        conn,
+    // Compute simulation of a train schedule
+    let (simulation, _) = train_simulation_batch(
+        &mut db_pool.get().await?,
         valkey_client,
-        core,
+        core_client,
         &[train_schedule],
-        infra,
+        &infra,
         electrical_profile_set_id,
     )
     .await?
     .pop()
-    .unwrap())
+    .unwrap(); // TODO handle unwrap ?
+
+    Ok(Json(simulation))
 }
 
-/// Compute in batch the simulation of a list of train schedule
+/// Compute in batch the simulation of a list of train schedule and consist parameters
 ///
 /// Note: The order of the returned simulations is the same as the order of the train schedules.
-pub async fn train_simulation_batch(
+pub async fn consist_train_simulation_batch(
     conn: &mut DbConnection,
     valkey_client: Arc<ValkeyClient>,
     core: Arc<CoreClient>,
     train_schedules: &[TrainSchedule],
     infra: &Infra,
     electrical_profile_set_id: Option<i64>,
+    physics_consist_parameters: PhysicsConsistParameters,
 ) -> Result<Vec<(SimulationResponse, PathfindingResult)>> {
-    let mut valkey_conn = valkey_client.get_connection().await?;
     // Compute path
-    let (rolling_stocks, _): (Vec<_>, _) = RollingStockModel::retrieve_batch(
-        conn,
-        train_schedules
-            .iter()
-            .map::<String, _>(|t| t.rolling_stock_name.clone()),
-    )
-    .await?;
-    let rolling_stocks: HashMap<_, _> = rolling_stocks
-        .into_iter()
-        .map(|rs| (rs.name.clone(), rs))
-        .collect();
+    let rolling_stocks =
+        retrieve_rolling_stocks_from_train_schedules(conn, train_schedules).await?;
+
+    let mut valkey_conn = valkey_client.get_connection().await?;
+
     let pathfinding_results = pathfinding_from_train_batch(
         conn,
         &mut valkey_conn,
@@ -428,6 +403,47 @@ pub async fn train_simulation_batch(
     )
     .await?;
 
+    get_train_batch_simulation(
+        &mut valkey_conn,
+        core.clone(),
+        infra,
+        train_schedules,
+        rolling_stocks,
+        pathfinding_results,
+        electrical_profile_set_id,
+        Some(physics_consist_parameters),
+    )
+    .await
+}
+
+async fn retrieve_rolling_stocks_from_train_schedules(
+    conn: &mut DbConnection,
+    train_schedules: &[TrainSchedule],
+) -> Result<HashMap<String, RollingStockModel>> {
+    let rolling_stocks_ids = train_schedules
+        .iter()
+        .map::<String, _>(|t| t.rolling_stock_name.clone());
+    let (rolling_stocks, _): (Vec<_>, _) =
+        RollingStockModel::retrieve_batch(conn, rolling_stocks_ids).await?;
+
+    let rolling_stocks: HashMap<_, _> = rolling_stocks
+        .into_iter()
+        .map(|rs| (rs.name.clone(), rs))
+        .collect();
+
+    Ok(rolling_stocks)
+}
+
+async fn get_train_batch_simulation(
+    valkey_conn: &mut ValkeyConnection,
+    core: Arc<CoreClient>,
+    infra: &Infra,
+    train_schedules: &[TrainSchedule],
+    rolling_stocks: HashMap<String, RollingStockModel>,
+    pathfinding_results: Vec<PathfindingResult>,
+    electrical_profile_set_id: Option<i64>,
+    physics_consist_parameters: Option<PhysicsConsistParameters>,
+) -> Result<Vec<(SimulationResponse, PathfindingResult)>> {
     let mut simulation_results = vec![SimulationResponse::default(); train_schedules.len()];
     let mut to_sim = Vec::with_capacity(train_schedules.len());
     for (index, (pathfinding, train_schedule)) in
@@ -459,13 +475,21 @@ pub async fn train_simulation_batch(
 
         // Build simulation request
         let rolling_stock = rolling_stocks[&train_schedule.rolling_stock_name].clone();
+
+        let physics_consist_parameters = physics_consist_parameters
+            .clone()
+            .map(|params| params.set_traction_engine(rolling_stock.clone().into()))
+            .unwrap_or_else(|| {
+                PhysicsConsistParameters::from_traction_engine(rolling_stock.into())
+            });
+
         let simulation_request = build_simulation_request(
             infra,
             train_schedule,
             path_item_positions,
             path,
-            rolling_stock,
             electrical_profile_set_id,
+            physics_consist_parameters.clone(),
         );
 
         // Compute unique hash of the simulation input
@@ -524,13 +548,53 @@ pub async fn train_simulation_batch(
         .collect())
 }
 
+/// Compute in batch the simulation of a list of train schedule
+///
+/// Note: The order of the returned simulations is the same as the order of the train schedules.
+pub async fn train_simulation_batch(
+    conn: &mut DbConnection,
+    valkey_client: Arc<ValkeyClient>,
+    core: Arc<CoreClient>,
+    train_schedules: &[TrainSchedule],
+    infra: &Infra,
+    electrical_profile_set_id: Option<i64>,
+) -> Result<Vec<(SimulationResponse, PathfindingResult)>> {
+    // Compute path
+    let rolling_stocks =
+        retrieve_rolling_stocks_from_train_schedules(conn, train_schedules).await?;
+
+    let mut valkey_conn = valkey_client.get_connection().await?;
+
+    let pathfinding_results = pathfinding_from_train_batch(
+        conn,
+        &mut valkey_conn,
+        core.clone(),
+        infra,
+        train_schedules,
+        &rolling_stocks,
+    )
+    .await?;
+
+    get_train_batch_simulation(
+        &mut valkey_conn,
+        core.clone(),
+        infra,
+        train_schedules,
+        rolling_stocks,
+        pathfinding_results,
+        electrical_profile_set_id,
+        None,
+    )
+    .await
+}
+
 fn build_simulation_request(
     infra: &Infra,
     train_schedule: &TrainSchedule,
     path_item_positions: &[u64],
     path: SimulationPath,
-    rolling_stock: RollingStockModel,
     electrical_profile_set_id: Option<i64>,
+    physics_consist_parameters: PhysicsConsistParameters,
 ) -> SimulationRequest {
     assert_eq!(path_item_positions.len(), train_schedule.path.len());
     // Project path items to path offset
@@ -590,7 +654,7 @@ fn build_simulation_request(
         speed_limit_tag: train_schedule.speed_limit_tag.clone(),
         power_restrictions,
         options: train_schedule.options.clone(),
-        rolling_stock: PhysicsConsistParameters::with_traction_engine(rolling_stock.into()).into(),
+        rolling_stock: physics_consist_parameters.into(),
         electrical_profile_set_id,
     }
 }