Modification to enable ESP pump viscosity corrections

docs/source/alfacase_definitions/PumpEquipmentDescription.txt

@@ -19,6 +19,7 @@
             esp_speed_curve: \ :class:`Curve <barril.curve.curve.Curve>`\  = Curve(Hz, s)[]
             esp_number_of_stages: int = 1
             esp_reference_density: \ :class:`Scalar <barril.units.Scalar>`\  = Scalar(0.0, 'kg/m3', 'density')
+            esp_viscosity_model: \ :class:`PumpViscosityModel <alfasim_sdk._internal.constants.PumpViscosityModel>`\  = PumpViscosityModel.NoModel
 
 .. tab:: Schema
 
@@ -54,3 +55,4 @@
             esp_reference_density:  # optional
                 value: number
                 unit: string
+            esp_viscosity_model: \ :class:`PumpViscosityModel <alfasim_sdk._internal.constants.PumpViscosityModel>`\  # optional

src/alfasim_sdk/__init__.py

@@ -290,6 +290,8 @@ def get_alfasim_sdk_api_path():
 from alfasim_sdk._internal.constants import CorrelationsOilViscosity
 from alfasim_sdk._internal.constants import CorrelationsGasViscosity
 from alfasim_sdk._internal.constants import CorrelationsSurfaceTension
+from alfasim_sdk._internal.constants import PumpViscosityModel
+
 
 # Plugins: Layouts imports
 from alfasim_sdk._internal.layout import tab
@@ -462,6 +464,7 @@ def get_alfasim_sdk_api_path():
     "ProfileDescription",
     "ProfileOutputDescription",
     "PumpEquipmentDescription",
+    "PumpViscosityModel",
     "PumpType",
     "PvtModelCombinedDescription",
     "PvtModelCompositionalDescription",

src/alfasim_sdk/_internal/alfacase/case_description.py

@@ -692,6 +692,7 @@ class TablePumpDescription:
                              )
 
 
+
     def __attrs_post_init__(self):
         expected_length = len(self.speeds)
         all_fields = list(attr.fields_dict(self.__class__).keys())
@@ -758,6 +759,8 @@ class PumpEquipmentDescription:
         category="density", default=Scalar(0.0, "kg/m3")
     )
 
+    esp_viscosity_model = attrib_enum(default=constants.PumpViscosityModel.NoModel)
+
 
 @attr.s(frozen=True, slots=True)
 class CompressorPressureTableDescription:

src/alfasim_sdk/_internal/alfacase/schema.py

@@ -945,6 +945,7 @@
         ),
         Optional("esp_number_of_stages"): Int(),
         Optional("esp_reference_density"): Map({"value": Float(), "unit": Str()}),
+        Optional("esp_viscosity_model"): Enum(['no_model', 'ansihi_2010']),
     }
 )
 pvt_model_combined_description_schema = Map(
@@ -1134,5 +1135,5 @@
         Optional("walls"): Seq(wall_description_schema),
     }
 )
-# [[[end]]] (checksum: 0d1c4f2308285b6cae0f4472304733f4)
+# [[[end]]] (checksum: 52573dd5a677a2b894836415f5330fa6)
 # fmt: on

src/alfasim_sdk/_internal/constants.py

@@ -407,6 +407,17 @@ class CompressorSpeedType(Enum):
     ConstantSpeed = "constant_speed"
 
 
+class PumpViscosityModel(Enum):
+    """
+    Defines a viscosity model correction to ESP
+        - ``NoModel``: Viscosity model correction is disabled.
+        - ``AnsiHi2010``: Applies the ANSI-HI(2010) viscosity correction.
+    """
+
+    NoModel = "no_model"
+    AnsiHi2010 = "ansihi_2010"
+
+
 class OutputAttachmentLocation(Enum):
     """
     Output Attachment Location will tell the location in which this attachment's data should be retrieved from.

tests/alfacase/test_generate_case_description_docstring/test_generate_case_description_docstring_PumpEquipmentDescription_.txt

@@ -19,6 +19,7 @@
             esp_speed_curve: \ :class:`Curve <barril.curve.curve.Curve>`\  = Curve(Hz, s)[]
             esp_number_of_stages: int = 1
             esp_reference_density: \ :class:`Scalar <barril.units.Scalar>`\  = Scalar(0.0, 'kg/m3', 'density')
+            esp_viscosity_model: \ :class:`PumpViscosityModel <alfasim_sdk._internal.constants.PumpViscosityModel>`\  = PumpViscosityModel.NoModel
 
 .. tab:: Schema
 
@@ -54,3 +55,4 @@
             esp_reference_density:  # optional
                 value: number
                 unit: string
+            esp_viscosity_model: \ :class:`PumpViscosityModel <alfasim_sdk._internal.constants.PumpViscosityModel>`\  # optional

tests/alfacase/test_generate_case_schema/test_generate_schema_for_all_cases.txt

@@ -999,6 +999,7 @@ pump_equipment_description_schema = Map(
         ),
         Optional("esp_number_of_stages"): Int(),
         Optional("esp_reference_density"): Map({"value": Float(), "unit": Str()}),
+        Optional("esp_viscosity_model"): Enum(['no_model', 'ansihi_2010']),
     }
 )
 

tests/common_testing/alfasim_sdk_common_testing/filled_case_descriptions.py

@@ -505,6 +505,7 @@
     ),
     esp_number_of_stages=2,
     esp_reference_density=Scalar(1000.0, "kg/m3"),
+    esp_viscosity_model=constants.PumpViscosityModel.AnsiHi2010,
 )
 VALVE_DESCRIPTION = case_description.ValveEquipmentDescription(
     position=Scalar(100.0, "m"),