diff --git a/src/pyrecest/filters/discrete_state/__init__.py b/src/pyrecest/filters/discrete_state/__init__.py index c70efe33a..bd2b119ec 100644 --- a/src/pyrecest/filters/discrete_state/__init__.py +++ b/src/pyrecest/filters/discrete_state/__init__.py @@ -341,6 +341,18 @@ def sparse_gaussian_transition_matrix( valid_state_mask=None, ): states = _validated_state_vectors(state_vectors) + if states.ndim == 1: + states = states[:, None] + elif states.ndim != 2: + raise ValueError( + "state_vectors must have shape (n_states,) or (n_states, state_dim)" + ) + n_states = states.shape[0] + if n_states == 0: + raise ValueError("state_vectors must contain at least one state") + if states.shape[1] == 0: + raise ValueError("state_vectors must contain at least one coordinate per state") + sigma = _validated_positive_scalar(sigma, "sigma") max_step_sigma = _validated_positive_scalar( max_step_sigma, @@ -348,9 +360,6 @@ def sparse_gaussian_transition_matrix( allow_infinite=True, ) - if states.ndim == 1: - states = states[:, None] - n_states = states.shape[0] valid_mask = _module_globals["_coerce_valid_state_mask"]( valid_state_mask, n_states, diff --git a/tests/filters/test_sparse_gaussian_transition_validation.py b/tests/filters/test_sparse_gaussian_transition_validation.py new file mode 100644 index 000000000..f0b38e778 --- /dev/null +++ b/tests/filters/test_sparse_gaussian_transition_validation.py @@ -0,0 +1,23 @@ +import numpy as np +import pytest + +from pyrecest.filters.discrete_state import sparse_gaussian_transition_matrix + + +@pytest.mark.parametrize( + ("state_vectors", "message"), + [ + (np.array(1.0), "state_vectors must have shape"), + (np.empty((2, 1, 1)), "state_vectors must have shape"), + (np.empty((0,)), "state_vectors must contain at least one state"), + ( + np.empty((2, 0)), + "state_vectors must contain at least one coordinate per state", + ), + ], +) +def test_sparse_gaussian_transition_matrix_rejects_invalid_state_shapes( + state_vectors, message +): + with pytest.raises(ValueError, match=message): + sparse_gaussian_transition_matrix(state_vectors, sigma=1.0)