refactor: remove double mapping pump model

frodehk · frodehk · commit 2fd78d65b62f · 2025-08-26T12:40:18.000+02:00
diff --git a/src/libecalc/domain/process/pump/factory.py b/src/libecalc/domain/process/pump/factory.py
diff --git a/src/libecalc/domain/process/pump/pump.py b/src/libecalc/domain/process/pump/pump.py
@@ -1,18 +1,13 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Literal
 
 import numpy as np
 from numpy.typing import NDArray
-from pydantic import ConfigDict
-from pydantic.dataclasses import dataclass
 from scipy.interpolate import interp1d
 
-from libecalc.common.energy_model_type import EnergyModelType
 from libecalc.common.list.adjustment import transform_linear
 from libecalc.common.logger import logger
-from libecalc.common.serializable_chart import SingleSpeedChartDTO, VariableSpeedChartDTO
 from libecalc.common.units import Unit, UnitConstants
 from libecalc.domain.process.core.results import PumpModelResult
 from libecalc.domain.process.value_objects.chart import SingleSpeedChart, VariableSpeedChart
@@ -90,15 +85,6 @@ def _calculate_power(
         )
 
 
-@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
-class PumpModelDTO:
-    energy_usage_adjustment_constant: float
-    energy_usage_adjustment_factor: float
-    chart: SingleSpeedChartDTO | VariableSpeedChartDTO
-    head_margin: float
-    typ: Literal[EnergyModelType.PUMP_MODEL] = EnergyModelType.PUMP_MODEL
-
-
 class PumpSingleSpeed(PumpModel):
     def __init__(
         self,
diff --git a/src/libecalc/presentation/yaml/domain/reference_service.py b/src/libecalc/presentation/yaml/domain/reference_service.py
@@ -4,7 +4,7 @@
 from libecalc.domain.infrastructure.energy_components.generator_set import GeneratorSetModel
 from libecalc.domain.infrastructure.energy_components.legacy_consumer.tabulated import TabularEnergyFunction
 from libecalc.domain.process.compressor.dto.model_types import CompressorModelTypes
-from libecalc.domain.process.pump.pump import PumpModelDTO
+from libecalc.domain.process.pump.pump import PumpModel
 from libecalc.dto import FuelType
 
 
@@ -24,6 +24,6 @@ def get_generator_set_model(self, reference: str) -> GeneratorSetModel: ...
 
     def get_compressor_model(self, reference: str) -> CompressorModelTypes: ...
 
-    def get_pump_model(self, reference: str) -> PumpModelDTO: ...
+    def get_pump_model(self, reference: str) -> PumpModel: ...
 
     def get_tabulated_model(self, reference: str) -> TabularEnergyFunction: ...
diff --git a/src/libecalc/presentation/yaml/mappers/consumer_function_mapper.py b/src/libecalc/presentation/yaml/mappers/consumer_function_mapper.py
@@ -40,7 +40,6 @@
     VariableSpeedCompressorTrainMultipleStreamsAndPressures,
 )
 from libecalc.domain.process.compressor.dto.model_types import CompressorModelTypes
-from libecalc.domain.process.pump.factory import create_pump_model
 from libecalc.domain.regularity import Regularity
 from libecalc.domain.time_series_flow_rate import TimeSeriesFlowRate
 from libecalc.domain.time_series_variable import TimeSeriesVariable
@@ -300,7 +299,7 @@ def _map_tabular(
     def _map_pump(
         self, model: YamlEnergyUsageModelPump, consumes: ConsumptionType, period: Period
     ) -> PumpConsumerFunction:
-        energy_model = self.__references.get_pump_model(model.energy_function)
+        pump_model = self.__references.get_pump_model(model.energy_function)
         period_regularity, period_evaluator = self._period_subsets[period]
         if consumes != ConsumptionType.ELECTRICITY:
             raise InvalidConsumptionType(actual=ConsumptionType.ELECTRICITY, expected=consumes)
@@ -333,7 +332,6 @@ def _map_pump(
         )
         discharge_pressure = ExpressionTimeSeriesPressure(time_series_expression=discharge_pressure_expression)
 
-        pump_model = create_pump_model(pump_model_dto=energy_model)
         return PumpConsumerFunction(
             power_loss_factor=power_loss_factor,
             pump_function=pump_model,
@@ -631,7 +629,7 @@ def _map_pump_system(
         pumps = []
         for pump in model.pumps:
             pump_model = self.__references.get_pump_model(pump.chart)
-            pumps.append(ConsumerSystemComponent(name=pump.name, facility_model=create_pump_model(pump_model)))
+            pumps.append(ConsumerSystemComponent(name=pump.name, facility_model=pump_model))
 
         operational_settings: list[ConsumerSystemOperationalSettingExpressions] = []
         for operational_setting in model.operational_settings:
diff --git a/src/libecalc/presentation/yaml/mappers/facility_input.py b/src/libecalc/presentation/yaml/mappers/facility_input.py
@@ -12,7 +12,8 @@
 from libecalc.domain.infrastructure.energy_components.legacy_consumer.tabulated import TabularEnergyFunction
 from libecalc.domain.process.compressor.dto import CompressorSampled as CompressorTrainSampledDTO
 from libecalc.domain.process.dto import EnergyModel
-from libecalc.domain.process.pump.pump import PumpModelDTO
+from libecalc.domain.process.pump.pump import PumpSingleSpeed, PumpVariableSpeed
+from libecalc.domain.process.value_objects.chart import SingleSpeedChart, VariableSpeedChart
 from libecalc.domain.resource import Resource, Resources
 from libecalc.presentation.yaml.file_context import FileContext, FileMark
 from libecalc.presentation.yaml.mappers.energy_model_factory import EnergyModelFactory
@@ -116,64 +117,62 @@ def _create_compressor_train_sampled_dto_model_data(
 
 def _create_pump_model_single_speed_dto_model_data(
     resource: Resource, facility_data: YamlPumpChartSingleSpeed, **kwargs
-) -> PumpModelDTO:
+) -> PumpSingleSpeed:
     chart_data = get_single_speed_chart_data(resource=resource)
-
-    chart = SingleSpeedChartDTO(
-        speed_rpm=chart_data.speed,
-        efficiency_fraction=convert_efficiency_to_fraction(
-            efficiency_values=chart_data.efficiency,
-            input_unit=YAML_UNIT_MAPPING[facility_data.units.efficiency],
-        ),
-        rate_actual_m3_hour=convert_rate_to_am3_per_hour(
-            rate_values=chart_data.rate, input_unit=YAML_UNIT_MAPPING[facility_data.units.rate]
-        ),
-        polytropic_head_joule_per_kg=convert_head_to_joule_per_kg(
-            head_values=chart_data.head, input_unit=YAML_UNIT_MAPPING[facility_data.units.head]
-        ),
-    )
-
-    head_margin = facility_data.head_margin
-
-    return PumpModelDTO(
-        chart=chart,
-        energy_usage_adjustment_constant=_get_adjustment_constant(facility_data),
-        energy_usage_adjustment_factor=_get_adjustment_factor(facility_data),
-        head_margin=head_margin,
-    )
-
-
-def _create_pump_chart_variable_speed_dto_model_data(
-    resource: Resource, facility_data: YamlPumpChartVariableSpeed, **kwargs
-) -> PumpModelDTO:
-    curves_data = chart_curves_as_resource_to_dto_format(resource=resource)
-
-    curves: list[ChartCurveDTO] = [
-        ChartCurveDTO(
-            speed_rpm=curve.speed,
+    pump_chart = SingleSpeedChart(
+        SingleSpeedChartDTO(
+            speed_rpm=chart_data.speed,
             rate_actual_m3_hour=convert_rate_to_am3_per_hour(
-                rate_values=curve.rate,
-                input_unit=YAML_UNIT_MAPPING[facility_data.units.rate],
+                rate_values=chart_data.rate, input_unit=YAML_UNIT_MAPPING[facility_data.units.rate]
             ),
             polytropic_head_joule_per_kg=convert_head_to_joule_per_kg(
-                head_values=curve.head,
-                input_unit=YAML_UNIT_MAPPING[facility_data.units.head],
+                head_values=chart_data.head, input_unit=YAML_UNIT_MAPPING[facility_data.units.head]
             ),
             efficiency_fraction=convert_efficiency_to_fraction(
-                efficiency_values=curve.efficiency,
+                efficiency_values=chart_data.efficiency,
                 input_unit=YAML_UNIT_MAPPING[facility_data.units.efficiency],
             ),
         )
-        for curve in curves_data
-    ]
+    )
+    return PumpSingleSpeed(
+        pump_chart=pump_chart,
+        energy_usage_adjustment_constant=_get_adjustment_constant(facility_data),
+        energy_usage_adjustment_factor=_get_adjustment_factor(facility_data),
+        head_margin=facility_data.head_margin,
+    )
 
-    head_margin = facility_data.head_margin
 
-    return PumpModelDTO(
-        chart=VariableSpeedChartDTO(curves=curves),
+def _create_pump_chart_variable_speed_dto_model_data(
+    resource: Resource, facility_data: YamlPumpChartVariableSpeed, **kwargs
+) -> PumpVariableSpeed:
+    curves_data = chart_curves_as_resource_to_dto_format(resource=resource)
+    pump_chart = VariableSpeedChart(
+        VariableSpeedChartDTO(
+            curves=[
+                ChartCurveDTO(
+                    speed_rpm=curve.speed,
+                    rate_actual_m3_hour=convert_rate_to_am3_per_hour(
+                        rate_values=curve.rate,
+                        input_unit=YAML_UNIT_MAPPING[facility_data.units.rate],
+                    ),
+                    polytropic_head_joule_per_kg=convert_head_to_joule_per_kg(
+                        head_values=curve.head,
+                        input_unit=YAML_UNIT_MAPPING[facility_data.units.head],
+                    ),
+                    efficiency_fraction=convert_efficiency_to_fraction(
+                        efficiency_values=curve.efficiency,
+                        input_unit=YAML_UNIT_MAPPING[facility_data.units.efficiency],
+                    ),
+                )
+                for curve in curves_data
+            ]
+        )
+    )
+    return PumpVariableSpeed(
+        pump_chart=pump_chart,
         energy_usage_adjustment_constant=_get_adjustment_constant(facility_data),
         energy_usage_adjustment_factor=_get_adjustment_factor(facility_data),
-        head_margin=head_margin,
+        head_margin=facility_data.head_margin,
     )
 
 
diff --git a/src/libecalc/presentation/yaml/yaml_entities.py b/src/libecalc/presentation/yaml/yaml_entities.py
@@ -19,7 +19,7 @@
 from libecalc.domain.infrastructure.energy_components.legacy_consumer.tabulated import TabularEnergyFunction
 from libecalc.domain.process.compressor.dto.model_types import CompressorModelTypes
 from libecalc.domain.process.dto.base import EnergyModel
-from libecalc.domain.process.pump.pump import PumpModelDTO
+from libecalc.domain.process.pump.pump import PumpModel
 from libecalc.domain.resource import Resource
 from libecalc.dto import FuelType
 from libecalc.presentation.yaml.domain.reference_service import InvalidReferenceException, ReferenceService
@@ -197,9 +197,9 @@ def get_compressor_model(self, reference: str) -> CompressorModelTypes:
             raise InvalidReferenceException("compressor model", reference)
         return model  # noqa
 
-    def get_pump_model(self, reference: str) -> PumpModelDTO:
+    def get_pump_model(self, reference: str) -> PumpModel:
         model = self._get_model_reference(reference, "compressor model")
-        if not isinstance(model, PumpModelDTO):
+        if not isinstance(model, PumpModel):
             raise InvalidReferenceException("pump model", reference)
         return model
 
diff --git a/tests/libecalc/input/mappers/test_consumer_chart.py b/tests/libecalc/input/mappers/test_consumer_chart.py
@@ -3,7 +3,8 @@
 from libecalc.common.errors.exceptions import InvalidResourceException
 from libecalc.common.serializable_chart import SingleSpeedChartDTO
 from libecalc.common.units import Unit
-from libecalc.domain.process.pump.pump import PumpModelDTO
+from libecalc.domain.process.pump.pump import PumpModel, PumpSingleSpeed, PumpVariableSpeed
+from libecalc.domain.process.value_objects.chart import SingleSpeedChart
 from libecalc.presentation.yaml.mappers.facility_input import (
     _create_pump_model_single_speed_dto_model_data,
 )
@@ -86,44 +87,21 @@ def pump_chart():
 class TestSingleSpeedChart:
     def test_valid_with_speed(self, pump_chart, chart_resource_with_speed):
         """Test that speed can be specified. Note: 1.0 and 1 is considered equal."""
-        pump_model_dto = _create_pump_model_single_speed_dto_model_data(
+        pump_model = _create_pump_model_single_speed_dto_model_data(
             resource=chart_resource_with_speed,
             facility_data=pump_chart,
         )
-        assert pump_model_dto == PumpModelDTO(
-            chart=SingleSpeedChartDTO(
-                rate_actual_m3_hour=[6.0, 6.0],
-                polytropic_head_joule_per_kg=[
-                    Unit.POLYTROPIC_HEAD_METER_LIQUID_COLUMN.to(Unit.POLYTROPIC_HEAD_JOULE_PER_KG)(x)
-                    for x in [7.0, 7.0]
-                ],
-                efficiency_fraction=[0.08, 0.08],
-                speed_rpm=5.0,
-            ),
-            energy_usage_adjustment_factor=1,
-            energy_usage_adjustment_constant=0,
-            head_margin=0,
-        )
+
+        assert pump_model.pump_chart.speed == 5.0
 
     def test_valid_without_speed(self, pump_chart, chart_resource_without_speed):
-        pump_model_dto = _create_pump_model_single_speed_dto_model_data(
+        pump_model = _create_pump_model_single_speed_dto_model_data(
             resource=chart_resource_without_speed,
             facility_data=pump_chart,
         )
-        assert pump_model_dto == PumpModelDTO(
-            chart=SingleSpeedChartDTO(
-                rate_actual_m3_hour=[6.0, 6.0],
-                polytropic_head_joule_per_kg=[
-                    Unit.POLYTROPIC_HEAD_METER_LIQUID_COLUMN.to(Unit.POLYTROPIC_HEAD_JOULE_PER_KG)(x)
-                    for x in [7.0, 7.0]
-                ],
-                efficiency_fraction=[0.08, 0.08],
-                speed_rpm=1,
-            ),
-            energy_usage_adjustment_constant=0,
-            energy_usage_adjustment_factor=1,
-            head_margin=0,
-        )
+
+        # Speed set to 1.0 if header not found
+        assert pump_model.pump_chart.speed == 1.0
 
     def test_invalid_unequal_speed(self, pump_chart, chart_resource_unequal_speed):
         with pytest.raises(InvalidResourceException) as exception_info:
