MAINT: Replace column strings with enums in inplace_volumes

equinor · Dec 19, 2024 · 2950512 · 2950512
1 parent 704f5cb
commit 2950512
Showing 1 changed file with 39 additions and 26 deletions.
diff --git a/src/fmu/dataio/export/rms/inplace_volumes.py b/src/fmu/dataio/export/rms/inplace_volumes.py
@@ -26,28 +26,31 @@
 _logger: Final = null_logger(__name__)
 
 
+VolumetricColumns = _enums.InplaceVolumes.VolumetricColumns
+TableIndexColumns = _enums.InplaceVolumes.TableIndexColumns
+
 # rename columns to FMU standard
 _RENAME_COLUMNS_FROM_RMS: Final = {
     "Proj. real.": "REAL",
-    "Zone": "ZONE",
-    "Segment": "REGION",
-    "Boundary": "LICENSE",
-    "Facies": "FACIES",
-    "BulkOil": "BULK_OIL",
-    "NetOil": "NET_OIL",
-    "PoreOil": "PORV_OIL",
-    "HCPVOil": "HCPV_OIL",
-    "STOIIP": "STOIIP_OIL",
-    "AssociatedGas": "ASSOCIATEDGAS_OIL",
-    "BulkGas": "BULK_GAS",
-    "NetGas": "NET_GAS",
-    "PoreGas": "PORV_GAS",
-    "HCPVGas": "HCPV_GAS",
-    "GIIP": "GIIP_GAS",
-    "AssociatedLiquid": "ASSOCIATEDOIL_GAS",
-    "Bulk": "BULK_TOTAL",
-    "Net": "NET_TOTAL",
-    "Pore": "PORV_TOTAL",
+    "Zone": TableIndexColumns.ZONE.value,
+    "Segment": TableIndexColumns.REGION.value,
+    "Boundary": TableIndexColumns.LICENSE.value,
+    "Facies": TableIndexColumns.FACIES.value,
+    "BulkOil": VolumetricColumns.BULK.value + "_OIL",
+    "NetOil": VolumetricColumns.NET.value + "_OIL",
+    "PoreOil": VolumetricColumns.PORV.value + "_OIL",
+    "HCPVOil": VolumetricColumns.HCPV.value + "_OIL",
+    "STOIIP": VolumetricColumns.STOIIP.value + "_OIL",
+    "AssociatedGas": VolumetricColumns.ASSOCIATEDGAS.value + "_OIL",
+    "BulkGas": VolumetricColumns.BULK.value + "_GAS",
+    "NetGas": VolumetricColumns.NET.value + "_GAS",
+    "PoreGas": VolumetricColumns.PORV.value + "_GAS",
+    "HCPVGas": VolumetricColumns.HCPV.value + "_GAS",
+    "GIIP": VolumetricColumns.GIIP.value + "_GAS",
+    "AssociatedLiquid": VolumetricColumns.ASSOCIATEDOIL.value + "_GAS",
+    "Bulk": VolumetricColumns.BULK.value + "_TOTAL",
+    "Net": VolumetricColumns.NET.value + "_TOTAL",
+    "Pore": VolumetricColumns.PORV.value + "_TOTAL",
 }
 
 
@@ -203,8 +206,14 @@ def _validate_table(self) -> None:
         """
         _logger.debug("Validating the dataframe...")
 
-        has_oil = "oil" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
-        has_gas = "gas" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        has_oil = (
+            _enums.InplaceVolumes.Fluid.oil.value
+            in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        )
+        has_gas = (
+            _enums.InplaceVolumes.Fluid.gas.value
+            in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        )
 
         # check that one of oil and gas fluids are present
         if not (has_oil or has_gas):
@@ -216,15 +225,19 @@ def _validate_table(self) -> None:
 
         # create list of missing or non-defined required columns
         missing_calculations = []
-        for col in ["BULK", "PORV", "HCPV"]:
+        for col in [
+            VolumetricColumns.BULK.value,
+            VolumetricColumns.PORV.value,
+            VolumetricColumns.HCPV.value,
+        ]:
             if self._is_column_missing_in_table(col):
                 missing_calculations.append(col)
 
-        if has_oil and self._is_column_missing_in_table("STOIIP"):
-            missing_calculations.append("STOIIP")
+        if has_oil and self._is_column_missing_in_table(VolumetricColumns.STOIIP.value):
+            missing_calculations.append(VolumetricColumns.STOIIP.value)
 
-        if has_gas and self._is_column_missing_in_table("GIIP"):
-            missing_calculations.append("GIIP")
+        if has_gas and self._is_column_missing_in_table(VolumetricColumns.GIIP.value):
+            missing_calculations.append(VolumetricColumns.GIIP.value)
 
         if missing_calculations:
             raise RuntimeError(