ai-nk-cce 0.1.0__py3-none-any.whl

nk_model/biased_prediction_agent.py

@@ -0,0 +1,389 @@
+from typing import Dict, List, Tuple
+
+import numpy as np
+from scipy.optimize import minimize
+
+from src.nk_model.models import Item, NKParams
+from src.nk_model.nk_landscape import NKLandscape
+
+
+class BiasedPredictionAgent:
+    """
+    Agent that generates biased improvement suggestions on NK landscapes.
+
+    The agent applies a consistent bias G to the base landscape F, creating
+    a biased landscape F' = normalize(F * G). It then generates suggestions
+    from a maximum entropy distribution that satisfies mean and variance
+    constraints relative to F'.
+    """
+
+    def __init__(
+        self,
+        landscape: NKLandscape,
+        bias_seed: int | None = None,
+        bias_power: float = 1.0,
+    ):
+        """
+        Initialize the biased prediction agent.
+
+        Args:
+            landscape: Base NK landscape F
+            bias_seed: Random seed for generating bias landscape G
+            bias_power: Power applied to bias landscape G. Higher values
+                make the bias more pronounced for regions with higher G(x)
+        """
+        self.landscape = landscape
+        self.bias_seed = bias_seed
+        self.bias_power = bias_power
+
+        # Generate bias landscape G (N'=N, K'=0)
+        self.bias_landscape = self._create_bias()
+
+        # Create biased landscape F' = normalize(F * G^power)
+        self.biased_landscape = self._create_biased_landscape()
+
+        # Cache for distributions: key=(x_tuple, radius, mu, sigma2),
+        # value=List of (coordinates, probability) tuples
+        self._distribution_cache: Dict[
+            Tuple[Tuple[int, ...], int, float, float],
+            List[Tuple[np.ndarray, float]],
+        ] = {}
+
+    def _create_bias(self) -> NKLandscape:
+        """
+        Create bias landscape G with N'=N, K'=0.
+
+        Returns:
+            NKLandscape with K=0 (no interactions)
+        """
+        if self.bias_seed is not None:
+            np.random.seed(self.bias_seed)
+
+        bias_params = NKParams(
+            n=self.landscape.N,
+            k=0,  # K'=0 as specified
+            m=self.landscape.N,  # Default m
+            power=self.bias_power,
+            max_val=1.0,
+            payoff_type="float",
+        )
+        return NKLandscape(params=bias_params)
+
+    def _create_biased_landscape(self) -> NKLandscape:
+        """
+        Create biased landscape F' = normalize(F * G^power).
+
+        Returns:
+            NKLandscape with payoffs F'(x) = normalize(F(x) * G(x)^power)
+        """
+        # Multiply payoffs: F'(x) = F(x) * G(x)^power
+        combined_items = []
+        for item_f, item_g in zip(
+            self.landscape.items, self.bias_landscape.items
+        ):
+            assert np.array_equal(
+                item_f.coordinates, item_g.coordinates
+            ), "Coordinates must match"
+            combined_payoff = item_f.payoff * (item_g.payoff)
+            combined_items.append(
+                Item(
+                    coordinates=item_f.coordinates.copy(),
+                    payoff=combined_payoff,
+                )
+            )
+
+        # Normalize to [0, 1]
+        payoffs = np.array([item.payoff for item in combined_items])
+        min_payoff = np.min(payoffs)
+        max_payoff = np.max(payoffs)
+        if max_payoff > min_payoff:
+            normalized_payoffs = (payoffs - min_payoff) / (
+                max_payoff - min_payoff
+            )
+        else:
+            normalized_payoffs = np.zeros_like(payoffs)
+
+        # Create new landscape with normalized payoffs
+        normalized_items = [
+            Item(
+                coordinates=item.coordinates.copy(),
+                payoff=normalized_payoffs[i],
+            )
+            for i, item in enumerate(combined_items)
+        ]
+
+        # Create a new NKLandscape-like object or store items
+        # For now, we'll store items and create lookup
+        biased_landscape = NKLandscape.__new__(NKLandscape)
+        biased_landscape.params = self.landscape.params
+        biased_landscape.N = self.landscape.N
+        biased_landscape.K = self.landscape.K
+        biased_landscape.M = self.landscape.M
+        biased_landscape.items = normalized_items
+        biased_landscape._payoff_lookup = {
+            tuple(item.coordinates): item.payoff for item in normalized_items
+        }
+        biased_landscape.uuid = f"{self.landscape.uuid}_biased"
+
+        return biased_landscape
+
+    def get_suggestion(
+        self,
+        x: np.ndarray,
+        radius: int,
+        delta_mu: float,
+        sigma2: float,
+    ) -> Item:
+        """
+        Generate a biased suggestion from distribution p_x.
+
+        Args:
+            x: Current point in landscape
+            radius: Locality parameter r (Hamming distance radius)
+            delta_mu: Expected improvement delta over F'(x)
+            sigma2: Variance of improvement suggestions
+
+        Returns:
+            Suggested Item from original (unbiased) landscape
+        """
+        # Compute absolute mu from delta_mu
+        current_fitness = self.landscape.get_payoff(x)
+        mu = current_fitness + delta_mu
+
+        # Get or compute distribution over all landscape items
+        item_distribution = self.generate_prediction_distribution(
+            x, radius, mu, sigma2
+        )
+
+        # Extract coordinates and probabilities
+        coordinates = [coord for coord, _ in item_distribution]
+        probabilities = np.array([prob for _, prob in item_distribution])
+
+        # Sample from distribution
+        sampled_idx = np.random.choice(len(coordinates), p=probabilities)
+        sampled_coords = coordinates[sampled_idx]
+
+        # Find and return the corresponding Item from original landscape
+        for item in self.landscape.items:
+            if np.array_equal(item.coordinates, sampled_coords):
+                return item
+
+        raise ValueError(
+            f"Item with coordinates {sampled_coords} not found"
+        )
+    
+    def generate_prediction_distribution(
+        self,
+        x: np.ndarray,
+        radius: int,
+        mu: float,
+        sigma2: float,
+    ) -> List[Tuple[np.ndarray, float]]:
+        """
+        Generate prediction distribution p_x using maximum entropy.
+
+        Returns a distribution over ALL landscape items, where items
+        outside B_r(x) have probability 0.
+
+        Args:
+            x: Current point
+            radius: Locality parameter r (Hamming distance radius)
+            mu: Target mean fitness F'(y)
+            sigma2: Target variance of fitness F'(y)
+
+        Returns:
+            List of (coordinates, probability) tuples for all items
+        """
+        # Check cache
+        x_tuple = tuple(x)
+        cache_key = (x_tuple, radius, mu, sigma2)
+        if cache_key in self._distribution_cache:
+            return self._distribution_cache[cache_key]
+
+        # Get ball items from original landscape
+        ball_items = self.closed_ball(x, radius)
+        ball_item_coordinates = [item.coordinates.copy() for item in ball_items]
+
+        # Check feasibility
+        self.feasible_mean_and_variance(ball_items, mu, sigma2)
+
+        # Initialize distribution with zeros for all items
+        item_distribution = [
+            (item.coordinates.copy(), 0.0) for item in self.landscape.items
+        ]
+
+        # Get probabilities for ball items
+        ball_distribution = self.solve_constraint_minimization(
+            mu, sigma2, ball_item_coordinates
+        )
+
+        # Create lookup for ball items
+        ball_coord_to_prob = {
+            tuple(coord): prob for coord, prob in ball_distribution
+        }
+
+        # Update distribution with ball item probabilities
+        item_distribution = [
+            (
+                coords,
+                ball_coord_to_prob.get(tuple(coords), 0.0),
+            )
+            for coords, _ in item_distribution
+        ]
+
+        # Cache result
+        self._distribution_cache[cache_key] = item_distribution
+
+        return item_distribution
+
+    def feasible_mean_and_variance(
+        self,
+        ball_items: list[Item],
+        mu: float,
+        sigma2: float,
+    ):
+        """
+        Get feasible min and max fitness values and validate mu, sigma2.
+
+        Args:
+            ball_items: List of items in B_r(x) from original landscape
+            mu: Target mean fitness F'(y)
+            sigma2: Target variance of fitness F'(y)
+
+        Returns:
+            Tuple of (min_fitness, max_fitness) in B_r(x)
+
+        Raises:
+            ValueError: If mu or sigma2 are not feasible
+        """
+        fitnesses = [
+            self.biased_landscape.get_payoff(item.coordinates)
+            for item in ball_items
+        ]
+        min_fitness = float(np.min(fitnesses))
+        max_fitness = float(np.max(fitnesses))
+
+        # Check mean feasibility
+        if not (min_fitness <= mu <= max_fitness):
+            raise ValueError(
+                f"mu={mu} not feasible. Must be in "
+                f"[{min_fitness}, {max_fitness}]"
+            )
+
+        # Check variance feasibility (Bhatia-Davis inequality)
+        max_var = (mu - min_fitness) * (max_fitness - mu)
+        if sigma2 > max_var:
+            raise ValueError(
+                f"sigma2={sigma2} not feasible. Must be <= {max_var}"
+            )
+
+    def solve_constraint_minimization(
+        self,
+        mu: float,
+        sigma2: float,
+        ball_item_coordinates: List[np.ndarray],
+    ) -> List[Tuple[np.ndarray, float]]:
+        """
+        Solve for maximum entropy distribution p_x.
+
+        Uses scipy.optimize.minimize to find distribution that maximizes
+        entropy subject to mean and variance constraints.
+
+        Args:
+            mu: Target mean
+            sigma2: Target variance
+            ball_item_coordinates: List of coordinate arrays in B_r(x)
+
+        Returns:
+            List of (coordinates, probability) tuples for ball items
+        """
+        # Get fitnesses from biased landscape
+        fitnesses = np.array(
+            [
+                self.biased_landscape.get_payoff(coordinate)
+                for coordinate in ball_item_coordinates
+            ]
+        )
+
+        n_points = len(ball_item_coordinates)
+
+        # Objective: minimize negative entropy (maximize entropy)
+        def objective(p: np.ndarray) -> float:
+            return np.sum(p * np.log(p + 1e-10))
+
+        # Constraints
+        constraints = [
+            {
+                "type": "eq",
+                "fun": lambda p: np.sum(p * fitnesses) - mu,
+            },
+            {
+                "type": "eq",
+                "fun": lambda p: np.sum(p * (fitnesses - mu) ** 2) - sigma2,
+            },
+            {"type": "eq", "fun": lambda p: np.sum(p) - 1},
+        ]
+
+        # Bounds: probabilities must be in [0, 1]
+        bounds = [(0, 1) for _ in range(n_points)]
+
+        # Initial guess: uniform distribution
+        p0 = np.ones(n_points) / n_points
+
+        # Solve
+        result = minimize(
+            objective,
+            p0,
+            bounds=bounds,
+            constraints=constraints,
+            method="SLSQP",
+        )
+
+        if not result.success:
+            raise RuntimeError(
+                f"Optimization failed: {result.message}"
+            )
+
+        # Return list of (coordinates, probability) tuples
+        return [
+            (ball_item_coordinates[i].copy(), float(result.x[i]))
+            for i in range(n_points)
+        ]
+
+    def hamming_distance(
+        self,
+        x: np.ndarray,
+        y: np.ndarray,
+    ) -> int:
+        """Compute Hamming distance between two points."""
+        return int(np.sum(x != y))
+
+    def closed_ball(
+        self,
+        x: np.ndarray,
+        r: int,
+    ) -> list[Item]:
+        """
+        Get all points y such that d_H(x, y) <= r.
+
+        Args:
+            x: Center point
+            r: Radius (Hamming distance)
+
+        Returns:
+            List of Item objects in B_r(x)
+
+        Raises:
+            ValueError: If no items found in ball B_r(x)
+        """
+        ball_items = [
+            item
+            for item in self.landscape.items
+            if self.hamming_distance(x, item.coordinates) <= r
+        ]
+        if len(ball_items) == 0:
+            raise ValueError(
+                f"No items found in ball B_{r}(x). "
+                f"Point x may not be in landscape."
+            )
+        return ball_items