def test_values(
distr: PiecewiseLinear,
target: List[float],
expected_target_cdf: List[float],
expected_target_crps: List[float],
serialize_fn,
):
distr = serialize_fn(distr)
target = mx.nd.array(target).reshape(shape=(len(target), ))
expected_target_cdf = np.array(expected_target_cdf).reshape(
(len(expected_target_cdf), ))
expected_target_crps = np.array(expected_target_crps).reshape(
(len(expected_target_crps), ))
assert all(np.isclose(distr.cdf(target).asnumpy(), expected_target_cdf))
assert all(np.isclose(distr.crps(target).asnumpy(), expected_target_crps))
# compare with empirical cdf from samples
num_samples = 100_000
samples = distr.sample(num_samples).asnumpy()
assert np.isfinite(samples).all()
emp_cdf, edges = empirical_cdf(samples)
calc_cdf = distr.cdf(mx.nd.array(edges)).asnumpy()
assert np.allclose(calc_cdf[1:, :], emp_cdf, atol=1e-2)
def test_simple_symmetric():
gamma = mx.nd.array([-1.0])
slopes = mx.nd.array([[2.0, 2.0]])
knot_spacings = mx.nd.array([[0.5, 0.5]])
distr = PiecewiseLinear(gamma=gamma,
slopes=slopes,
knot_spacings=knot_spacings)
assert distr.cdf(mx.nd.array([-2.0])).asnumpy().item() == 0.0
assert distr.cdf(mx.nd.array([+2.0])).asnumpy().item() == 1.0
expected_crps = np.array([1.0 + 2.0 / 3.0])
assert np.allclose(
distr.crps(mx.nd.array([-2.0])).asnumpy(), expected_crps)
assert np.allclose(distr.crps(mx.nd.array([2.0])).asnumpy(), expected_crps)