Guard spatial lag enrichment for small universes

openavmkit/data.py

@@ -2673,6 +2673,16 @@ def _enrich_universe_spatial_lag(
 
     # we TRAIN on these coordinates -- coordinates that are NOT in the test set
     coords_train = df_train_univ[["latitude", "longitude"]].values
+
+    # Guard: with no training parcels there is nothing to build a neighbour tree
+    # from (cKDTree of an empty array raises), so return the universe unchanged --
+    # no spatial-lag columns are produced. This mirrors the sales-side guard in
+    # _enrich_sup_spatial_lag_for_model_group, which `continue`s when its training
+    # set is too small; the universe side was missing the equivalent check.
+    n_train = len(df_train_univ)
+    if n_train == 0:
+        return df
+
     tree = cKDTree(coords_train)
 
     # we PREDICT on these coordinates -- all the coordinates in the universe
@@ -2682,8 +2692,14 @@ def _enrich_universe_spatial_lag(
         if value_field not in df:
             continue
 
-        # Choose the number of nearest neighbors to use
-        k = value_fields[value_field]
+        # Choose the number of nearest neighbors to use, clamped to the number of
+        # training parcels available. cKDTree.query with k greater than the number
+        # of points returns the sentinel index `n_train` (out of range) for the
+        # missing neighbours, which then raises IndexError at
+        # `parcel_values[indices]` below. Clamping keeps the lag well-defined for a
+        # small universe (e.g. a sparse subject with only a handful of comps)
+        # instead of crashing the whole prediction.
+        k = min(value_fields[value_field], n_train)
 
         # Query the tree: for each parcel in df_universe, find the k nearest parcels
         # distances: shape (n_universe, k); indices: corresponding indices in df_sales

tests/test_data.py

@@ -1276,4 +1276,37 @@ def test_canonical_split_explicit_none_disables_default_rule():
 	# Legacy fill: all 53 test slots come from lookback; only 22 of 75 lookback
 	# sales remain for training.
 	assert test_2025 == 53
-	assert train_2025 == 22
+	assert train_2025 == 22
+
+
+def test_enrich_universe_spatial_lag_small_universe_does_not_crash():
+	# Regression (Fabrica ENG-3031): cKDTree.query returns an out-of-range
+	# sentinel index when the universe has fewer training parcels than k.
+	from openavmkit.data import _enrich_universe_spatial_lag
+
+	df_univ = pd.DataFrame({
+		"key": ["a", "b"],
+		"latitude": [38.89, 38.90],
+		"longitude": [-90.18, -90.19],
+		"model_group": ["vacant_land", "vacant_land"],
+		"land_area_sqft": [21780.0, 10890.0],
+		"bldg_area_finished_sqft": [0.0, 0.0],
+	})
+	df_test = pd.DataFrame({"key": pd.Series([], dtype=object)})
+	out = _enrich_universe_spatial_lag(
+		df_univ, df_test, "vacant_land", ["vacant_land"], {}
+	)
+	assert "spatial_lag_land_area_sqft" in out.columns
+
+	df_one = df_univ.iloc[[0]].copy()
+	df_test_one = pd.DataFrame({"key": pd.Series([], dtype=object)})
+	out_one = _enrich_universe_spatial_lag(
+		df_one, df_test_one, "vacant_land", ["vacant_land"], {}
+	)
+	assert "spatial_lag_land_area_sqft" in out_one.columns
+
+	df_test_all = pd.DataFrame({"key": ["a", "b"]})
+	out_empty = _enrich_universe_spatial_lag(
+		df_univ, df_test_all, "vacant_land", ["vacant_land"], {}
+	)
+	assert not any(c.startswith("spatial_lag_") for c in out_empty.columns)