diff --git a/timely_beliefs/beliefs/utils.py b/timely_beliefs/beliefs/utils.py index 1ca3e049..5470d71d 100644 --- a/timely_beliefs/beliefs/utils.py +++ b/timely_beliefs/beliefs/utils.py @@ -1067,11 +1067,15 @@ def convert_to_instantaneous( """ df2 = df.copy() df2.index = df2.index + df.event_resolution - df = df.reset_index().set_index(["event_start", "belief_time", "source", "cumulative_probability"]) - df2 = df2.reset_index().set_index(["event_start", "belief_time", "source", "cumulative_probability"]) + df = df.reset_index().set_index( + ["event_start", "belief_time", "source", "cumulative_probability"] + ) + df2 = df2.reset_index().set_index( + ["event_start", "belief_time", "source", "cumulative_probability"] + ) df = pd.concat([df, df2], axis=1) if boundary_policy == "first": - s = df.fillna(method='bfill', axis=1).iloc[:, 0] + s = df.fillna(method="bfill", axis=1).iloc[:, 0] else: s = getattr(df, boundary_policy)(axis=1).rename("event_value") df = s.sort_index().reset_index().set_index("event_start") diff --git a/timely_beliefs/tests/test_df_resampling.py b/timely_beliefs/tests/test_df_resampling.py index 35b24370..8bee0b86 100644 --- a/timely_beliefs/tests/test_df_resampling.py +++ b/timely_beliefs/tests/test_df_resampling.py @@ -449,6 +449,6 @@ def test_upsample_to_instantaneous(df_4111, test_source_a: BeliefSource): ] pd.testing.assert_index_equal( df.index.get_level_values(level="event_start"), - pd.DatetimeIndex(expected_event_starts, name="event_start") + pd.DatetimeIndex(expected_event_starts, name="event_start"), ) assert df["event_value"].values.tolist() == expected_values