diff --git a/src/fmu/dataio/export/rms/_enums.py b/src/fmu/dataio/export/rms/_enums.py new file mode 100644 index 000000000..2bbdfe78f --- /dev/null +++ b/src/fmu/dataio/export/rms/_enums.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from enum import Enum +from typing import Final + + +class InplaceVolumes: + """Enumerations relevant to inplace volumes tables.""" + + class Fluid(str, Enum): + """Fluid types""" + + OIL = "OIL" + GAS = "GAS" + WATER = "WATER" + + class TableIndexColumns(str, Enum): + """The index columns for an inplace volumes table.""" + + FLUID = "FLUID" + ZONE = "ZONE" + REGION = "REGION" + FACIES = "FACIES" + LICENSE = "LICENSE" + + FLUID_COLUMN: Final = TableIndexColumns.FLUID + """The column name and value used to indicate the index value for fluid type.""" + + class VolumetricColumns(str, Enum): + """The value columns for an inplace volumes table.""" + + BULK = "BULK" + NET = "NET" + PORV = "PORV" + HCPV = "HCPV" + STOIIP = "STOIIP" + GIIP = "GIIP" + ASSOCIATEDGAS = "ASSOCIATEDGAS" + ASSOCIATEDOIL = "ASSOCIATEDOIL" + + @staticmethod + def index_columns() -> list[str]: + """Returns a list of the index columns.""" + return [k.value for k in InplaceVolumes.TableIndexColumns] + + @staticmethod + def value_columns() -> list[str]: + """Returns a list of the value columns.""" + return [k.value for k in InplaceVolumes.VolumetricColumns] + + @staticmethod + def table_columns() -> list[str]: + """Returns a list of all table columns.""" + return InplaceVolumes.index_columns() + InplaceVolumes.value_columns() diff --git a/src/fmu/dataio/export/rms/inplace_volumes.py b/src/fmu/dataio/export/rms/inplace_volumes.py index f7b215eef..2c7b46e26 100644 --- a/src/fmu/dataio/export/rms/inplace_volumes.py +++ b/src/fmu/dataio/export/rms/inplace_volumes.py @@ -2,7 +2,6 @@ import warnings from dataclasses import dataclass -from enum import Enum from pathlib import Path from typing import Any, Final @@ -14,6 +13,7 @@ from fmu.dataio._model.enums import Classification from fmu.dataio.export._decorators import experimental from fmu.dataio.export._export_result import ExportResult, ExportResultItem +from fmu.dataio.export.rms import _enums from fmu.dataio.export.rms._conditional_rms_imports import import_rms_package from fmu.dataio.export.rms._utils import ( check_rmsapi_version, @@ -25,27 +25,6 @@ _logger: Final = null_logger(__name__) -_FLUID_COLUMN: Final = "FLUID" -_TABLE_INDEX_COLUMNS: Final = [_FLUID_COLUMN, "ZONE", "REGION", "FACIES", "LICENSE"] -_VOLUMETRIC_COLUMNS: Final = [ - "BULK", - "NET", - "PORV", - "HCPV", - "STOIIP", - "GIIP", - "ASSOCIATEDGAS", - "ASSOCIATEDOIL", -] - - -class _Fluid(str, Enum): - """Fluid types""" - - OIL = "OIL" - GAS = "GAS" - WATER = "WATER" - # rename columns to FMU standard _RENAME_COLUMNS_FROM_RMS: Final = { @@ -149,7 +128,7 @@ def _convert_table_from_rms_to_legacy_format(table: pd.DataFrame) -> pd.DataFram def _add_missing_columns_to_table(table: pd.DataFrame) -> pd.DataFrame: """Add columns with nan values if not present in table.""" _logger.debug("Add table index columns to table if missing...") - for col in _TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS: + for col in _enums.InplaceVolumes.table_columns(): if col not in table: table[col] = np.nan return table @@ -158,7 +137,7 @@ def _add_missing_columns_to_table(table: pd.DataFrame) -> pd.DataFrame: def _set_table_column_order(table: pd.DataFrame) -> pd.DataFrame: """Set the column order in the table.""" _logger.debug("Settting the table column order...") - return table[_TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS] + return table[_enums.InplaceVolumes.table_columns()] @staticmethod def _transform_and_add_fluid_column_to_table( @@ -170,10 +149,15 @@ def _transform_and_add_fluid_column_to_table( are renamed into 'BULK' and 'PORV' columns. To separate the data an additional FLUID column is added that indicates the type of fluid the row represents. """ - table_index = [col for col in _TABLE_INDEX_COLUMNS if col in table] + table_index = [ + col for col in _enums.InplaceVolumes.index_columns() if col in table + ] tables = [] - for fluid in [_Fluid.GAS.value, _Fluid.OIL.value]: + for fluid in [ + _enums.InplaceVolumes.Fluid.GAS.value, + _enums.InplaceVolumes.Fluid.OIL.value, + ]: fluid_columns = [col for col in table.columns if col.endswith(f"_{fluid}")] if fluid_columns: fluid_table = table[table_index + fluid_columns].copy() @@ -182,7 +166,7 @@ def _transform_and_add_fluid_column_to_table( fluid_table.columns = fluid_table.columns.str.replace(f"_{fluid}", "") # add the fluid as column entry instead - fluid_table[_FLUID_COLUMN] = fluid.lower() + fluid_table[_enums.InplaceVolumes.FLUID_COLUMN] = fluid.lower() tables.append(fluid_table) @@ -196,7 +180,9 @@ def _convert_table_from_legacy_to_standard_format( product. The standard format has a fluid column, and all table_index and volumetric columns are present with a standard order in the table. """ - table_index = [col for col in _TABLE_INDEX_COLUMNS if col in table] + table_index = [ + col for col in _enums.InplaceVolumes.index_columns() if col in table + ] table = self._transform_and_add_fluid_column_to_table(table, table_index) table = self._add_missing_columns_to_table(table) return self._set_table_column_order(table) @@ -212,8 +198,8 @@ def _validate_table(self) -> None: """ _logger.debug("Validating the dataframe...") - has_oil = "oil" in self._dataframe[_FLUID_COLUMN].values - has_gas = "gas" in self._dataframe[_FLUID_COLUMN].values + has_oil = "oil" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values + has_gas = "gas" 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): @@ -255,7 +241,7 @@ def _export_volume_table(self) -> ExportResult: classification=self._classification, name=self.grid_name, rep_include=False, - table_index=_TABLE_INDEX_COLUMNS, + table_index=_enums.InplaceVolumes.index_columns(), ) absolute_export_path = edata.export(self._dataframe) _logger.debug("Volume result to: %s", absolute_export_path) diff --git a/tests/test_export_rms/test_export_rms_volumetrics.py b/tests/test_export_rms/test_export_rms_volumetrics.py index 285fbd8aa..f34c9bc5e 100644 --- a/tests/test_export_rms/test_export_rms_volumetrics.py +++ b/tests/test_export_rms/test_export_rms_volumetrics.py @@ -94,13 +94,13 @@ def test_rms_volumetrics_export_class(exportvolumetrics): def test_rms_volumetrics_export_class_table_index(voltable_standard, exportvolumetrics): """See mocks in local conftest.py""" - from fmu.dataio.export.rms.inplace_volumes import _TABLE_INDEX_COLUMNS + from fmu.dataio.export.rms import _enums out = exportvolumetrics._export_volume_table() metadata = dataio.read_metadata(out.items[0].absolute_path) # check that the table index is set correctly - assert metadata["data"]["table_index"] == _TABLE_INDEX_COLUMNS + assert metadata["data"]["table_index"] == _enums.InplaceVolumes.index_columns() # should fail if missing table index exportvolumetrics._dataframe = voltable_standard.drop(columns="ZONE") @@ -120,10 +120,8 @@ def test_convert_table_from_legacy_to_standard_format( """Test that a voltable with legacy format is converted to the expected standard format""" - from fmu.dataio.export.rms.inplace_volumes import ( - _FLUID_COLUMN, - _ExportVolumetricsRMS, - ) + from fmu.dataio.export.rms import _enums + from fmu.dataio.export.rms.inplace_volumes import _ExportVolumetricsRMS monkeypatch.chdir(rmssetup_with_fmuconfig) @@ -147,8 +145,11 @@ def test_convert_table_from_legacy_to_standard_format( pd.testing.assert_frame_equal(voltable_standard, exported_table) # check that the fluid column exists and contains oil and gas - assert _FLUID_COLUMN in exported_table - assert set(exported_table[_FLUID_COLUMN].unique()) == {"oil", "gas"} + assert _enums.InplaceVolumes.FLUID_COLUMN in exported_table + assert set(exported_table[_enums.InplaceVolumes.FLUID_COLUMN].unique()) == { + "oil", + "gas", + } # check the column order assert list(exported_table.columns) == EXPECTED_COLUMN_ORDER @@ -298,11 +299,8 @@ def test_rms_volumetrics_export_function( ): """Test the public function.""" - from fmu.dataio.export.rms import export_inplace_volumes - from fmu.dataio.export.rms.inplace_volumes import ( - _TABLE_INDEX_COLUMNS, - _ExportVolumetricsRMS, - ) + from fmu.dataio.export.rms import _enums, export_inplace_volumes + from fmu.dataio.export.rms.inplace_volumes import _ExportVolumetricsRMS monkeypatch.chdir(rmssetup_with_fmuconfig) @@ -327,7 +325,7 @@ def test_rms_volumetrics_export_function( assert "volumes" in metadata["data"]["content"] assert metadata["access"]["classification"] == "restricted" - assert metadata["data"]["table_index"] == _TABLE_INDEX_COLUMNS + assert metadata["data"]["table_index"] == _enums.InplaceVolumes.index_columns() @inside_rms @@ -360,14 +358,10 @@ def test_inplace_volumes_payload_validates_against_schema( @inside_rms def test_inplace_volumes_export_and_result_columns_are_the_same( - mock_project_variable, mocked_rmsapi_modules, ) -> None: - from fmu.dataio.export.rms.inplace_volumes import ( - _TABLE_INDEX_COLUMNS, - _VOLUMETRIC_COLUMNS, - ) + from fmu.dataio.export.rms import _enums - export_columns = _TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS - result_columns = InplaceVolumesResultRow.model_fields.keys() - assert set(export_columns) == set(result_columns) + assert _enums.InplaceVolumes.table_columns() == list( + InplaceVolumesResultRow.model_fields.keys() + )