morphml 1.0.0__py3-none-any.whl

morphml/optimizers/bayesian/acquisition.py

@@ -0,0 +1,387 @@
+"""Acquisition functions for Bayesian optimization.
+
+Acquisition functions balance exploration (uncertain regions) and exploitation
+(promising regions) by quantifying the value of sampling at a given point.
+
+Common acquisition functions:
+- Expected Improvement (EI): Expected gain over current best
+- Upper Confidence Bound (UCB): Optimistic estimate
+- Probability of Improvement (PI): Probability of beating current best
+
+Author: Eshan Roy <eshanized@proton.me>
+Organization: TONMOY INFRASTRUCTURE & VISION
+"""
+
+from typing import Callable, List, Optional, Tuple
+
+import numpy as np
+from scipy.optimize import differential_evolution, minimize
+from scipy.stats import norm
+
+from morphml.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+def expected_improvement(
+    mu: np.ndarray, sigma: np.ndarray, f_best: float, xi: float = 0.01
+) -> np.ndarray:
+    """
+    Expected Improvement acquisition function.
+
+    EI balances exploration and exploitation by computing the expected
+    amount by which a point improves over the current best.
+
+    Formula:
+        EI(x) = E[max(f(x) - f*, 0)]
+              = (μ - f* - ξ) * Φ(Z) + σ * φ(Z)
+        where Z = (μ - f* - ξ) / σ
+
+    Args:
+        mu: Predicted mean(s)
+        sigma: Predicted standard deviation(s)
+        f_best: Current best fitness value
+        xi: Exploration parameter (higher = more exploration)
+
+    Returns:
+        Expected improvement value(s)
+
+    Example:
+        >>> mu = np.array([0.5, 0.7, 0.3])
+        >>> sigma = np.array([0.1, 0.2, 0.05])
+        >>> f_best = 0.6
+        >>> ei = expected_improvement(mu, sigma, f_best, xi=0.01)
+    """
+    with np.errstate(divide="warn", invalid="warn"):
+        # Compute improvement
+        imp = mu - f_best - xi
+
+        # Compute Z-score
+        Z = imp / sigma
+
+        # Expected improvement
+        ei = imp * norm.cdf(Z) + sigma * norm.pdf(Z)
+
+        # Handle zero sigma (no uncertainty)
+        ei[sigma == 0.0] = 0.0
+
+    return ei
+
+
+def upper_confidence_bound(mu: np.ndarray, sigma: np.ndarray, kappa: float = 2.576) -> np.ndarray:
+    """
+    Upper Confidence Bound acquisition function.
+
+    UCB provides an optimistic estimate by adding a multiple of the
+    uncertainty to the predicted mean. The kappa parameter controls
+    the exploration-exploitation trade-off.
+
+    Formula:
+        UCB(x) = μ(x) + κ * σ(x)
+
+    Common kappa values:
+    - 1.96: 95% confidence (moderate exploration)
+    - 2.576: 99% confidence (high exploration)
+    - 1.0: Balanced
+
+    Args:
+        mu: Predicted mean(s)
+        sigma: Predicted standard deviation(s)
+        kappa: Exploration parameter (higher = more exploration)
+
+    Returns:
+        UCB value(s)
+
+    Example:
+        >>> mu = np.array([0.5, 0.7])
+        >>> sigma = np.array([0.1, 0.2])
+        >>> ucb = upper_confidence_bound(mu, sigma, kappa=2.0)
+    """
+    return mu + kappa * sigma
+
+
+def lower_confidence_bound(mu: np.ndarray, sigma: np.ndarray, kappa: float = 2.576) -> np.ndarray:
+    """
+    Lower Confidence Bound acquisition function (for minimization).
+
+    LCB is the opposite of UCB, used for minimization problems.
+    It balances exploitation (low mean) with exploration (high uncertainty).
+
+    Args:
+        mu: Predicted means from the surrogate model
+        sigma: Predicted standard deviations
+        kappa: Exploration-exploitation trade-off parameter.
+               Higher values favor exploration.
+               Default 2.576 corresponds to 99% confidence interval.
+
+    Returns:
+        LCB values for each point
+
+    Example:
+        >>> mu = np.array([0.5, 0.7])
+        >>> sigma = np.array([0.1, 0.2])
+        >>> lcb = lower_confidence_bound(mu, sigma, kappa=2.0)
+    """
+    return mu - kappa * sigma
+
+
+def probability_of_improvement(
+    mu: np.ndarray, sigma: np.ndarray, f_best: float, xi: float = 0.01
+) -> np.ndarray:
+    """
+    Probability of Improvement acquisition function.
+
+    PI computes the probability that a point will improve over the
+    current best. More conservative than EI.
+
+    Formula:
+        PI(x) = P(f(x) > f*)
+              = Φ((μ - f* - ξ) / σ)
+
+    Args:
+        mu: Predicted mean(s)
+        sigma: Predicted standard deviation(s)
+        f_best: Current best fitness value
+        xi: Exploration parameter
+
+    Returns:
+        Probability of improvement value(s)
+
+    Example:
+        >>> mu = np.array([0.5, 0.7])
+        >>> sigma = np.array([0.1, 0.2])
+        >>> f_best = 0.6
+        >>> pi = probability_of_improvement(mu, sigma, f_best)
+    """
+    with np.errstate(divide="warn", invalid="warn"):
+        Z = (mu - f_best - xi) / sigma
+        pi = norm.cdf(Z)
+
+        # Handle zero sigma
+        pi[sigma == 0.0] = 0.0
+
+    return pi
+
+
+def thompson_sampling(
+    mu: np.ndarray, sigma: np.ndarray, random_state: Optional[int] = None
+) -> np.ndarray:
+    """
+    Thompson Sampling for acquisition.
+
+    Sample from the posterior distribution and select the point
+    with the highest sample. Naturally balances exploration/exploitation.
+
+    Args:
+        mu: Predicted mean(s)
+        sigma: Predicted standard deviation(s)
+        random_state: Random seed for reproducibility
+
+    Returns:
+        Sampled values from posterior
+    """
+    if random_state is not None:
+        np.random.seed(random_state)
+
+    samples = np.random.normal(mu, sigma)
+    return samples
+
+
+class AcquisitionOptimizer:
+    """
+    Optimizer for acquisition functions.
+
+    Finds the point that maximizes the acquisition function value,
+    which determines where to sample next.
+
+    Attributes:
+        method: Optimization method ('lbfgs', 'de', 'random')
+        n_restarts: Number of random restarts for local optimization
+        n_samples: Number of random samples for 'random' method
+
+    Example:
+        >>> def my_acquisition(x):
+        ...     return expected_improvement(gp.predict(x), f_best=0.8)
+        >>> optimizer = AcquisitionOptimizer(method='lbfgs', n_restarts=10)
+        >>> x_next = optimizer.optimize(my_acquisition, bounds)
+    """
+
+    def __init__(
+        self,
+        method: str = "lbfgs",
+        n_restarts: int = 10,
+        n_samples: int = 1000,
+        random_state: Optional[int] = None,
+    ):
+        """
+        Initialize acquisition optimizer.
+
+        Args:
+            method: Optimization method ('lbfgs', 'de', 'random')
+            n_restarts: Number of random restarts for 'lbfgs'
+            n_samples: Number of samples for 'random' method
+            random_state: Random seed for reproducibility
+        """
+        self.method = method
+        self.n_restarts = n_restarts
+        self.n_samples = n_samples
+        self.random_state = random_state
+
+        if random_state is not None:
+            np.random.seed(random_state)
+
+    def optimize(
+        self,
+        acquisition_fn: Callable[[np.ndarray], float],
+        bounds: List[Tuple[float, float]],
+        n_candidates: int = 1,
+    ) -> np.ndarray:
+        """
+        Find point(s) that maximize acquisition function.
+
+        Args:
+            acquisition_fn: Function to maximize (takes array, returns scalar)
+            bounds: List of (min, max) tuples for each dimension
+            n_candidates: Number of candidates to return
+
+        Returns:
+            Best point(s) as numpy array of shape (n_candidates, n_dims)
+
+        Raises:
+            ValueError: If method is unknown
+        """
+        if self.method == "lbfgs":
+            return self._optimize_lbfgs(acquisition_fn, bounds, n_candidates)
+        elif self.method == "de":
+            return self._optimize_differential_evolution(acquisition_fn, bounds)
+        elif self.method == "random":
+            return self._optimize_random_search(acquisition_fn, bounds, n_candidates)
+        else:
+            raise ValueError(f"Unknown optimization method: {self.method}")
+
+    def _optimize_lbfgs(
+        self, acquisition_fn: Callable, bounds: List[Tuple[float, float]], n_candidates: int
+    ) -> np.ndarray:
+        """
+        Multi-start L-BFGS-B optimization.
+
+        Performs multiple local optimizations from random starting points
+        and returns the best result.
+        """
+        best_x = None
+        best_value = -np.inf
+
+        for _ in range(self.n_restarts):
+            # Random starting point
+            x0 = np.array([np.random.uniform(low, high) for low, high in bounds])
+
+            # Minimize negative (to maximize)
+            try:
+                result = minimize(
+                    lambda x: -acquisition_fn(x.reshape(1, -1))[0],
+                    x0=x0,
+                    bounds=bounds,
+                    method="L-BFGS-B",
+                    options={"maxiter": 100},
+                )
+
+                value = -result.fun
+                if value > best_value:
+                    best_value = value
+                    best_x = result.x
+
+            except Exception as e:
+                logger.warning(f"L-BFGS-B optimization failed: {e}")
+                continue
+
+        if best_x is None:
+            # Fallback to random sample
+            best_x = np.array([np.random.uniform(low, high) for low, high in bounds])
+
+        return best_x.reshape(1, -1) if n_candidates == 1 else best_x
+
+    def _optimize_differential_evolution(
+        self, acquisition_fn: Callable, bounds: List[Tuple[float, float]]
+    ) -> np.ndarray:
+        """
+        Global optimization using Differential Evolution.
+
+        More robust than L-BFGS but slower. Good for multimodal acquisitions.
+        """
+        try:
+            result = differential_evolution(
+                lambda x: -acquisition_fn(x.reshape(1, -1))[0],
+                bounds=bounds,
+                maxiter=100,
+                seed=self.random_state,
+                workers=1,
+                polish=True,
+            )
+            return result.x.reshape(1, -1)
+
+        except Exception as e:
+            logger.warning(f"Differential evolution failed: {e}")
+            # Fallback to random
+            return np.array([[np.random.uniform(low, high) for low, high in bounds]])
+
+    def _optimize_random_search(
+        self, acquisition_fn: Callable, bounds: List[Tuple[float, float]], n_candidates: int
+    ) -> np.ndarray:
+        """
+        Random search: sample many points and pick best.
+
+        Simple but surprisingly effective baseline.
+        """
+        # Generate random candidates
+        n_dims = len(bounds)
+        candidates = np.zeros((self.n_samples, n_dims))
+
+        for i, (low, high) in enumerate(bounds):
+            candidates[:, i] = np.random.uniform(low, high, self.n_samples)
+
+        # Evaluate all candidates
+        values = np.array([acquisition_fn(x.reshape(1, -1))[0] for x in candidates])
+
+        # Return top n_candidates
+        top_indices = np.argsort(values)[-n_candidates:][::-1]
+
+        return candidates[top_indices]
+
+
+def get_acquisition_function(name: str, **kwargs) -> Callable:
+    """
+    Factory function for acquisition functions.
+
+    Args:
+        name: Acquisition function name ('ei', 'ucb', 'pi', 'ts')
+        **kwargs: Additional parameters for the acquisition function
+
+    Returns:
+        Acquisition function
+
+    Example:
+        >>> acq_fn = get_acquisition_function('ei', f_best=0.8, xi=0.01)
+        >>> value = acq_fn(mu=0.9, sigma=0.1)
+    """
+    if name.lower() == "ei":
+        f_best = kwargs.get("f_best", 0.0)
+        xi = kwargs.get("xi", 0.01)
+        return lambda mu, sigma: expected_improvement(mu, sigma, f_best, xi)
+
+    elif name.lower() == "ucb":
+        kappa = kwargs.get("kappa", 2.576)
+        return lambda mu, sigma: upper_confidence_bound(mu, sigma, kappa)
+
+    elif name.lower() == "pi":
+        f_best = kwargs.get("f_best", 0.0)
+        xi = kwargs.get("xi", 0.01)
+        return lambda mu, sigma: probability_of_improvement(mu, sigma, f_best, xi)
+
+    elif name.lower() == "ts":
+        random_state = kwargs.get("random_state", None)
+        return lambda mu, sigma: thompson_sampling(mu, sigma, random_state)
+
+    else:
+        raise ValueError(
+            f"Unknown acquisition function: {name}. " f"Choose from: 'ei', 'ucb', 'pi', 'ts'"
+        )

morphml/optimizers/bayesian/base.py

@@ -0,0 +1,319 @@
+"""Base class for Bayesian optimization algorithms.
+
+This module provides the foundation for sample-efficient Bayesian optimization
+methods that use surrogate models to guide the search process.
+
+Author: Eshan Roy <eshanized@proton.me>
+Organization: TONMOY INFRASTRUCTURE & VISION
+"""
+
+from abc import abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from morphml.core.dsl import SearchSpace
+from morphml.core.graph import ModelGraph
+from morphml.core.search import Individual
+from morphml.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class BaseBayesianOptimizer:
+    """
+    Base class for Bayesian optimization algorithms.
+
+    Bayesian optimization uses a surrogate model to approximate the
+    expensive-to-evaluate fitness function, enabling intelligent
+    exploration-exploitation trade-offs.
+
+    Key components:
+    1. **Surrogate Model**: Approximates f(x) (e.g., GP, RF, TPE)
+    2. **Acquisition Function**: Decides where to sample next
+    3. **Architecture Encoding**: Maps graphs to continuous/discrete vectors
+
+    Attributes:
+        search_space: SearchSpace defining architecture options
+        config: Algorithm configuration
+        generation: Current generation/iteration
+        history: List of all evaluated architectures
+        best_individual: Best architecture found so far
+
+    Example:
+        >>> from morphml.optimizers.bayesian import GaussianProcessOptimizer
+        >>> optimizer = GaussianProcessOptimizer(
+        ...     search_space=space,
+        ...     config={'acquisition': 'ei', 'n_initial_points': 10}
+        ... )
+        >>> best = optimizer.optimize(evaluator)
+    """
+
+    def __init__(self, search_space: SearchSpace, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize Bayesian optimizer.
+
+        Args:
+            search_space: SearchSpace to sample architectures from
+            config: Algorithm configuration dictionary
+        """
+        self.search_space = search_space
+        self.config = config or {}
+
+        self.generation = 0
+        self.history: List[Dict[str, Any]] = []
+        self.best_individual: Optional[Individual] = None
+
+        # Configuration parameters
+        self.n_initial_points = self.config.get("n_initial_points", 10)
+        self.max_iterations = self.config.get("max_iterations", 100)
+        self.random_state = self.config.get("random_state", None)
+
+        if self.random_state is not None:
+            np.random.seed(self.random_state)
+
+        logger.info(
+            f"Initialized {self.__class__.__name__} with " f"{self.n_initial_points} initial points"
+        )
+
+    def optimize(
+        self, evaluator: Any, max_evaluations: Optional[int] = None, callback: Optional[Any] = None
+    ) -> Individual:
+        """
+        Run Bayesian optimization loop.
+
+        Args:
+            evaluator: Function that evaluates ModelGraph fitness
+            max_evaluations: Maximum number of evaluations (overrides config)
+            callback: Optional callback function called each iteration
+
+        Returns:
+            Best Individual found
+
+        Example:
+            >>> def my_evaluator(graph):
+            ...     return train_and_evaluate(graph)
+            >>> best = optimizer.optimize(my_evaluator, max_evaluations=50)
+        """
+        max_evals = max_evaluations or self.max_iterations
+
+        logger.info(f"Starting Bayesian optimization for {max_evals} evaluations")
+
+        for iteration in range(max_evals):
+            # Ask: Get next candidate(s) to evaluate
+            candidates = self.ask()
+
+            # Evaluate candidates
+            results = []
+            for graph in candidates:
+                fitness = evaluator(graph)
+                results.append((graph, fitness))
+
+                # Track best
+                individual = Individual(graph)
+                individual.fitness = fitness
+
+                if self.best_individual is None or fitness > self.best_individual.fitness:
+                    self.best_individual = individual
+                    logger.info(f"Iteration {iteration}: New best fitness = {fitness:.4f}")
+
+            # Tell: Update surrogate model with results
+            self.tell(results)
+
+            # Callback
+            if callback is not None:
+                callback(iteration, self.best_individual, self.history)
+
+            self.generation += 1
+
+        logger.info(f"Optimization complete. Best fitness: {self.best_individual.fitness:.4f}")
+
+        return self.best_individual
+
+    @abstractmethod
+    def ask(self) -> List[ModelGraph]:
+        """
+        Generate next candidate architecture(s) to evaluate.
+
+        Uses the surrogate model and acquisition function to select
+        promising architectures.
+
+        Returns:
+            List of ModelGraph candidates
+        """
+        pass
+
+    @abstractmethod
+    def tell(self, results: List[Tuple[ModelGraph, float]]) -> None:
+        """
+        Update surrogate model with evaluation results.
+
+        Args:
+            results: List of (graph, fitness) tuples
+        """
+        pass
+
+    def get_best(self) -> Optional[Individual]:
+        """
+        Get best individual found so far.
+
+        Returns:
+            Best Individual or None if no evaluations yet
+        """
+        return self.best_individual
+
+    def get_history(self) -> List[Dict[str, Any]]:
+        """
+        Get optimization history.
+
+        Returns:
+            List of dictionaries with generation, graph, fitness
+        """
+        return self.history
+
+    def reset(self) -> None:
+        """Reset optimizer to initial state."""
+        self.generation = 0
+        self.history = []
+        self.best_individual = None
+        logger.info(f"{self.__class__.__name__} reset to initial state")
+
+    def _encode_architecture(self, graph: ModelGraph) -> np.ndarray:
+        """
+        Encode ModelGraph as fixed-length numerical vector.
+
+        This is a critical method that maps complex graph structures
+        to continuous/discrete vectors suitable for surrogate models.
+
+        Encoding strategies:
+        1. **Positional Encoding**: Represent nodes by position in topological order
+        2. **Operation One-Hot**: Encode operation types
+        3. **Hyperparameters**: Include numerical parameters
+        4. **Connectivity**: Encode edge structure
+
+        Args:
+            graph: ModelGraph to encode
+
+        Returns:
+            Fixed-length numpy array
+        """
+        # Get topological ordering
+        try:
+            topo_order = list(graph.topological_sort())
+        except Exception:
+            # Handle invalid graphs
+            topo_order = list(graph.nodes.values())
+
+        # Define operation vocabulary
+        operation_types = [
+            "input",
+            "output",
+            "conv2d",
+            "dense",
+            "relu",
+            "sigmoid",
+            "tanh",
+            "maxpool",
+            "avgpool",
+            "batchnorm",
+            "dropout",
+            "flatten",
+            "add",
+            "concat",
+        ]
+
+        # Fixed encoding length (max depth)
+        max_depth = 20
+        encoding_per_node = 3  # operation_id, param1, param2
+
+        encoding = []
+
+        for i in range(max_depth):
+            if i < len(topo_order):
+                node = topo_order[i]
+
+                # Encode operation type
+                if node.operation in operation_types:
+                    op_id = operation_types.index(node.operation)
+                else:
+                    op_id = 0  # Unknown operation
+                encoding.append(float(op_id))
+
+                # Encode key hyperparameters
+                if node.operation == "conv2d":
+                    filters = node.params.get("filters", 32)
+                    kernel_size = node.params.get("kernel_size", 3)
+                    encoding.extend([float(filters), float(kernel_size)])
+
+                elif node.operation == "dense":
+                    units = node.params.get("units", 128)
+                    encoding.extend([float(units), 0.0])
+
+                elif node.operation == "dropout":
+                    rate = node.params.get("rate", 0.5)
+                    encoding.extend([float(rate * 100), 0.0])
+
+                elif node.operation in ["maxpool", "avgpool"]:
+                    pool_size = node.params.get("pool_size", 2)
+                    encoding.extend([float(pool_size), 0.0])
+
+                else:
+                    encoding.extend([0.0, 0.0])
+            else:
+                # Padding for shorter architectures
+                encoding.extend([0.0] * encoding_per_node)
+
+        return np.array(encoding, dtype=np.float64)
+
+    def _decode_architecture(self, x: np.ndarray) -> ModelGraph:
+        """
+        Decode numerical vector back to ModelGraph.
+
+        This is challenging because the mapping is many-to-one
+        (many vectors may decode to similar graphs).
+
+        Strategy:
+        1. Sample random architecture from search space
+        2. Use vector to guide mutations toward desired structure
+
+        Args:
+            x: Numerical encoding
+
+        Returns:
+            Decoded ModelGraph
+        """
+        # Simplified decoding: sample from search space
+        # In practice, this would use more sophisticated reconstruction
+        graph = self.search_space.sample()
+
+        # TODO: Could add logic to mutate graph toward target encoding
+        # For now, return sampled graph (acquisition still guides search)
+
+        return graph
+
+    def _get_encoding_bounds(self) -> List[Tuple[float, float]]:
+        """
+        Get bounds for architecture encoding dimensions.
+
+        Returns:
+            List of (min, max) tuples for each encoding dimension
+        """
+        operation_types_count = 14  # Number of supported operations
+        max_filters = 512
+        max_kernel_size = 7
+
+        max_depth = 20
+
+        bounds = []
+        for _ in range(max_depth):
+            bounds.append((0, operation_types_count))  # Operation ID
+            bounds.append((0, max_filters))  # Param 1 (filters/units)
+            bounds.append((0, max_kernel_size))  # Param 2 (kernel/pool size)
+
+        return bounds
+
+
+class BayesianOptimizationError(Exception):
+    """Exception raised for errors in Bayesian optimization."""
+
+    pass