Skip to content

Commit

Permalink
Fix summary.pandas frame time index type error
Browse files Browse the repository at this point in the history

---------

Co-authored-by: Eivind Jahren <[email protected]>
  • Loading branch information
HakonSohoel and eivindjahren authored Dec 4, 2024
1 parent 4f0df4d commit c76fecd
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 30 deletions.
15 changes: 10 additions & 5 deletions python/resdata/summary/rd_sum.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import ctypes
import pandas as pd
import re
from typing import Sequence, List, Tuple, Optional
from typing import Sequence, List, Tuple, Optional, Union

# Observe that there is some convention conflict with the C code
# regarding order of arguments: The C code generally takes the time
Expand Down Expand Up @@ -477,7 +477,9 @@ def get_values(self, key, report_only=False):
else:
raise KeyError("Summary object does not have key:%s" % key)

def _make_time_vector(self, time_index):
def _make_time_vector(
self, time_index: Sequence[Union[CTime, datetime.datetime, int, datetime.date]]
) -> TimeVector:
time_points = TimeVector()
for t in time_index:
time_points.append(t)
Expand Down Expand Up @@ -558,7 +560,11 @@ def report_dates(self):
dates.append(self.get_report_time(report))
return dates

def pandas_frame(self, time_index=None, column_keys=None):
def pandas_frame(
self,
time_index: Optional[Sequence[datetime.datetime]] = None,
column_keys: Optional[Sequence[str]] = None,
) -> pd.DataFrame:
"""Will create a pandas frame with summary data.
By default you will get all time points in the summary case, but by
Expand Down Expand Up @@ -615,8 +621,7 @@ def pandas_frame(self, time_index=None, column_keys=None):
data.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
)

frame = pd.DataFrame(index=time_index, columns=list(keywords), data=data)
return frame
return pd.DataFrame(index=list(time_index), columns=list(keywords), data=data)

@staticmethod
def _compile_headers_list(
Expand Down
58 changes: 33 additions & 25 deletions python/tests/rd_tests/test_sum.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import datetime
import os
import os.path
import pytest
import shutil
import stat
import datetime
Expand All @@ -26,6 +27,7 @@ def assert_frame_equal(a, b):
from resdata.summary import Summary, SummaryKeyWordVector, SummaryVarType
from resdata.util.test import TestAreaContext
from resdata.util.test.mock import createSummary
from resdata.util.util import CTime, TimeVector
from tests import ResdataTest


Expand Down Expand Up @@ -596,30 +598,6 @@ def test_wells_and_groups(self):
self.assertEqual(case.wells(), [])
self.assertEqual(case.groups(), [])

def test_pandas(self):
case = create_case()
dates = (
[datetime.datetime(2000, 1, 1)]
+ case.dates
+ [datetime.datetime(2020, 1, 1)]
)
frame = case.pandas_frame(column_keys=["FOPT", "FOPR"], time_index=dates)

fopr = frame["FOPR"]
fopt = frame["FOPT"]

self.assertEqual(fopr[0], 0)
self.assertEqual(fopr[-1], 0)

self.assertEqual(fopt[0], 0)
self.assertEqual(fopt[0], case.first_value("FOPT"))
self.assertEqual(fopt[-1], case.last_value("FOPT"))

frame = case.pandas_frame()
rows, columns = frame.shape
self.assertEqual(len(case.keys()), columns)
self.assertEqual(len(case), rows)

def test_csv_load(self):
case = create_case2()
frame = case.pandas_frame()
Expand Down Expand Up @@ -691,7 +669,6 @@ def test_resample_extrapolate(self):
"""
Test resampling of summary with extrapolate option of lower and upper boundaries enabled
"""
from resdata.util.util import CTime, TimeVector

time_points = TimeVector()

Expand Down Expand Up @@ -757,6 +734,37 @@ def test_resample_extrapolate(self):
)


def create_time_vector(lst):
vec = TimeVector()
for l in lst:
vec.append(l)
return vec


@pytest.mark.parametrize("time_index_type", [list, create_time_vector, tuple])
def test_summary_to_pandas_frame(time_index_type):
case = create_case()
dates = time_index_type(
[datetime.datetime(2000, 1, 1)] + case.dates + [datetime.datetime(2020, 1, 1)]
)
frame = case.pandas_frame(column_keys=["FOPT", "FOPR"], time_index=dates)

fopr = frame["FOPR"]
fopt = frame["FOPT"]

assert fopr[0] == 0
assert fopr[-1] == 0

assert fopt[0] == 0
assert fopt[0] == case.first_value("FOPT")
assert fopt[-1] == case.last_value("FOPT")

frame = case.pandas_frame()
rows, columns = frame.shape
assert len(case.keys()) == columns
assert len(case) == rows


def test_t_step():
sum = createSummary(
"CASE",
Expand Down

0 comments on commit c76fecd

Please sign in to comment.