def autocorrelation(arg: List[numpy_types], index: Tuple[int, ...] = ()) -> ndarray:
check_all_floats(arg)
check_all_shapes_match([elem.shape for elem in arg])
check_index_is_valid(arg[0].shape, index)
arg_array = stack(arg, axis=-1)[index]
autocorr = k.jvm_view().Autocorrelation.calculate(k.to_java_array(arg_array))
return fromiter(autocorr, arg[0].dtype)
def test_throws_on_index_longer_than_shape() -> None:
with pytest.raises(
ValueError,
match=
r"Length of desired index \(3, 1, 3\) must match the length of the shape \(2, 2\)"
):
check_index_is_valid((2, 2), (3, 1, 3))
def test_throws_on_index_longer_than_shape() -> None:
with pytest.raises(ValueError):
check_index_is_valid((2, 2), (3, 1, 3))
def test_doesnt_throw_on_valid_index() -> None:
check_index_is_valid((3, 3), (2, 2))
def test_throws_on_index_out_of_bounds() -> None:
with pytest.raises(ValueError):
check_index_is_valid((2, 2, 6), (1, 1, 8))
def test_throws_on_index_out_of_bounds() -> None:
with pytest.raises(
ValueError,
match=r"Invalid index \(1, 1, 8\) for shape \(2, 2, 6\)"):
check_index_is_valid((2, 2, 6), (1, 1, 8))