diff --git a/hydrolib/core/dflowfm/tim/models.py b/hydrolib/core/dflowfm/tim/models.py index fc72de66f..1cdadce62 100644 --- a/hydrolib/core/dflowfm/tim/models.py +++ b/hydrolib/core/dflowfm/tim/models.py @@ -1,6 +1,7 @@ from pathlib import Path from typing import Callable, Dict, List +from pandas import DataFrame from pydantic.v1 import Field from pydantic.v1.class_validators import validator @@ -169,3 +170,24 @@ def _raise_error_if_duplicate_time(timeseries: List[TimRecord]) -> None: f"Timeseries cannot contain duplicate times. Time: {timrecord.time} is duplicate." ) seen_times.add(timrecord.time) + + def as_dataframe(self) -> DataFrame: + """Return the timeseries as a pandas DataFrame. + + Returns: + DataFrame: The timeseries as a pandas DataFrame. + + Examples: + --------- + >>> from hydrolib.core.dflowfm.tim.models import TimModel + >>> tim_model = TimModel(filepath="tests/data/input/tim/triple_data_for_timeseries.tim") + >>> df = tim_model.as_dataframe() + >>> print(df) + 0 1 2 + 10.0 1.232 2.343 3.454 + 20.0 4.565 5.676 6.787 + 30.0 1.500 2.600 3.700 + """ + time_series = [record.data for record in self.timeseries] + index = [record.time for record in self.timeseries] + return DataFrame(time_series, index=index) diff --git a/tests/dflowfm/test_tim.py b/tests/dflowfm/test_tim.py index 9f7769be9..a70ea6d7d 100644 --- a/tests/dflowfm/test_tim.py +++ b/tests/dflowfm/test_tim.py @@ -104,6 +104,28 @@ def test_initialization(self): model = TimModel() assert len(model.comments) == 0 assert len(model.timeseries) == 0 + assert model.as_dataframe().empty + + def test_as_dataframe(self): + model = TimModel( + timeseries=self.single_data_for_timeseries_floats, + comments=["this file", "contains", "stuff"], + ) + df = model.as_dataframe() + assert all( + df.index + == [ + self.single_data_for_timeseries_floats[i].time + for i in range(len(self.single_data_for_timeseries_floats)) + ] + ) + assert df.shape == (13, 1) + assert df.columns.to_list() == [0] + vals = [ + self.single_data_for_timeseries_floats[i].data[0] + for i in range(len(self.single_data_for_timeseries_floats)) + ] + assert df.loc[:, 0].to_list() == vals @pytest.mark.parametrize( "input_data, reference_path",