From ba8d9b71d17ec8a88336ef6461cead3fb060642c Mon Sep 17 00:00:00 2001 From: Victor Garcia Reolid Date: Fri, 22 Mar 2024 11:42:12 +0100 Subject: [PATCH] Add options floor_event_start, ceil_event_start and round_event_start to read_csv to floor, ceil round the event_start to the sensor resolution, respectively. Signed-off-by: Victor Garcia Reolid --- timely_beliefs/beliefs/utils.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/timely_beliefs/beliefs/utils.py b/timely_beliefs/beliefs/utils.py index e7d4dc8d..7967d296 100644 --- a/timely_beliefs/beliefs/utils.py +++ b/timely_beliefs/beliefs/utils.py @@ -574,6 +574,9 @@ def read_csv( # noqa C901 :param event_starts_before: Optionally, keep only events that start before this datetime. Exclusive for non-instantaneous events, inclusive for instantaneous events. Note that the last event may transpire partially after this datetime. + :param floor_event_start: Whether to floor the event_start datetime to the sensor event_resolution. + :param ceil_event_start: Whether to ceil the event_start datetime to the sensor event_resolution. + :param round_event_start: Whether to round the event_start datetime to the sensor event_resolution. :param datetime_column_split: Optionally, help parse the datetime column by splitting according to some string. For example: "1 jan 2022 00:00 - 1 jan 2022 01:00" @@ -633,6 +636,9 @@ def read_csv( # noqa C901 ext = find_out_extension(path) dayfirst = kwargs.pop("dayfirst", None) + floor_event_start = kwargs.pop("floor_event_start", False) + ceil_event_start = kwargs.pop("ceil_event_start", False) + round_event_start = kwargs.pop("round_event_start", False) if ext.lower() == "csv": df = pd.read_csv(path, **kwargs) @@ -700,6 +706,13 @@ def read_csv( # noqa C901 if cumulative_probability is not None: df["cumulative_probability"] = cumulative_probability + if ceil_event_start: + df["event_start"] = df["event_start"].dt.ceil(sensor.event_resolution) + elif floor_event_start: + df["event_start"] = df["event_start"].dt.floor(sensor.event_resolution) + elif round_event_start: + df["event_start"] = df["event_start"].dt.round(sensor.event_resolution) + # Construct BeliefsDataFrame bdf = classes.BeliefsDataFrame(df, sensor=sensor)