diff --git a/src/xtgeo/well/_welldata.py b/src/xtgeo/well/_welldata.py index a6b1ea31f..2aa41aaa2 100644 --- a/src/xtgeo/well/_welldata.py +++ b/src/xtgeo/well/_welldata.py @@ -8,15 +8,17 @@ 4 463256.916 5930542.302 -46.1533 2.0000 NaN NaN NaN ... ... ... ... ... ... ... ... -Where each log beside the 3 first has a wlogtypes dictoniary, telling if the logs are +Where each log has a wlogtypes dictionary, telling if the logs are treated as discrete (DISC) or continuous (CONT). In addition there is a wlogrecords dict, storing the unit for continuous logs (defaulted to None) or a dictionary of codes if the log in DISC type. -The purpose here is to automate; if a column is added to the dataframe, then the -class methods here will try to guess the wlogtype and wlogtype, and add those; similarly -of a column is removed, the corresponding entries in wlogtypes and wlogrecords will be -deleted. +The 3 first columsn are the coordinates. + +The purpose here is to automate and improve; if a column is added to the dataframe, then +the class methods here will try to guess the wlogtype and wlogtype, and add those; +similarly of a column is removed, the corresponding entries in wlogtypes and wlogrecords +will be deleted. """ from dataclasses import dataclass, field from enum import Enum, EnumMeta, unique @@ -25,7 +27,7 @@ class methods here will try to guess the wlogtype and wlogtype, and add those; s class _LogTypeMeta(EnumMeta): - """For 'in' method, cf https://stackoverflow.com/questions/43634618""" + """For enabling 'in' method, cf https://stackoverflow.com/questions/43634618""" def __contains__(cls, item): try: @@ -252,3 +254,21 @@ def set_wlogrecord(self, name: str, record: dict) -> None: raise ValueError("Something went wrong when setting logrecord.") self.ensure_consistency() + + def get_dataframe_copy(self, infer_dtype: bool = False): + """Get a deep copy of the dataframe, with options. + + If infer_dtype is True, then DISC columns will be of "int32" type + """ + + if infer_dtype: + dfr = self.data.copy() + + for name, wtype in self.wlogtypes.items(): + if "DISC" in wtype: + dfr[name] = dfr[name].astype("int32") + + return dfr + + else: + return self.data.copy() diff --git a/tests/test_well/test_welldata_class.py b/tests/test_well/test_welldata_class.py index 091d51b78..0e330e9bc 100644 --- a/tests/test_well/test_welldata_class.py +++ b/tests/test_well/test_welldata_class.py @@ -153,3 +153,38 @@ def test_logtype_class(): with pytest.raises(ValueError, match="is not a valid"): _LogType("FOO") + + +def test_welldata_dataframe_copy(generate_data: pd.DataFrame): + """Test get dataframe method, with option""" + + instance = _WellData(generate_data) + + copy = instance.get_dataframe_copy() + col = list(copy) + + dtypes = [str(entry) for entry in copy[col].dtypes] + assert dtypes == [ + "float64", + "float64", + "float64", + "float32", + "float32", + "float32", + "float32", + ] + + copy = instance.get_dataframe_copy(infer_dtype=True) + + dtypes = [str(entry) for entry in copy[col].dtypes] + assert dtypes == [ + "float64", + "float64", + "float64", + "float32", + "float32", + "int32", + "int32", + ] + + instance.data =