refactor: remove choke from process system

Author: jsolaas

src/libecalc/domain/process/process_solver/solvers/downstream_choke_solver.py

@@ -1,22 +1,21 @@
+from libecalc.domain.process.entities.process_units.choke import Choke
 from libecalc.domain.process.process_solver.solver import Solver
 from libecalc.domain.process.process_system.process_system import ProcessSystem
 from libecalc.domain.process.value_objects.fluid_stream import FluidStream
 
 
 class DownstreamChokeSolver(Solver):
-    def __init__(self, target_pressure: float):
+    def __init__(self, target_pressure: float, choke: Choke):
         self._target_pressure = target_pressure
+        self._choke = choke
 
     def solve(self, process_system: ProcessSystem, inlet_stream: FluidStream) -> FluidStream | None:
-        downstream_choke = process_system.get_downstream_choke()
-        assert downstream_choke is not None, "DownstreamChokeSolver needs a downstream choke"
-
         outlet_stream = process_system.propagate_stream(inlet_stream=inlet_stream)
         if outlet_stream.pressure_bara <= self._target_pressure:
             # Don't use choke if outlet pressure is below target
             return outlet_stream
         else:
             # Calculate needed pressure change in downstream choke
             pressure_change = outlet_stream.pressure_bara - self._target_pressure
-            downstream_choke.set_pressure_change(pressure_change=pressure_change)
+            self._choke.set_pressure_change(pressure_change=pressure_change)
             return process_system.propagate_stream(inlet_stream=inlet_stream)

src/libecalc/domain/process/process_solver/solvers/upstream_choke_solver.py

@@ -1,19 +1,18 @@
 from libecalc.domain.process.compressor.core.train.utils.numeric_methods import find_root
+from libecalc.domain.process.entities.process_units.choke import Choke
 from libecalc.domain.process.process_solver.solver import Solver
 from libecalc.domain.process.process_system.process_system import ProcessSystem
 from libecalc.domain.process.value_objects.fluid_stream import FluidService, FluidStream
 
 
 class UpstreamChokeSolver(Solver):
-    def __init__(self, target_pressure: float, fluid_service: FluidService, minimum_pressure: float):
+    def __init__(self, target_pressure: float, fluid_service: FluidService, minimum_pressure: float, choke: Choke):
         self._target_pressure = target_pressure
         self._fluid_service = fluid_service
         self._minimum_pressure = minimum_pressure
+        self._choke = choke
 
     def solve(self, process_system: ProcessSystem, inlet_stream: FluidStream) -> FluidStream | None:
-        upstream_choke = process_system.get_upstream_choke()
-        assert upstream_choke is not None, "UpstreamChokeSolver needs an upstream choke"
-
         outlet_stream = process_system.propagate_stream(inlet_stream)
         if outlet_stream.pressure_bara <= self._target_pressure:
             # Don't use choke if outlet pressure is below target
@@ -36,5 +35,5 @@ def get_outlet_pressure(inlet_pressure: float) -> float:
         )
 
         pressure_change = inlet_stream.pressure_bara - choked_inlet_pressure
-        upstream_choke.set_pressure_change(pressure_change=pressure_change)
+        self._choke.set_pressure_change(pressure_change=pressure_change)
         return process_system.propagate_stream(inlet_stream=inlet_stream)

src/libecalc/domain/process/process_system/process_system.py

@@ -1,4 +1,3 @@
-from libecalc.domain.process.entities.process_units.choke import Choke
 from libecalc.domain.process.entities.shaft import Shaft
 from libecalc.domain.process.process_system.process_unit import ProcessUnit
 from libecalc.domain.process.value_objects.fluid_stream import FluidStream
@@ -10,33 +9,14 @@ class ProcessSystem:
     def __init__(
         self,
         process_units: list[ProcessUnit],
-        upstream_choke: Choke | None = None,
-        downstream_choke: Choke | None = None,
     ):
         self._process_units = process_units
-        self._upstream_choke = upstream_choke
-        self._downstream_choke = downstream_choke
 
     def get_shaft(self) -> Shaft:
         return self._shaft
 
-    def get_downstream_choke(self) -> Choke | None:
-        return self._downstream_choke
-
-    def get_upstream_choke(self) -> Choke | None:
-        return self._upstream_choke
-
     def propagate_stream(self, inlet_stream: FluidStream) -> FluidStream:
-        process_units: list[ProcessUnit] = []
-        if self._upstream_choke is not None:
-            process_units.append(self._upstream_choke)
-
-        process_units.extend(self._process_units)
-
-        if self._downstream_choke is not None:
-            process_units.append(self._downstream_choke)
-
         current_inlet = inlet_stream
-        for process_unit in process_units:
+        for process_unit in self._process_units:
             current_inlet = process_unit.propagate_stream(inlet_stream=current_inlet)
         return current_inlet

tests/libecalc/domain/process/conftest.py

@@ -1,7 +1,6 @@
 import pytest
 
 from libecalc.domain.process.compressor.core.train.stage import CompressorTrainStage
-from libecalc.domain.process.entities.process_units.choke import Choke
 from libecalc.domain.process.entities.shaft import Shaft
 from libecalc.domain.process.process_solver.stream_constraint import PressureStreamConstraint
 from libecalc.domain.process.process_system.process_error import OutsideCapacityError
@@ -175,9 +174,6 @@ class StageProcessUnit(ProcessUnit):
     def __init__(self, compressor_stage: CompressorTrainStage):
         self._compressor_stage = compressor_stage
 
-    def get_shaft(self) -> Shaft:
-        return self._compressor_stage.compressor.shaft
-
     def propagate_stream(self, inlet_stream: FluidStream) -> FluidStream:
         result = self._compressor_stage.evaluate(inlet_stream_stage=inlet_stream)
         if not result.within_capacity:
@@ -189,15 +185,11 @@ def propagate_stream(self, inlet_stream: FluidStream) -> FluidStream:
 def process_system_factory(compressor_stage_factory, fluid_service):
     def create_process_system(
         process_units: list[ProcessUnit] | None = None,
-        downstream_choke: Choke | None = None,
-        upstream_choke: Choke | None = None,
     ):
         if process_units is None:
             process_units = [simple_process_unit_factory(fluid_service)]
         return ProcessSystem(
             process_units=process_units,
-            downstream_choke=downstream_choke,
-            upstream_choke=upstream_choke,
         )
 
     return create_process_system

tests/libecalc/domain/process/solvers/test_downstream_choke_solver.py

@@ -22,13 +22,12 @@ def test_downstream_choke_solver(
 ):
     downstream_choke = Choke(fluid_service=fluid_service)
     process_system = process_system_factory(
-        downstream_choke=downstream_choke,
-        process_units=[simple_process_unit_factory(pressure_multiplier=1)],
+        process_units=[simple_process_unit_factory(pressure_multiplier=1), downstream_choke],
     )
 
     inlet_stream = stream_factory(standard_rate_m3_per_day=1000, pressure_bara=inlet_pressure)
 
-    downstream_choke_solver = DownstreamChokeSolver(target_pressure=target_pressure)
+    downstream_choke_solver = DownstreamChokeSolver(target_pressure=target_pressure, choke=downstream_choke)
     outlet_stream = downstream_choke_solver.solve(process_system, inlet_stream)
 
     assert outlet_stream.pressure_bara == expected_pressure

tests/libecalc/domain/process/solvers/test_upstream_choke_solver.py

@@ -23,8 +23,7 @@ def test_upstream_choke_solver(
 ):
     choke = Choke(fluid_service=fluid_service)
     process_system = process_system_factory(
-        upstream_choke=choke,
-        process_units=[simple_process_unit_factory(pressure_multiplier=1)],
+        process_units=[choke, simple_process_unit_factory(pressure_multiplier=1)],
     )
 
     inlet_stream = stream_factory(standard_rate_m3_per_day=1000, pressure_bara=inlet_pressure)
@@ -33,6 +32,7 @@ def test_upstream_choke_solver(
         fluid_service=fluid_service,
         target_pressure=target_pressure,
         minimum_pressure=EPSILON,
+        choke=choke,
     )
     outlet_stream = upstream_choke_solver.solve(process_system, inlet_stream)
 

tests/libecalc/domain/process/test_process_solver.py

@@ -1,7 +1,6 @@
 import pytest
 from inline_snapshot import snapshot
 
-from libecalc.domain.process.entities.process_units.choke import Choke
 from libecalc.domain.process.entities.shaft import VariableSpeedShaft
 from libecalc.domain.process.process_solver.boundary import Boundary
 from libecalc.domain.process.process_solver.process_solver import (
@@ -38,7 +37,6 @@ def test_speed_solver(
     stream_constraint = stream_constraint_factory(pressure=target_pressure)
     process_system = process_system_factory(
         process_units=[compressor_train_stage_process_unit_factory(chart_data=chart_data, shaft=shaft)],
-        downstream_choke=Choke(fluid_service=fluid_service),
     )
 
     process_solver = ProcessSolver(