def test_collocation_flag(test_data, key):
ref_frame, test_frame, expected_nan = setup_data(test_data, key)
flag = np.random.choice([True, False], len(ref_frame))
# with array
res = tmatching.temporal_collocation(
ref_frame,
test_frame,
pd.Timedelta(6, "H"),
flag=flag,
)
compare_with_nan(res.iloc[:, 0].values[~flag],
test_frame.iloc[:, 0].values[~flag])
assert np.all(np.isnan(res.values[:, 0][flag]))
# with array, using invalid as replacement
res = tmatching.temporal_collocation(
ref_frame,
test_frame,
pd.Timedelta(6, "H"),
flag=flag,
use_invalid=True,
)
compare_with_nan(res.iloc[:, 0].values, test_frame.iloc[:, 0].values)
# with dataframe
test_frame["flag"] = flag
res = tmatching.temporal_collocation(
ref_frame,
test_frame,
pd.Timedelta(6, "H"),
flag="flag",
)
compare_with_nan(res.iloc[:, 0].values[~flag],
test_frame.iloc[:, 0].values[~flag])
assert np.all(np.isnan(res.iloc[:, 0].values[flag]))
def test_collocation_nearest_neighbour(test_data, key):
ref_frame, test_frame, expected_nan = setup_data(test_data, key)
res = tmatching.temporal_collocation(ref_frame, test_frame,
pd.Timedelta(6, "H"))
assert_equal_except_nan(res, test_frame, expected_nan)