Return an empty array from chained vector lookups that select no rows

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.

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,15 +420,16 @@ 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):
                             scalar_vals[i] = float(numpy.asarray(v).flat[0])
                         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_)

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:

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)]