fix: ensure shaft speed is set correctly after root-finding

Author: frodehk

src/libecalc/domain/process/compressor/core/train/compressor_train_common_shaft.py

@@ -453,6 +453,7 @@ def _calculate_train_result_given_speed_at_stone_wall(
                     func=lambda x: _calculate_train_result_given_speed_at_stone_wall(speed=x).discharge_pressure
                     - target_discharge_pressure,
                 )
+                self.shaft.set_speed(result_speed)
                 compressor_train_result = _calculate_train_result_given_speed_at_stone_wall(speed=result_speed)
 
                 rate_to_return = compressor_train_result.mass_rate_kg_per_hour * (1 - RATE_CALCULATION_TOLERANCE)
@@ -1024,6 +1025,7 @@ def _calculate_compressor_train(_speed: float) -> CompressorTrainResultSingleTim
 
         if not train_result_for_maximum_speed.within_capacity:
             # will not find valid result - the rate is above maximum rate, return invalid results at maximum speed
+            self.shaft.set_speed(maximum_speed)
             return maximum_speed
         if not train_result_for_minimum_speed.within_capacity:
             # rate is above maximum rate for minimum speed. Find the lowest minimum speed which gives a valid result
@@ -1048,15 +1050,17 @@ def _calculate_compressor_train(_speed: float) -> CompressorTrainResultSingleTim
                 upper_bound=maximum_speed,
                 func=lambda x: _calculate_compressor_train(_speed=x).discharge_pressure - target_discharge_pressure,
             )
-
+            self.shaft.set_speed(speed)
             return speed
 
         # Solution 2, target pressure is too low:
         if (
             constraints.discharge_pressure is not None
             and constraints.discharge_pressure < train_result_for_minimum_speed.discharge_pressure
         ):
+            self.shaft.set_speed(minimum_speed)
             return minimum_speed
 
         # Solution 3, target discharge pressure is too high
+        self.shaft.set_speed(maximum_speed)
         return maximum_speed

src/libecalc/domain/process/compressor/core/train/compressor_train_common_shaft_multiple_streams_and_pressures.py

@@ -589,13 +589,12 @@ def find_and_calculate_for_compressor_train_with_two_pressure_requirements(
             stream_rates=std_rates_last_part,
         )
 
-        compressor_train_first_part.find_fixed_shaft_speed_given_constraints(
+        # Get the speed found for the first part - will be used to compare with the last part
+        shaft_speed_first_part = compressor_train_first_part.find_fixed_shaft_speed_given_constraints(
             constraints=constraints_first_part,
             lower_bound_for_speed=self.minimum_speed,  # Only search for a solution within the bounds of the
             upper_bound_for_speed=self.maximum_speed,  # original, complete compressor train
         )
-        # Get the speed found for the first part - will be used to compare with the last part
-        shaft_speed_first_part = compressor_train_first_part.shaft.get_speed()
 
         if math.isclose(shaft_speed_first_part, self.minimum_speed, rel_tol=EPSILON):
             compressor_train_results_first_part_with_optimal_speed_result = (
@@ -625,14 +624,13 @@ def find_and_calculate_for_compressor_train_with_two_pressure_requirements(
         assert isinstance(compressor_train_last_part._fluid_model, list)  # for mypy
         compressor_train_last_part._fluid_model[0] = compressor_train_last_part.streams[0].fluid_model
 
-        compressor_train_last_part.find_fixed_shaft_speed_given_constraints(
+        shaft_speed_last_part = compressor_train_last_part.find_fixed_shaft_speed_given_constraints(
             constraints=constraints_last_part,
             lower_bound_for_speed=self.minimum_speed,
             upper_bound_for_speed=self.maximum_speed,
         )
-        # Get the speed found for the last part - will be used to compare with the first part
-        shaft_speed_last_part = compressor_train_last_part.shaft.get_speed()
 
+        # Get the speed found for the last part - will be used to compare with the first part
         if math.isclose(shaft_speed_last_part, self.minimum_speed, rel_tol=EPSILON):
             compressor_train_results_last_part_with_optimal_speed_result = (
                 compressor_train_last_part.evaluate_with_pressure_control_given_constraints(

tests/libecalc/core/models/compressor_modelling/test_compressor_train_multiple_streams.py

@@ -514,7 +514,7 @@ def test_evaluate_variable_speed_compressor_train_multiple_streams_and_pressures
 
     result_first_stage = result.stage_results[0]
     result_last_stage = result.stage_results[1]
-    assert result_first_stage.speed == pytest.approx([10825.72, 11100.26, 11313.19], abs=0.01)
+    assert result_first_stage.speed == pytest.approx([10825.77, 11100.32, 11313.25], abs=0.01)
     assert result_first_stage.outlet_stream_condition.pressure == pytest.approx([30.0, 30.0, 30.0], abs=0.01)
     assert result_last_stage.asv_recirculation_loss_mw == pytest.approx([4.22, 4.39, 4.46], abs=0.01)
     assert result_first_stage.speed == result_last_stage.speed