add new implementation for data_collector module

Author: SimonBlanke

src/surfaces/data_collector_new/__init__.py

@@ -0,0 +1,2 @@
+from .search_data_collector import SearchDataCollector
+from .search_data_lookup import SearchDataLookup

src/surfaces/data_collector_new/grid_generator.py

@@ -0,0 +1,37 @@
+import numpy as np
+from itertools import product
+from typing import Dict, List, Tuple, Any
+import time
+import os
+
+
+class GridGenerator:
+    """
+    Generates parameter grids from search space definitions.
+
+    This class takes a search space dictionary and creates all possible
+    combinations of parameters for grid search evaluation.
+    """
+
+    def __init__(self, search_space: Dict[str, List[Any]]):
+        self.search_space = search_space
+        self.param_names = list(search_space.keys())
+        self.param_values = [search_space[name] for name in self.param_names]
+
+    def generate_grid(self) -> Tuple[np.ndarray, List[str]]:
+        """
+        Generate complete parameter grid.
+
+        Returns:
+            grid: numpy array of shape (n_combinations, n_params)
+            param_names: list of parameter names
+        """
+        # Create all combinations
+        combinations = list(product(*self.param_values))
+        grid = np.array(combinations)
+
+        return grid, self.param_names
+
+    def get_grid_size(self) -> int:
+        """Calculate total number of grid points."""
+        return np.prod([len(values) for values in self.param_values])

src/surfaces/data_collector_new/search_data_collector.py

@@ -0,0 +1,72 @@
+from typing import Dict, Any, List
+import numpy as np
+import time
+
+from .grid_generator import GridGenerator
+
+
+class SearchDataCollector:
+    """
+    Collects search data by evaluating objective functions on parameter grids.
+
+    This class handles the expensive computation of evaluating ML models
+    across a parameter grid and saves the results for future use.
+    """
+
+    def __init__(self, objective_function, search_space: Dict[str, List[Any]]):
+        self.objective_function = objective_function
+        self.search_space = search_space
+        self.grid_generator = GridGenerator(search_space)
+
+    def collect(self, verbose: bool = True) -> Dict[str, np.ndarray]:
+        """
+        Evaluate objective function on entire grid.
+
+        Returns dictionary containing:
+            - 'parameters': parameter grid
+            - 'scores': objective function values
+            - 'times': evaluation times in seconds
+            - 'param_names': parameter names
+        """
+        grid, param_names = self.grid_generator.generate_grid()
+        n_points = len(grid)
+
+        scores = np.zeros(n_points)
+        times = np.zeros(n_points)
+
+        for i, params in enumerate(grid):
+            if verbose and i % 100 == 0:
+                print(f"Evaluating point {i+1}/{n_points}")
+
+            # Convert to dictionary for objective function
+            param_dict = {name: value for name, value in zip(param_names, params)}
+
+            # Time the evaluation
+            start_time = time.perf_counter()
+            scores[i] = self.objective_function(param_dict)
+            times[i] = time.perf_counter() - start_time
+
+        return {
+            "parameters": grid,
+            "scores": scores,
+            "times": times,
+            "param_names": np.array(param_names, dtype="U"),  # Unicode string array
+        }
+
+    def save(self, filepath: str, verbose: bool = True):
+        """Collect data and save to file."""
+        data = self.collect(verbose=verbose)
+
+        # Add metadata
+        data["search_space_keys"] = np.array(list(self.search_space.keys()), dtype="U")
+        data["search_space_sizes"] = np.array(
+            [len(v) for v in self.search_space.values()]
+        )
+
+        # Save as compressed numpy archive
+        np.savez_compressed(filepath, **data)
+
+        if verbose:
+            print(f"Saved search data to {filepath}")
+            print(f"Total points: {len(data['scores'])}")
+            print(f"Total evaluation time: {np.sum(data['times']):.2f} seconds")

src/surfaces/data_collector_new/search_data_lookup.py

@@ -0,0 +1,76 @@
+import os
+from typing import Dict, Any, List
+
+
+class SearchDataLookup:
+    """
+    Provides fast lookup of pre-computed objective function values.
+
+    This class loads search data from disk and provides O(1) lookup
+    for parameter combinations that were evaluated during grid search.
+    """
+
+    def __init__(self, filepath: str):
+        self.filepath = filepath
+        self._load_data()
+        self._build_lookup_table()
+
+    def _load_data(self):
+        """Load search data from file."""
+        if not os.path.exists(self.filepath):
+            raise FileNotFoundError(f"Search data file not found: {self.filepath}")
+
+        data = np.load(self.filepath)
+        self.parameters = data["parameters"]
+        self.scores = data["scores"]
+        self.times = data["times"]
+        self.param_names = data["param_names"]
+
+        # Reconstruct search space structure
+        self.search_space_keys = data["search_space_keys"]
+        self.search_space_sizes = data["search_space_sizes"]
+
+    def _build_lookup_table(self):
+        """Build hash table for fast parameter lookup."""
+        self.lookup_table = {}
+
+        for i, params in enumerate(self.parameters):
+            # Create hashable key from parameters
+            key = tuple(params)
+            self.lookup_table[key] = {"score": self.scores[i], "time": self.times[i]}
+
+    def evaluate(self, param_dict: Dict[str, Any]) -> float:
+        """
+        Look up objective function value for given parameters.
+
+        Args:
+            param_dict: Dictionary of parameter names to values
+
+        Returns:
+            Objective function value
+
+        Raises:
+            KeyError: If parameter combination not found in search data
+        """
+        # Convert dict to tuple in correct order
+        param_values = [param_dict[name] for name in self.param_names]
+        key = tuple(param_values)
+
+        if key not in self.lookup_table:
+            raise KeyError(
+                f"Parameter combination not found in search data: {param_dict}"
+            )
+
+        return self.lookup_table[key]["score"]
+
+    def get_evaluation_time(self, param_dict: Dict[str, Any]) -> float:
+        """Get the original evaluation time for given parameters."""
+        param_values = [param_dict[name] for name in self.param_names]
+        key = tuple(param_values)
+
+        if key not in self.lookup_table:
+            raise KeyError(
+                f"Parameter combination not found in search data: {param_dict}"
+            )
+
+        return self.lookup_table[key]["time"]