diff --git a/changelog.d/fix-empty-vector-parameter-lookup.fixed.md b/changelog.d/fix-empty-vector-parameter-lookup.fixed.md new file mode 100644 index 000000000..f0a63ad80 --- /dev/null +++ b/changelog.d/fix-empty-vector-parameter-lookup.fixed.md @@ -0,0 +1 @@ +Looking up parameters with empty key arrays now returns an empty array at every level of a chained lookup, instead of raising `IndexError` once a level holds numeric values; after such a selection a one-element key gives no rows if it names a child and raises `ParameterNotFoundError` if it does not. diff --git a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py index ea7df3d2a..41563bdbd 100644 --- a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py +++ b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py @@ -341,7 +341,18 @@ def __getitem__(self, key: str) -> Any: dtypes_match = all(val.dtype == values[0].dtype for val in values) v0_len = len(values[0]) - if v0_len <= 1: + if v0_len == 0: + # An earlier selection left no rows. Broadcast the key + # over them as the leaf path does: no rows, or an error + # for a key of two or more. + numpy.broadcast_shapes(numpy.shape(idx), (0,)) + if dtypes_match: + result = values[0][:0] + else: + result = numpy.zeros( + 0, dtype=_unify_structured_dtypes(values)[0] + ) + elif v0_len == 1: # 1-element structured arrays: simple concat + index if not dtypes_match: unified_dtype, all_fields, values_cast = ( @@ -409,7 +420,7 @@ def __getitem__(self, key: str) -> Any: # or N-element vectors (after prior vectorial indexing). if values: v0 = numpy.asarray(values[0]) - if v0.ndim == 0 or v0.shape[0] <= 1: + if v0.ndim == 0 or v0.shape[0] == 1: # Scalar per child: 1D lookup scalar_vals = numpy.empty(len(values) + 1, dtype=numpy.float64) for i, v in enumerate(values): @@ -417,7 +428,8 @@ def __getitem__(self, key: str) -> Any: scalar_vals[-1] = numpy.nan result = scalar_vals[idx] else: - # N-element vectors: stack into (K+1, N) matrix + # N-element vectors (N may be 0 after an empty + # selection): stack into (K+1, N) matrix m = v0.shape[0] stacked = numpy.empty((len(values) + 1, m), dtype=numpy.float64) for i, v in enumerate(values): @@ -427,8 +439,10 @@ def __getitem__(self, key: str) -> Any: else: result = numpy.full(n, numpy.nan) - # Check for unexpected keys - if helpers.contains_nan(result): + # Check for unexpected keys. A key that names no child points at + # the NaN sentinel, but the result can hide it: a one-element key + # broadcast over zero rows selects nothing. + if helpers.contains_nan(result) or (idx == SENTINEL).any(): unexpected_keys = set( numpy.asarray(key, dtype=str) if not numpy.issubdtype(numpy.asarray(key).dtype, numpy.str_) diff --git a/tests/core/parameters_fancy_indexing/test_fancy_indexing.py b/tests/core/parameters_fancy_indexing/test_fancy_indexing.py index a3eaed648..d496daaef 100644 --- a/tests/core/parameters_fancy_indexing/test_fancy_indexing.py +++ b/tests/core/parameters_fancy_indexing/test_fancy_indexing.py @@ -85,6 +85,49 @@ def test_triple_fancy_indexing(): ) +def test_empty_fancy_indexing_at_every_level(): + # Selecting no rows used to raise IndexError at the third level: the + # numeric leaf held zero values and the lookup read its first one. + none = np.asarray([], dtype=str) + for result in [ + P.single.owner[none], + P.single[none][none], + P[none][none][none], + ]: + assert result.shape == (0,) + assert result.dtype == np.float64 + + +def test_empty_fancy_indexing_on_a_single_path(): + # The audit witness: one leaf a.b.c and three empty key arrays. + node = ParameterNode(data={"a": {"b": {"c": {"2020-01-01": 0}}}})("2020-01-01") + none = np.asarray([], dtype=str) + assert node[none][none][none].shape == (0,) + + +def test_one_key_after_an_empty_selection(): + # A one-element key broadcasts over the zero rows an empty selection + # leaves. A key naming a child gives no rows, at a node or a leaf; one + # naming no child raises, as it does for rows that exist. (At a node, a + # valid key used to return one row of NaN.) + node = ParameterNode( + data={"a": {"b": {"c": {"2020-01-01": 7}}, "d": {"c": {"2020-01-01": 8}}}} + )("2020-01-01") + none = np.asarray([], dtype=str) + + at_node = node[none][np.asarray(["b"])] + assert len(at_node.vector) == 0 + assert at_node[np.asarray(["c"])].shape == (0,) + assert node[none][none][np.asarray(["c"])].shape == (0,) + for lookup in ( + lambda: node[none][np.asarray(["missing"])], + lambda: node[none][none][np.asarray(["missing"])], + ): + with pytest.raises(ParameterNotFoundError) as e: + lookup() + assert "missing' was not found" in get_message(e.value) + + def test_wrong_key(): zone = np.asarray(["z1", "z2", "z2", "toto"]) with pytest.raises(ParameterNotFoundError) as e: diff --git a/tests/core/parameters_fancy_indexing/test_fancy_indexing_property.py b/tests/core/parameters_fancy_indexing/test_fancy_indexing_property.py new file mode 100644 index 000000000..b1c1a3564 --- /dev/null +++ b/tests/core/parameters_fancy_indexing/test_fancy_indexing_property.py @@ -0,0 +1,99 @@ +"""Property: vector parameter lookups agree row by row with scalar lookups. + +For any parameter tree of depth 1 to 3 (1 to 3 children per node, numeric +leaves) and any list of 0 to 20 paths through it, looking the paths up one +level at a time with key arrays returns an array of one value per path, each +equal to the scalar lookup of that path. Zero paths give an empty array; that +case used to raise ``IndexError`` once a level held numeric leaves. A key that +names no child raises ``ParameterNotFoundError`` at any level, including a +one-element key after a selection of no rows. Examples are in +``test_fancy_indexing.py``. +""" + +from __future__ import annotations + +import itertools + +import numpy as np +import pytest + +# The smoke job installs Core without the dev extra but collects every module. +pytest.importorskip("hypothesis") + +from hypothesis import given, settings # noqa: E402 +from hypothesis import strategies as st # noqa: E402 + +from policyengine_core.parameters import ( # noqa: E402 + ParameterNode, + ParameterNotFoundError, +) + + +@st.composite +def _tree_and_paths(draw): + widths = draw(st.lists(st.integers(1, 3), min_size=1, max_size=3)) + leaves = list(itertools.product(*[range(width) for width in widths])) + values = draw( + st.lists( + st.integers(-10_000, 10_000), + min_size=len(leaves), + max_size=len(leaves), + ) + ) + data = {} + for leaf, value in zip(leaves, values): + node = data + for depth, child in enumerate(leaf[:-1]): + node = node.setdefault(f"k{depth}_{child}", {}) + node[f"k{len(leaf) - 1}_{leaf[-1]}"] = {"2020-01-01": value} + paths = draw(st.lists(st.sampled_from(leaves), max_size=20)) + return widths, data, paths + + +@settings(max_examples=300, deadline=None) +@given(_tree_and_paths()) +def test_vector_lookup_matches_scalar_lookups(case): + widths, data, paths = case + node = ParameterNode(data=data)("2020-01-01") + + result = node + for depth in range(len(widths)): + result = result[ + np.asarray([f"k{depth}_{path[depth]}" for path in paths], dtype=str) + ] + + expected = [] + for path in paths: + scalar = node + for depth, child in enumerate(path): + scalar = scalar[f"k{depth}_{child}"] + expected.append(scalar) + + assert np.shape(result) == (len(paths),) + np.testing.assert_array_equal(result, np.asarray(expected, dtype=float)) + + +@settings(max_examples=300, deadline=None) +@given(_tree_and_paths(), st.data()) +def test_unknown_key_at_any_level_raises(case, data): + # The paths select rows (possibly none); then one level is asked for a + # key that names no child, among valid keys for the selected rows, or + # alone when no row is selected. + widths, data_tree, paths = case + node = ParameterNode(data=data_tree)("2020-01-01") + level = data.draw(st.integers(0, len(widths) - 1)) + + result = node + for depth in range(level): + result = result[ + np.asarray([f"k{depth}_{path[depth]}" for path in paths], dtype=str) + ] + keys = [f"k{level}_{path[level]}" for path in paths] + if keys: + keys[data.draw(st.integers(0, len(keys) - 1))] = "unknown" + else: + # No rows selected: a one-element key broadcasts over them. + keys = ["unknown"] + + with pytest.raises(ParameterNotFoundError): + result[np.asarray(keys, dtype=str)]