refactor: unify simplified compressor train models and improve code quality

- Consolidate CompressorTrainSimplifiedKnownStages and
  CompressorTrainSimplifiedUnknownStages into single unified model
- Replace separate energy model types with unified COMPRESSOR_TRAIN_SIMPLIFIED type
- Add KnownStagesConfig and UnknownStagesConfig for flexible stage configuration
- Update factory functions to create unified models with proper configuration
- Improve SimplifiedTrainBuilder API by removing redundant parameters
- Fix unused variables and parameter shadowing issues
- Update consumer functions to use unified model approach
- Maintain backward compatibility through configuration-based stage preparation

Author: kjbrak

src/libecalc/common/energy_model_type.py

@@ -6,8 +6,7 @@ class EnergyModelType(str, Enum):
     TABULATED = "TABULATED"
     COMPRESSOR_SAMPLED = "COMPRESSOR_SAMPLED"
     PUMP_MODEL = "PUMP_MODEL"
-    COMPRESSOR_TRAIN_SIMPLIFIED_WITH_KNOWN_STAGES = "COMPRESSOR_TRAIN_SIMPLIFIED_WITH_KNOWN_NUMBER_OF_COMPRESSORS"
-    COMPRESSOR_TRAIN_SIMPLIFIED_WITH_UNKNOWN_STAGES = "COMPRESSOR_TRAIN_SIMPLIFIED_WITH_UNKNOWN_NUMBER_OF_COMPRESSORS"
+    COMPRESSOR_TRAIN_SIMPLIFIED = "COMPRESSOR_TRAIN_SIMPLIFIED"
     VARIABLE_SPEED_COMPRESSOR_TRAIN_COMMON_SHAFT = "VARIABLE_SPEED_COMPRESSOR_TRAIN_COMMON_SHAFT"
     SINGLE_SPEED_COMPRESSOR_TRAIN_COMMON_SHAFT = "SINGLE_SPEED_COMPRESSOR_TRAIN_COMMON_SHAFT"
     VARIABLE_SPEED_COMPRESSOR_TRAIN_MULTIPLE_STREAMS_AND_PRESSURES = (

src/libecalc/domain/infrastructure/energy_components/legacy_consumer/consumer_function/compressor_consumer_function.py

@@ -5,10 +5,7 @@
     ConsumerFunctionResult,
 )
 from libecalc.domain.process.compressor.core.base import CompressorModel, CompressorWithTurbineModel
-from libecalc.domain.process.compressor.core.train.simplified_train import (
-    CompressorTrainSimplifiedKnownStages,
-    CompressorTrainSimplifiedUnknownStages,
-)
+from libecalc.domain.process.compressor.core.train.simplified_train import CompressorTrainSimplified
 from libecalc.domain.process.compressor.core.train.simplified_train_builder import SimplifiedTrainBuilder
 from libecalc.domain.process.compressor.core.train.variable_speed_compressor_train_common_shaft_multiple_streams_and_pressures import (
     VariableSpeedCompressorTrainCommonShaftMultipleStreamsAndPressures,
@@ -69,28 +66,21 @@ def __init__(
             else None
         )
 
-        # Use existing _compressor_model property to avoid code duplication
-        prepared_stages = None
-        if isinstance(self._compressor_model, CompressorTrainSimplifiedKnownStages):
-            prepared_stages = SimplifiedTrainBuilder.prepare_stages_for_known_stages_model(
-                original_dto_stages=self._compressor_model._original_dto_stages,
-                fluid_factory=self._compressor_model.fluid_factory,
-                rate=stream_day_rate,
-                suction_pressure=suction_pressure,
-                discharge_pressure=discharge_pressure,
-            )
-        elif isinstance(self._compressor_model, CompressorTrainSimplifiedUnknownStages):
-            prepared_stages = SimplifiedTrainBuilder.prepare_stages_for_unknown_stages_model(
-                stage_template=self._compressor_model.stage,
-                maximum_pressure_ratio_per_stage=self._compressor_model.maximum_pressure_ratio_per_stage,
-                fluid_factory=self._compressor_model.fluid_factory,
-                rate=stream_day_rate,
-                suction_pressure=suction_pressure,
-                discharge_pressure=discharge_pressure,
-            )
-
-        if prepared_stages is not None:
-            self._compressor_model.stages = prepared_stages
+        # Prepare stages for simplified compressor train using configuration approach
+        if isinstance(self._compressor_model, CompressorTrainSimplified):
+            if (
+                self._compressor_model._stage_config is not None
+                and suction_pressure is not None
+                and discharge_pressure is not None
+            ):
+                builder = SimplifiedTrainBuilder(fluid_factory=self._compressor_model.fluid_factory)
+                prepared_stages, _ = builder.prepare_stages_from_config(
+                    config=self._compressor_model._stage_config,
+                    rate=stream_day_rate,
+                    suction_pressure=suction_pressure,
+                    discharge_pressure=discharge_pressure,
+                )
+                self._compressor_model.set_prepared_stages(prepared_stages)
 
         compressor_function.set_evaluation_input(
             rate=stream_day_rate,

src/libecalc/domain/infrastructure/energy_components/legacy_consumer/system/consumer_function.py

@@ -27,10 +27,7 @@
     get_operational_settings_number_used_from_model_results,
 )
 from libecalc.domain.process.compressor.core.base import CompressorModel, CompressorWithTurbineModel
-from libecalc.domain.process.compressor.core.train.simplified_train import (
-    CompressorTrainSimplifiedKnownStages,
-    CompressorTrainSimplifiedUnknownStages,
-)
+from libecalc.domain.process.compressor.core.train.simplified_train import CompressorTrainSimplified
 from libecalc.domain.process.compressor.core.train.simplified_train_builder import SimplifiedTrainBuilder
 from libecalc.domain.process.pump.pump import PumpModel
 from libecalc.domain.time_series_power_loss_factor import TimeSeriesPowerLossFactor
@@ -351,27 +348,17 @@ def evaluate_consumers(  # type: ignore[override]
                 actual_model = self._get_actual_compressor_model(consumer_model)
             else:
                 actual_model = consumer_model
-            prepared_stages = None
-            if isinstance(actual_model, CompressorTrainSimplifiedKnownStages):
-                prepared_stages = SimplifiedTrainBuilder.prepare_stages_for_known_stages_model(
-                    original_dto_stages=actual_model._original_dto_stages,
-                    fluid_factory=actual_model.fluid_factory,
-                    rate=np.asarray(consumer_rates[i]),
-                    suction_pressure=np.asarray(operational_setting.suction_pressures[i]),
-                    discharge_pressure=np.asarray(operational_setting.discharge_pressures[i]),
-                )
-            elif isinstance(actual_model, CompressorTrainSimplifiedUnknownStages):
-                prepared_stages = SimplifiedTrainBuilder.prepare_stages_for_unknown_stages_model(
-                    stage_template=actual_model.stage,
-                    maximum_pressure_ratio_per_stage=actual_model.maximum_pressure_ratio_per_stage,
-                    fluid_factory=actual_model.fluid_factory,
-                    rate=np.asarray(consumer_rates[i]),
-                    suction_pressure=np.asarray(operational_setting.suction_pressures[i]),
-                    discharge_pressure=np.asarray(operational_setting.discharge_pressures[i]),
-                )
-
-            if prepared_stages is not None:
-                actual_model.stages = prepared_stages
+            # Prepare stages for simplified compressor train using configuration approach
+            if isinstance(actual_model, CompressorTrainSimplified):
+                if actual_model._stage_config is not None:
+                    builder = SimplifiedTrainBuilder(fluid_factory=actual_model.fluid_factory)
+                    prepared_stages, _ = builder.prepare_stages_from_config(
+                        config=actual_model._stage_config,
+                        rate=np.asarray(consumer_rates[i]),
+                        suction_pressure=np.asarray(operational_setting.suction_pressures[i]),
+                        discharge_pressure=np.asarray(operational_setting.discharge_pressures[i]),
+                    )
+                    actual_model.set_prepared_stages(prepared_stages)
 
             consumer_model.set_evaluation_input(
                 rate=np.asarray(consumer_rates[i]),

src/libecalc/domain/process/compressor/core/factory.py

@@ -1,8 +1,10 @@
 from libecalc.domain.process.compressor.core.sampled import CompressorModelSampled
 from libecalc.domain.process.compressor.core.train.simplified_train import (
-    CompressorTrainSimplifiedKnownStages,
-    CompressorTrainSimplifiedUnknownStages,
+    CompressorTrainSimplified,
+    KnownStagesConfig,
+    UnknownStagesConfig,
 )
+from libecalc.domain.process.compressor.core.train.simplified_train_builder import SimplifiedTrainBuilder
 from libecalc.domain.process.compressor.core.train.single_speed_compressor_train_common_shaft import (
     SingleSpeedCompressorTrainCommonShaft,
 )
@@ -51,25 +53,6 @@ def _create_variable_speed_compressor_train_multiple_streams_and_pressures_strea
     )
 
 
-def _create_compressor_train_simplified_with_known_stages(
-    compressor_model_dto: CompressorTrainSimplifiedWithKnownStages,
-) -> CompressorTrainSimplifiedKnownStages:
-    # Energy usage adjustment not supported for this model (yet)
-    # Issue error if factors are not default (and not changing the energy usage result)
-
-    fluid_factory = _create_fluid_factory(compressor_model_dto.fluid_model)
-    if fluid_factory is None:
-        raise ValueError("Fluid model is required for compressor train")
-    return CompressorTrainSimplifiedKnownStages(
-        fluid_factory=fluid_factory,
-        energy_usage_adjustment_constant=compressor_model_dto.energy_usage_adjustment_constant,
-        energy_usage_adjustment_factor=compressor_model_dto.energy_usage_adjustment_factor,
-        stages=compressor_model_dto.stages,
-        calculate_max_rate=compressor_model_dto.calculate_max_rate,
-        maximum_power=compressor_model_dto.maximum_power,
-    )
-
-
 def _create_single_speed_compressor_train(
     compressor_model_dto: SingleSpeedCompressorTrain,
 ) -> SingleSpeedCompressorTrainCommonShaft:
@@ -148,31 +131,69 @@ def _create_variable_speed_compressor_train_multiple_streams_and_pressures(
     )
 
 
-def _create_compressor_train_simplified_with_unknown_stages(
-    compressor_model_dto: CompressorTrainSimplifiedWithUnknownStages,
-) -> CompressorTrainSimplifiedUnknownStages:
+def _create_compressor_sampled(compressor_model_dto: CompressorSampled) -> CompressorModelSampled:
+    return CompressorModelSampled(
+        energy_usage_adjustment_constant=compressor_model_dto.energy_usage_adjustment_constant,
+        energy_usage_adjustment_factor=compressor_model_dto.energy_usage_adjustment_factor,
+        energy_usage_type=compressor_model_dto.energy_usage_type,
+        energy_usage_values=compressor_model_dto.energy_usage_values,
+        rate_values=compressor_model_dto.rate_values,
+        suction_pressure_values=compressor_model_dto.suction_pressure_values,
+        discharge_pressure_values=compressor_model_dto.discharge_pressure_values,
+        power_interpolation_values=compressor_model_dto.power_interpolation_values,
+    )
+
+
+def create_compressor_train_simplified_from_known_stages(
+    compressor_model_dto: CompressorTrainSimplifiedWithKnownStages,
+) -> CompressorTrainSimplified:
+    """Create unified simplified compressor train model from known stages DTO."""
     fluid_factory = _create_fluid_factory(compressor_model_dto.fluid_model)
     if fluid_factory is None:
         raise ValueError("Fluid model is required for compressor train")
-    return CompressorTrainSimplifiedUnknownStages(
+
+    # Create configuration for known stages
+    config = KnownStagesConfig(stages=compressor_model_dto.stages)
+
+    # Use centralized chart type detection
+    supports_max_rate = SimplifiedTrainBuilder.supports_max_rate_calculation(compressor_model_dto.stages)
+
+    return CompressorTrainSimplified(
         fluid_factory=fluid_factory,
         energy_usage_adjustment_constant=compressor_model_dto.energy_usage_adjustment_constant,
         energy_usage_adjustment_factor=compressor_model_dto.energy_usage_adjustment_factor,
-        stage=compressor_model_dto.stage,
-        maximum_pressure_ratio_per_stage=compressor_model_dto.maximum_pressure_ratio_per_stage,
+        stages=[],  # Stages will be prepared later using the configuration
         calculate_max_rate=compressor_model_dto.calculate_max_rate,
         maximum_power=compressor_model_dto.maximum_power,
+        supports_max_rate=supports_max_rate,
+        stage_config=config,
     )
 
 
-def _create_compressor_sampled(compressor_model_dto: CompressorSampled) -> CompressorModelSampled:
-    return CompressorModelSampled(
+def create_compressor_train_simplified_from_unknown_stages(
+    compressor_model_dto: CompressorTrainSimplifiedWithUnknownStages,
+) -> CompressorTrainSimplified:
+    """Create unified simplified compressor train model from unknown stages DTO."""
+    fluid_factory = _create_fluid_factory(compressor_model_dto.fluid_model)
+    if fluid_factory is None:
+        raise ValueError("Fluid model is required for compressor train")
+
+    # Create configuration for unknown stages
+    config = UnknownStagesConfig(
+        stage_template=compressor_model_dto.stage,
+        maximum_pressure_ratio_per_stage=compressor_model_dto.maximum_pressure_ratio_per_stage,
+    )
+
+    # Use centralized chart type detection
+    supports_max_rate = SimplifiedTrainBuilder.supports_max_rate_calculation([compressor_model_dto.stage])
+
+    return CompressorTrainSimplified(
+        fluid_factory=fluid_factory,
         energy_usage_adjustment_constant=compressor_model_dto.energy_usage_adjustment_constant,
         energy_usage_adjustment_factor=compressor_model_dto.energy_usage_adjustment_factor,
-        energy_usage_type=compressor_model_dto.energy_usage_type,
-        energy_usage_values=compressor_model_dto.energy_usage_values,
-        rate_values=compressor_model_dto.rate_values,
-        suction_pressure_values=compressor_model_dto.suction_pressure_values,
-        discharge_pressure_values=compressor_model_dto.discharge_pressure_values,
-        power_interpolation_values=compressor_model_dto.power_interpolation_values,
+        stages=[],  # Stages will be prepared later using the configuration
+        calculate_max_rate=compressor_model_dto.calculate_max_rate,
+        maximum_power=compressor_model_dto.maximum_power,
+        supports_max_rate=supports_max_rate,
+        stage_config=config,
     )