Skip to content

Commit e2cb846

Browse files
committed
Make _observations.py use polars
1 parent 87d7a72 commit e2cb846

2 files changed

Lines changed: 27 additions & 33 deletions

File tree

src/ert/config/_observations.py

Lines changed: 24 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
)
1919

2020
import numpy as np
21-
import pandas as pd
2221
import polars as pl
2322
import scipy as sp
2423
from pydantic import BaseModel, ConfigDict, Field, model_serializer
@@ -686,11 +685,7 @@ def from_csv(
686685
f"The CSV file ({filename}) does not exist or is not accessible.",
687686
filename,
688687
)
689-
csv_file = pd.read_csv(
690-
filename,
691-
encoding="utf-8",
692-
on_bad_lines="error",
693-
)
688+
csv_file = pl.read_csv(filename, encoding="utf-8")
694689

695690
required_columns = {
696691
"WELL_NAME",
@@ -701,7 +696,9 @@ def from_csv(
701696
"EAST",
702697
"TVD",
703698
}
704-
missing_required_columns = required_columns - set(csv_file.keys())
699+
700+
columns = set(csv_file.columns)
701+
missing_required_columns = required_columns - columns
705702
if missing_required_columns:
706703
raise ObservationConfigError.with_context(
707704
f"The rft observations file {filename} is missing required column(s) "
@@ -711,9 +708,9 @@ def from_csv(
711708

712709
rft_observations = []
713710
invalid_observations = []
714-
for row in csv_file.itertuples(index=True):
715-
east_val = validate_float(str(row.EAST), "EAST")
716-
north_val = validate_float(str(row.NORTH), "NORTH")
711+
for index, row in enumerate(csv_file.iter_rows(named=True)):
712+
east_val = validate_float(str(row["EAST"]), "EAST")
713+
north_val = validate_float(str(row["NORTH"]), "NORTH")
717714
radius = radius if radius is not None else DEFAULT_LOCALIZATION_RADIUS
718715

719716
shape_id = shape_registry.register(
@@ -725,22 +722,20 @@ def from_csv(
725722
)
726723

727724
rft_observation = cls(
728-
name=f"{observation_dict['name']}[{row.Index}]",
729-
well=str(row.WELL_NAME),
730-
date=str(row.DATE),
725+
name=f"{observation_dict['name']}[{index}]",
726+
well=str(row["WELL_NAME"]),
727+
date=str(row["DATE"]),
731728
property=observed_property,
732-
value=validate_float(
733-
str(getattr(row, observed_property)), observed_property
734-
),
735-
error=validate_float(str(row.ERROR), "ERROR"),
729+
value=validate_float(str(row[observed_property]), observed_property),
730+
error=validate_float(str(row["ERROR"]), "ERROR"),
736731
east=east_val,
737732
north=north_val,
738733
shape_id=shape_id,
739-
tvd=validate_float(str(row.TVD), "TVD"),
740-
md=validate_float(str(row.MD), "MD") if "MD" in csv_file else None,
734+
tvd=validate_float(str(row["TVD"]), "TVD"),
735+
md=validate_float(str(row["MD"]), "MD") if "MD" in columns else None,
741736
zone=(
742-
str(row.ZONE)
743-
if "ZONE" in csv_file and row.ZONE is not None
737+
str(row["ZONE"])
738+
if "ZONE" in columns and row["ZONE"] is not None
744739
else None
745740
),
746741
)
@@ -1007,11 +1002,10 @@ class SeismicObservation(BaseObservation):
10071002
TOLERANCE: ClassVar[float] = 0.1
10081003

10091004
@staticmethod
1010-
def _load_observations(filepath: Path) -> pd.DataFrame:
1011-
df = pd.read_csv(
1005+
def _load_observations(filepath: Path) -> pl.DataFrame:
1006+
df = pl.read_csv(
10121007
filepath,
10131008
encoding="utf-8",
1014-
on_bad_lines="error",
10151009
)
10161010

10171011
required_columns = {
@@ -1020,7 +1014,7 @@ def _load_observations(filepath: Path) -> pd.DataFrame:
10201014
"OBS",
10211015
"OBS_ERROR",
10221016
}
1023-
missing_required_columns = required_columns - set(df.keys())
1017+
missing_required_columns = required_columns - set(df.columns)
10241018
if missing_required_columns:
10251019
raise ObservationConfigError.with_context(
10261020
f"The seismic observations file {filepath} "
@@ -1119,11 +1113,11 @@ def from_obs_dict(
11191113
boundary_id = shape_registry.register(boundary)
11201114

11211115
seismic_observations = []
1122-
for row in df.itertuples():
1123-
east = validate_float(str(row.X_UTME), "X_UTME")
1124-
north = validate_float(str(row.Y_UTMN), "Y_UTMN")
1125-
value = validate_float(str(row.OBS), "OBS")
1126-
error = validate_float(str(row.OBS_ERROR), "OBS_ERROR")
1116+
for row in df.iter_rows(named=True):
1117+
east = validate_float(str(row["X_UTME"]), "X_UTME")
1118+
north = validate_float(str(row["Y_UTMN"]), "Y_UTMN")
1119+
value = validate_float(str(row["OBS"]), "OBS")
1120+
error = validate_float(str(row["OBS_ERROR"]), "OBS_ERROR")
11271121

11281122
# Currently supports only default localization radius as behavior of
11291123
# LOCALIZATION keyword is undefined. All shapes are being registered

tests/ert/unit_tests/config/test_observation_declaration.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -841,9 +841,9 @@ def test_that_seismic_observation_instantiates(file_context_token):
841841
create_seismic_observation(
842842
name="NAME",
843843
filepath=Path("obs.csv"),
844-
east=461231.5537527473,
844+
east=461231.55375274725,
845845
north=5933187.729869121,
846-
value=-0.0003566695393886,
846+
value=-0.00035666953938864876,
847847
error=0.005,
848848
shape_id=0,
849849
boundary_id=None,
@@ -853,7 +853,7 @@ def test_that_seismic_observation_instantiates(file_context_token):
853853
filepath=Path("obs.csv"),
854854
east=461156.9532936567,
855855
north=5933317.28138355,
856-
value=-0.0005293887515127,
856+
value=-0.0005293887515127136,
857857
error=0.005,
858858
shape_id=1,
859859
boundary_id=None,

0 commit comments

Comments
 (0)