From 631f869219c73cbce343f39c59dad5a921fbe79f Mon Sep 17 00:00:00 2001 From: Matthew Ferrera Date: Thu, 19 Dec 2024 08:08:11 +0100 Subject: [PATCH] ENH: Enumerate inplace volumes table columns (#926) --- src/fmu/dataio/_products/inplace_volumes.py | 4 +- src/fmu/dataio/export/_enums.py | 54 ++++++++++++++++++ src/fmu/dataio/export/rms/inplace_volumes.py | 57 ++++++++----------- .../test_export_rms_volumetrics.py | 34 ++++------- 4 files changed, 93 insertions(+), 56 deletions(-) create mode 100644 src/fmu/dataio/export/_enums.py diff --git a/src/fmu/dataio/_products/inplace_volumes.py b/src/fmu/dataio/_products/inplace_volumes.py index 76a9d3aab..f3acf4350 100644 --- a/src/fmu/dataio/_products/inplace_volumes.py +++ b/src/fmu/dataio/_products/inplace_volumes.py @@ -5,6 +5,8 @@ from pydantic import BaseModel, Field, RootModel from pydantic.json_schema import GenerateJsonSchema +from fmu.dataio.export._enums import InplaceVolumes + if TYPE_CHECKING: from typing import Any, Mapping @@ -23,7 +25,7 @@ class InplaceVolumesResultRow(BaseModel): should increase the version number in a way that corresponds to the schema versioning specification (i.e. they are a patch, minor, or major change).""" - FLUID: Literal["oil", "gas", "water"] + FLUID: InplaceVolumes.Fluid ZONE: str REGION: Optional[str] = Field(default=None) FACIES: Optional[str] = Field(default=None) diff --git a/src/fmu/dataio/export/_enums.py b/src/fmu/dataio/export/_enums.py new file mode 100644 index 000000000..8fdfd4bc5 --- /dev/null +++ b/src/fmu/dataio/export/_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 used as values in the FLUID column.""" + + 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..cae505f35 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 @@ -12,6 +11,7 @@ import fmu.dataio as dio from fmu.dataio._logging import null_logger from fmu.dataio._model.enums import Classification +from fmu.dataio.export import _enums from fmu.dataio.export._decorators import experimental from fmu.dataio.export._export_result import ExportResult, ExportResultItem from fmu.dataio.export.rms._conditional_rms_imports import import_rms_package @@ -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,19 +149,29 @@ 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]: - fluid_columns = [col for col in table.columns if col.endswith(f"_{fluid}")] + for fluid in ( + _enums.InplaceVolumes.Fluid.gas.value, + _enums.InplaceVolumes.Fluid.oil.value, + ): + fluid_suffix = fluid.upper() + fluid_columns = [ + col for col in table.columns if col.endswith(f"_{fluid_suffix}") + ] if fluid_columns: fluid_table = table[table_index + fluid_columns].copy() # drop fluid suffix from columns to get standard names - fluid_table.columns = fluid_table.columns.str.replace(f"_{fluid}", "") + fluid_table.columns = fluid_table.columns.str.replace( + f"_{fluid_suffix}", "" + ) # add the fluid as column entry instead - fluid_table[_FLUID_COLUMN] = fluid.lower() + fluid_table[_enums.InplaceVolumes.FLUID_COLUMN] = fluid tables.append(fluid_table) @@ -196,7 +185,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 +203,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 +246,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..6a7cd4223 100644 --- a/tests/test_export_rms/test_export_rms_volumetrics.py +++ b/tests/test_export_rms/test_export_rms_volumetrics.py @@ -15,6 +15,7 @@ InplaceVolumesResultRow, dump, ) +from fmu.dataio.export import _enums from tests.utils import inside_rms logger = null_logger(__name__) @@ -94,13 +95,11 @@ 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 - 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 +119,7 @@ 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.inplace_volumes import _ExportVolumetricsRMS monkeypatch.chdir(rmssetup_with_fmuconfig) @@ -147,8 +143,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 @@ -299,10 +298,7 @@ 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.inplace_volumes import _ExportVolumetricsRMS monkeypatch.chdir(rmssetup_with_fmuconfig) @@ -327,7 +323,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 +356,8 @@ 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, + assert _enums.InplaceVolumes.table_columns() == list( + InplaceVolumesResultRow.model_fields.keys() ) - - export_columns = _TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS - result_columns = InplaceVolumesResultRow.model_fields.keys() - assert set(export_columns) == set(result_columns)