use decision tree to rule out invalid parameter combinations

Author: SimonBlanke

src/surfaces/_surrogates/README.md

@@ -223,6 +223,38 @@ Each ONNX model has an accompanying `.meta.json` file:
 }
 ```
 
+## Validity Model
+
+Some hyperparameter combinations are invalid (e.g., `n_neighbors > dataset_size` in KNN). The real function returns `NaN` for these cases.
+
+The surrogate system handles this by training a **validity classifier** alongside the regression model:
+
+1. During training, both valid and invalid samples are collected
+2. A binary classifier learns to predict validity
+3. During inference, validity is checked first
+4. Invalid combinations return `NaN`, just like the real function
+
+```python
+# Surrogate correctly returns NaN for invalid combinations
+func = KNeighborsClassifierFunction(use_surrogate=True)
+
+# Valid: returns score
+result = func({'n_neighbors': 5, 'cv': 5, 'dataset': digits_data, ...})
+# 0.9560
+
+# Invalid: returns NaN (n_neighbors too large for dataset)
+result = func({'n_neighbors': 140, 'cv': 5, 'dataset': iris_data, ...})
+# nan
+```
+
+Files for a function with validity model:
+```
+models/
+├── k_neighbors_classifier.onnx              # Regression model
+├── k_neighbors_classifier.validity.onnx     # Validity classifier
+└── k_neighbors_classifier.onnx.meta.json    # Metadata (has_validity_model: true)
+```
+
 ## Limitations
 
 1. **Interpolation only**: Surrogates work best within the training search space

src/surfaces/_surrogates/_surrogate_loader.py

@@ -44,13 +44,15 @@ def __init__(
         self.metadata_path = metadata_path or self.model_path.with_suffix(
             self.model_path.suffix + ".meta.json"
         )
+        self.validity_model_path = self.model_path.with_suffix(".validity.onnx")
 
         self._session = None
+        self._validity_session = None
         self._metadata = None
 
     @property
     def session(self):
-        """Lazy-load ONNX runtime session."""
+        """Lazy-load ONNX runtime session for regression model."""
         if self._session is None:
             try:
                 import onnxruntime as ort
@@ -66,6 +68,35 @@ def session(self):
             )
         return self._session
 
+    @property
+    def validity_session(self):
+        """Lazy-load ONNX runtime session for validity model."""
+        if self._validity_session is None:
+            if not self.has_validity_model:
+                return None
+
+            try:
+                import onnxruntime as ort
+            except ImportError:
+                raise ImportError(
+                    "onnxruntime is required for surrogate models. "
+                    "Install it with: pip install onnxruntime"
+                )
+
+            self._validity_session = ort.InferenceSession(
+                str(self.validity_model_path),
+                providers=["CPUExecutionProvider"],
+            )
+        return self._validity_session
+
+    @property
+    def has_validity_model(self) -> bool:
+        """Check if a validity model exists."""
+        return (
+            self.metadata.get("has_validity_model", False)
+            and self.validity_model_path.exists()
+        )
+
     @property
     def metadata(self) -> Dict[str, Any]:
         """Load metadata from JSON file."""
@@ -114,6 +145,30 @@ def _encode_params(self, params: Dict[str, Any]) -> np.ndarray:
 
         return np.array([values], dtype=np.float32)
 
+    def is_valid(self, params: Dict[str, Any]) -> bool:
+        """Check if parameter combination is valid.
+
+        Parameters
+        ----------
+        params : dict
+            Parameter dictionary.
+
+        Returns
+        -------
+        bool
+            True if valid, False if invalid (would return NaN).
+        """
+        if not self.has_validity_model:
+            return True  # No validity model, assume all valid
+
+        input_array = self._encode_params(params)
+        input_name = self.validity_session.get_inputs()[0].name
+        output = self.validity_session.run(None, {input_name: input_array})
+
+        # Output is class label (0=invalid, 1=valid)
+        predicted_class = int(output[0][0])
+        return predicted_class == 1
+
     def predict(self, params: Dict[str, Any]) -> float:
         """Run inference on the surrogate model.
 
@@ -125,8 +180,12 @@ def predict(self, params: Dict[str, Any]) -> float:
         Returns
         -------
         float
-            Predicted objective value.
+            Predicted objective value, or NaN if parameters are invalid.
         """
+        # Check validity first
+        if not self.is_valid(params):
+            return float("nan")
+
         input_array = self._encode_params(params)
         input_name = self.session.get_inputs()[0].name
         output = self.session.run(None, {input_name: input_array})

src/surfaces/_surrogates/_surrogate_trainer.py

@@ -53,9 +53,12 @@ def __init__(
 
         self.X: Optional[np.ndarray] = None
         self.y: Optional[np.ndarray] = None
+        self.X_all: Optional[np.ndarray] = None  # All samples (valid + invalid)
+        self.y_valid: Optional[np.ndarray] = None  # Validity labels (0/1)
         self.param_names: List[str] = []
         self.param_encodings: Dict[str, Dict[str, int]] = {}
         self.model = None
+        self.validity_model = None
 
         self._training_time: float = 0
         self._collection_time: float = 0
@@ -145,8 +148,10 @@ def collect_samples_grid(
             grid_points = [grid_points[i] for i in indices]
 
         n_samples = len(grid_points)
-        X_list = []
+        X_valid_list = []
         y_list = []
+        X_all_list = []
+        validity_list = []
 
         if verbose:
             print(f"Collecting {n_samples} samples...")
@@ -169,29 +174,42 @@ def collect_samples_grid(
             # Evaluate function (use pure_objective_function to get raw value)
             try:
                 score = self.function.pure_objective_function(params)
-                # Skip NaN values (can happen with invalid hyperparameter combos)
+
+                # Track all samples for validity model
+                X_all_list.append(x_row)
+
                 if np.isnan(score):
-                    continue
-                X_list.append(x_row)
-                y_list.append(score)
+                    # Invalid combination
+                    validity_list.append(0)
+                else:
+                    # Valid combination
+                    validity_list.append(1)
+                    X_valid_list.append(x_row)
+                    y_list.append(score)
 
                 if verbose and (i + 1) % 100 == 0:
                     print(f"  Collected {len(y_list)}/{n_samples} valid samples")
             except Exception as e:
+                # Treat exceptions as invalid
+                X_all_list.append(x_row)
+                validity_list.append(0)
                 if verbose:
                     print(f"  Error at sample {i}: {e}")
 
-        self.X = np.array(X_list, dtype=np.float32)
+        self.X = np.array(X_valid_list, dtype=np.float32)
         self.y = np.array(y_list, dtype=np.float32)
+        self.X_all = np.array(X_all_list, dtype=np.float32)
+        self.y_valid = np.array(validity_list, dtype=np.int32)
 
         self._collection_time = time.time() - start_time
 
+        n_valid = len(self.y)
+        n_invalid = len(self.y_valid) - n_valid
+
         if verbose:
-            n_valid = len(self.y)
-            n_skipped = n_samples - n_valid
             print(f"Collected {n_valid} valid samples in {self._collection_time:.1f}s")
-            if n_skipped > 0:
-                print(f"  Skipped {n_skipped} samples (NaN or errors)")
+            if n_invalid > 0:
+                print(f"  Invalid samples: {n_invalid} (will train validity model)")
             if n_valid > 0:
                 print(f"  y range: [{self.y.min():.4f}, {self.y.max():.4f}]")
 
@@ -205,6 +223,8 @@ def train(
     ):
         """Train an MLP regressor on collected samples.
 
+        Also trains a validity classifier if invalid samples were found.
+
         Parameters
         ----------
         hidden_layer_sizes : tuple
@@ -225,11 +245,13 @@ def train(
 
         start_time = time.time()
 
-        # Normalize inputs
+        # Normalize inputs for regression model
         self.scaler_X = StandardScaler()
         X_scaled = self.scaler_X.fit_transform(self.X)
 
-        # Train MLP
+        # Train regression MLP
+        if verbose:
+            print("Training regression model...")
         self.model = MLPRegressor(
             hidden_layer_sizes=hidden_layer_sizes,
             max_iter=max_iter,
@@ -240,17 +262,43 @@ def train(
         )
         self.model.fit(X_scaled, self.y)
 
-        self._training_time = time.time() - start_time
-
-        # Evaluate on training data
+        # Evaluate regression on training data
         y_pred = self.model.predict(X_scaled)
         mse = np.mean((self.y - y_pred) ** 2)
         r2 = 1 - mse / np.var(self.y)
 
+        # Train validity classifier if there are invalid samples
+        n_invalid = np.sum(self.y_valid == 0)
+        if n_invalid > 0:
+            if verbose:
+                print("\nTraining validity classifier (DecisionTree)...")
+
+            from sklearn.tree import DecisionTreeClassifier
+
+            # Decision tree doesn't need scaling, but we keep scaler for API consistency
+            self.scaler_X_validity = None
+
+            self.validity_model = DecisionTreeClassifier(
+                max_depth=10,
+                min_samples_leaf=5,
+                random_state=42,
+            )
+            self.validity_model.fit(self.X_all, self.y_valid)
+
+            # Evaluate validity classifier
+            validity_pred = self.validity_model.predict(self.X_all)
+            validity_acc = np.mean(validity_pred == self.y_valid)
+
+            if verbose:
+                print(f"  Validity classifier accuracy: {validity_acc:.4f}")
+                print(f"  Tree depth: {self.validity_model.get_depth()}")
+
+        self._training_time = time.time() - start_time
+
         if verbose:
             print(f"\nTraining completed in {self._training_time:.1f}s")
-            print(f"  MSE: {mse:.6f}")
-            print(f"  R2:  {r2:.4f}")
+            print(f"  Regression MSE: {mse:.6f}")
+            print(f"  Regression R2:  {r2:.4f}")
 
     def export(
         self,
@@ -294,7 +342,7 @@ def export(
             ("mlp", self.model),
         ])
 
-        # Convert to ONNX
+        # Convert regression model to ONNX
         n_features = self.X.shape[1]
         initial_type = [("input", FloatTensorType([None, n_features]))]
         onnx_model = convert_sklearn(pipeline, initial_types=initial_type)
@@ -303,12 +351,33 @@ def export(
         with open(output_path, "wb") as f:
             f.write(onnx_model.SerializeToString())
 
+        # Export validity model if it exists
+        has_validity_model = self.validity_model is not None
+        if has_validity_model:
+            validity_path = output_path.with_suffix(".validity.onnx")
+
+            # DecisionTree doesn't need a scaler pipeline
+            onnx_validity = convert_sklearn(
+                self.validity_model,
+                initial_types=initial_type,
+                options={id(self.validity_model): {"zipmap": False}},
+            )
+
+            with open(validity_path, "wb") as f:
+                f.write(onnx_validity.SerializeToString())
+
+            if verbose:
+                print(f"Exported validity model to: {validity_path}")
+
         # Save metadata
+        n_invalid = int(np.sum(self.y_valid == 0))
         metadata = {
             "function_name": getattr(self.function, "_name_", self.function.__class__.__name__),
             "param_names": self.param_names,
             "param_encodings": self.param_encodings,
             "n_samples": len(self.y),
+            "n_invalid_samples": n_invalid,
+            "has_validity_model": has_validity_model,
             "y_range": [float(self.y.min()), float(self.y.max())],
             "training_time": self._training_time,
             "collection_time": self._collection_time,

src/surfaces/_surrogates/models/gradient_boosting_regressor.onnx

@@ -12,10 +12,12 @@
     }
   },
   "n_samples": 1000,
+  "n_invalid_samples": 0,
+  "has_validity_model": false,
   "y_range": [
-    -0.18430274724960327,
-    0.4645007848739624
+    -0.188394695520401,
+    0.46403050422668457
   ],
-  "training_time": 0.9513494968414307,
-  "collection_time": 571.8018696308136
+  "training_time": 0.9096114635467529,
+  "collection_time": 581.1047446727753
 }

src/surfaces/_surrogates/models/gradient_boosting_regressor.onnx.meta.json

@@ -19,11 +19,13 @@
       "iris_data": 2
     }
   },
-  "n_samples": 866,
+  "n_samples": 868,
+  "n_invalid_samples": 132,
+  "has_validity_model": true,
   "y_range": [
     0.39886364340782166,
     0.9802631735801697
   ],
-  "training_time": 1.0907559394836426,
-  "collection_time": 40.40663194656372
+  "training_time": 0.5762979984283447,
+  "collection_time": 41.40630507469177
 }