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

nk_model/models.py

@@ -0,0 +1,172 @@
+import json
+from typing import Literal, Union
+
+import numpy as np
+from pydantic import (
+    BaseModel,
+    Field,
+    NonNegativeInt,
+    PositiveFloat,
+    PositiveInt,
+    field_validator,
+    model_validator,
+)
+
+from src.nk_model.enums import ConvolutionMethod, NeighborhoodMethod
+
+
+class NKParams(BaseModel):
+    n: PositiveInt
+    k: NonNegativeInt
+    power: PositiveFloat
+    max_val: PositiveFloat = Field(
+        default=1.0,
+        description="Maximum value for NK-landscape values. All values "
+        "will be scaled so that the maximum value equals this float. "
+        "Example: 1000.0 means the maximum value will be 1000.0.",
+    )
+    m: PositiveInt = Field(
+        description=(
+            "Number of convolutions to apply to the NK model landscape. "
+            "Default is m=n, which means no convolutions are applied."
+        ),
+    )
+    neighborhood: NeighborhoodMethod = Field(
+        default=NeighborhoodMethod.RANDOM,
+        description=(
+            "Method to determine the nearest neighbor for a given node in "
+            "the NK model."
+        ),
+    )
+    convolution: ConvolutionMethod = Field(
+        default=ConvolutionMethod.RANDOM,
+        description=(
+            "Method to determine the convolution method to apply to the NK "
+            "model landscape."
+        ),
+    )
+    payoff_type: Literal["int", "float"] = Field(
+        description=(
+            "Type for payoff values. If 'int', all payoffs will be integers. "
+            "If 'float', all payoffs will be floats."
+        ),
+    )
+
+    @field_validator("payoff_type", mode="before")
+    @classmethod
+    def validate_payoff_type(cls, v):
+        """Convert type objects to strings and validate."""
+        # Allow both type objects and strings
+        if v == int or v == "int":
+            return "int"
+        elif v == float or v == "float":
+            return "float"
+        else:
+            raise ValueError(
+                "payoff_type must be int, float, 'int', or 'float'"
+            )
+
+    @property
+    def payoff_type_class(self) -> type[int] | type[float]:
+        """Get the actual Python type for payoff_type."""
+        return int if self.payoff_type == "int" else float
+
+    @model_validator(mode="before")
+    @classmethod
+    def set_m_default(cls, data: dict) -> dict:
+        """Set m to n if not provided or None."""
+        if isinstance(data, dict):
+            data = data.copy()
+            # Set m to n if m is missing or None, and n is present
+            if ("m" not in data or data.get("m") is None) and "n" in data:
+                data["m"] = data["n"]
+        return data
+
+    @model_validator(mode="after")
+    def validate_and_set_defaults(self) -> "NKParams":
+        """Validate k and m."""
+        if self.k >= self.n:
+            raise ValueError("The value of k must be less than n.")
+
+        if self.m > self.n:
+            raise ValueError("The value of m must be less than or equal to n.")
+        return self
+
+
+class Item(BaseModel):
+    """
+    Represents an item in an NK model landscape.
+
+    Attributes:
+        coordinates (np.ndarray):
+            Binary coordinates of the item (e.g., np.array([0, 1, 1, 0])).
+        payoff (Union[int, float]):
+            Payoff value. Type depends on NKParams.payoff_type.
+    """
+
+    # Allows np.ndarray as attribute
+    model_config = {"arbitrary_types_allowed": True}
+
+    coordinates: np.ndarray
+    payoff: Union[int, float] = Field(..., ge=0)
+
+    @field_validator("coordinates")
+    def _validate_state(cls, coordinates: np.ndarray):
+        assert coordinates.ndim == 1, "Coordinates must be a 1D array."
+        assert all(
+            x in [0, 1] for x in coordinates
+        ), "Coordinates must be binary."
+        return coordinates
+
+    def __lt__(self, other: "Item") -> bool:
+        "Compare two items based on their payoffs"
+        return self.payoff < other.payoff
+
+    def model_dump(self) -> dict:
+        """
+        Return a dictionary of the Item,
+        because arrays cannot be serialized to JSON.
+        """
+        return {
+            "coordinates": self.coordinates.tolist(),
+            "payoff": self.payoff,
+        }
+
+
+class NKLandscapeCache(BaseModel):
+    """
+    Pydantic model for caching NKLandscape data.
+
+    This model simplifies the serialization and deserialization
+    of cached landscape data by using Pydantic's built-in
+    serialization capabilities.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    params: NKParams
+    items: list[Item]
+
+    def model_dump_json(self) -> str:
+        """Serialize to JSON string."""
+        return json.dumps(
+            {
+                "params": self.params.model_dump(mode="json"),
+                "items": [item.model_dump() for item in self.items],
+            }
+        )
+
+    @classmethod
+    def model_validate_json(cls, json_str: str) -> "NKLandscapeCache":
+        """Deserialize from JSON string."""
+        data = json.loads(json_str)
+        return cls(
+            params=NKParams(**data["params"]),
+            items=[
+                Item(
+                    coordinates=np.array(item["coordinates"]),
+                    payoff=item["payoff"],
+                )
+                for item in data["items"]
+            ],
+        )

nk_model/nk_landscape.py

@@ -0,0 +1,498 @@
+import itertools
+import random
+from math import ceil, floor
+from typing import Dict, List, Optional
+from uuid import uuid4
+
+import numpy as np
+
+from src.nk_model.assembler import assembler_v1_sym, ensemble_builder
+from src.nk_model.enums import ConvolutionMethod, NeighborhoodMethod
+from src.nk_model.landscape_cache import LandscapeCache
+from src.nk_model.models import Item, NKLandscapeCache, NKParams
+from src.utils.binary_conversion import binary_array_to_int
+
+
+class NKLandscape:
+    """
+    NK model fitness landscape class.
+
+    This class generates NK model landscapes, applies transformations
+    (normalization, power scaling, max_val scaling), and stores the
+    results as a list of Item objects.
+
+    The class always creates new landscapes in __init__ and provides
+    a class method from_cache() to load existing landscapes from cache.
+    """
+
+    # Class-level cache
+    _cache = LandscapeCache()
+
+    def __init__(self, params: NKParams):
+        """
+        Initialize NKLandscape with given parameters.
+
+        Always creates a new landscape and caches it with a generated UUID.
+
+        Args:
+            params: Parameters for the NK model containing:
+                - n: Number of components in the system
+                - k: Number of interactions per component
+                - m: Number of subcombinations to convolute
+                - power: Power scaling factor for payoffs
+                - max_val: Maximum value for scaled payoffs
+                - neighborhood: Method to determine neighbors
+                - convolution: Method for convolution
+
+        Attributes:
+            params: NKParams instance
+            items: List of Item objects with transformed payoffs
+            uuid: Unique identifier for this landscape
+            _payoff_lookup: Dictionary mapping coordinate tuples to payoffs
+        """
+        self.params = params
+        self.N = params.n
+        self.K = params.k
+        self.M = params.m
+
+        # Generate UUID for this landscape
+        self.uuid = str(uuid4())
+
+        # Create the landscape with transformed payoffs
+        self.items = self._create_landscape()
+
+        # Create lookup dictionary for O(1) payoff access
+        self._payoff_lookup = {
+            tuple(item.coordinates): item.payoff for item in self.items
+        }
+
+        # Cache the landscape
+        cache_data = NKLandscapeCache(params=self.params, items=self.items)
+        self._cache.save(self.uuid, cache_data)
+
+    @classmethod
+    def from_cache(cls, uuid: str) -> "NKLandscape":
+        """
+        Load a landscape from cache by UUID.
+
+        Args:
+            uuid: The UUID of the landscape to retrieve
+
+        Returns:
+            NKLandscape instance with cached data
+
+        Raises:
+            ValueError: If the landscape with the given UUID is not found
+        """
+        cached_data = cls._cache.get(uuid)
+        if cached_data is None:
+            raise ValueError(f"Landscape with uuid {uuid} not found in cache.")
+
+        # Create instance without calling __init__
+        instance = cls.__new__(cls)
+        instance.params = cached_data.params
+        instance.N = cached_data.params.n
+        instance.K = cached_data.params.k
+        instance.M = cached_data.params.m
+        instance.uuid = uuid
+        instance.items = cached_data.items
+
+        # Create lookup dictionary
+        instance._payoff_lookup = {
+            tuple(item.coordinates): item.payoff for item in instance.items
+        }
+
+        return instance
+
+    @classmethod
+    def from_dict(
+        cls,
+        data: Dict[str, int | float],
+        k: Optional[int] = None,
+        m: Optional[int] = None,
+        power: float = 1.0,
+        neighborhood: NeighborhoodMethod = (
+            NeighborhoodMethod.RANDOM
+        ),
+        convolution: ConvolutionMethod = ConvolutionMethod.RANDOM,
+    ) -> "NKLandscape":
+        """
+        Load a landscape from dictionary data.
+
+        Args:
+            data: Dictionary mapping binary string locations
+                (e.g., '01101010') to payoff values (int or float).
+            n: Number of components. If None, derived from binary
+                string length.
+            k: Number of interactions per component. If None,
+                defaults to n-1 (ensuring k < n).
+            m: Number of subcombinations to convolute. If None,
+                defaults to n.
+            power: Power scaling factor for payoffs. Defaults to 1.0.
+            neighborhood: Method to determine neighbors.
+                Defaults to RANDOM.
+            convolution: Method for convolution. Defaults to RANDOM.
+            payoff_type: Type for payoff values ('int' or 'float').
+                If None, derived from data (int if all values are int,
+                else float).
+
+        Returns:
+            NKLandscape instance with data from dictionary.
+
+        Raises:
+            ValueError: If binary strings have inconsistent lengths
+                or if derived parameters are invalid.
+        """
+        if not data:
+            raise ValueError("Data dictionary cannot be empty.")
+
+        # Extract binary strings and values
+        binary_strings = list(data.keys())
+        values = list(data.values())
+        
+        # Derive parameters from data
+        payoff_type = "float" if any(isinstance(v, float) for v in values) else "int"
+        n = len(binary_strings[0])
+        max_val = max(values)
+
+        # Set default parameters
+        k = k or max(0, n - 1)
+        m = m or n
+
+        for bs in binary_strings:
+            # Validate binary strings contain only '0' and '1'
+            if not all(bit in "01" for bit in bs):
+                raise ValueError(
+                    f"Binary string '{bs}' contains invalid characters. "
+                    + "Only '0' and '1' are allowed."
+                )
+            if len(bs) != n:
+                raise ValueError(
+                    f"Binary string '{bs}' has length {len(bs)}, "
+                    + f"but expected length {n}."
+                )
+
+        # Create NKParams
+        params = NKParams(
+            n=n,
+            k=k,
+            m=m,
+            power=power,
+            max_val=max_val,
+            neighborhood=neighborhood,
+            convolution=convolution,
+            payoff_type=payoff_type,
+        )
+
+        # Convert dictionary data to Item objects
+        items = []
+        for binary_str, payoff in data.items():
+            coordinates = np.array([int(bit) for bit in binary_str])
+            payoff_type_class = params.payoff_type_class
+            items.append(
+                Item(
+                    coordinates=coordinates,
+                    payoff=payoff_type_class(payoff),
+                )
+            )
+
+        # Create instance without calling __init__
+        instance = cls.__new__(cls)
+        instance.params = params
+        instance.N = params.n
+        instance.K = params.k
+        instance.M = params.m
+        instance.uuid = str(uuid4())
+        instance.items = items
+
+        # Create lookup dictionary
+        instance._payoff_lookup = {
+            tuple(item.coordinates): item.payoff for item in instance.items
+        }
+
+        return instance
+
+    @classmethod
+    def clear_cache(cls) -> None:
+        """Clear the entire landscape cache."""
+        cls._cache.clear()
+
+    def get_payoff(self, state: np.ndarray) -> int | float:
+        """
+        Get the payoff for a given state.
+
+        Args:
+            state: Binary state vector of length N
+
+        Returns:
+            Payoff value for the given state
+
+        Raises:
+            AssertionError: If state length doesn't match N
+        """
+        assert (
+            len(state) == self.params.n
+        ), f"""State length ({len(state)})
+        must be equal to N ({self.params.n})."""
+        return self._payoff_lookup[tuple(state)]
+
+    def get_ordered_payoffs(self) -> List[int | float]:
+        """
+        Get the payoffs for all states in the landscape,
+        ordered by ascending coordinates.
+        """
+        return [
+            payoff
+            for _, payoff in sorted(
+                self._payoff_lookup.items(),
+                key=lambda x: binary_array_to_int(np.array(x[0])),
+            )
+        ]
+
+    def _create_landscape(self) -> list[Item]:
+        """
+        Create the NK landscape with transformed payoffs.
+
+        This method applies the following transformations:
+        1. Generate raw NK model payoffs
+        2. Normalize (min=0, max=1)
+        3. Apply power scaling
+        4. Scale to max_val if specified
+
+        Returns:
+            List of Item objects with transformed payoffs
+        """
+        # Generate raw NK model items
+        raw_items = self._create_nk_model()
+
+        # Extract payoffs
+        payoffs = np.array([item.payoff for item in raw_items])
+
+        # Normalize payoffs
+        min_payoff = np.min(payoffs)
+        max_payoff = np.max(payoffs)
+        if max_payoff > min_payoff:
+            payoffs = (payoffs - min_payoff) / (max_payoff - min_payoff)
+        else:
+            payoffs = np.zeros_like(payoffs)
+
+        # Apply power scaling
+        scaled_payoffs = payoffs**self.params.power
+
+        # Scale to max_val if specified
+        if self.params.max_val is not None:
+            max_value = np.max(scaled_payoffs)
+            if max_value > 0:
+                scaled_payoffs = scaled_payoffs * (
+                    self.params.max_val / max_value
+                )
+
+        # Convert payoffs to the specified type
+        if self.params.payoff_type == "int":
+            scaled_payoffs = np.round(scaled_payoffs).astype(int)
+        else:
+            scaled_payoffs = scaled_payoffs.astype(float)
+
+        # Create items with transformed payoffs
+        payoff_type_class = self.params.payoff_type_class
+        return [
+            Item(
+                coordinates=item.coordinates, payoff=payoff_type_class(payoff)
+            )
+            for item, payoff in zip(raw_items, scaled_payoffs)
+        ]
+
+    def _create_nk_model(self) -> list[Item]:
+        """
+        Create the NK model by generating interaction matrices and
+        computing items list for all states.
+
+        Returns:
+            List of Item objects containing coordinates and raw payoff
+            values for all possible states in the NK model.
+        """
+        # Generate node interaction matrix
+        interaction_method_map = {
+            NeighborhoodMethod.RANDOM: (self._get_node_interaction_matrix),
+            NeighborhoodMethod.NEAREST: (
+                self._get_nearest_neighbor_interaction_matrix
+            ),
+            NeighborhoodMethod.RING: (self._get_ring_node_interaction_matrix),
+        }
+        node_interaction_matrix = interaction_method_map[
+            self.params.neighborhood
+        ]()
+
+        # Generate subcombination payoff matrix
+        base = np.random.uniform(
+            low=0,
+            high=1,
+            size=((self.M,) + (2,) * (self.K + 1)),
+        )
+        convolution_method_map = {
+            ConvolutionMethod.RANDOM: self._get_random_convolution_matrix,
+            ConvolutionMethod.SYMMETRIC: (
+                self._get_symmetric_convolution_matrix
+            ),
+        }
+        subcombination_payoff_matrix = convolution_method_map[
+            self.params.convolution
+        ](base=base)
+
+        # Compute items list
+        coordinates = itertools.product([0, 1], repeat=self.N)
+
+        def get_payoff(
+            coordinate: tuple[int, ...],
+            interaction_matrix: np.ndarray,
+            subcombination_matrix: np.ndarray,
+        ) -> float:
+            # Sum up the payoffs for all digits of the coordinates
+            return sum(
+                # For each dimension, get the interacting node IDs,
+                # extract their coordinates values, and look up the
+                # payoff from the subcombination matrix
+                subcombination_matrix[dimension][
+                    tuple(
+                        [
+                            coordinate[interact_node]
+                            for interact_node in interaction_matrix[dimension]
+                        ]
+                    )
+                ]
+                for dimension in range(len(coordinate))
+            )
+
+        return [
+            Item(
+                coordinates=np.array(coordinate),
+                payoff=get_payoff(
+                    coordinate=coordinate,
+                    interaction_matrix=node_interaction_matrix,
+                    subcombination_matrix=subcombination_payoff_matrix,
+                ),
+            )
+            for coordinate in coordinates
+        ]
+
+    def _get_random_convolution_matrix(
+        self,
+        base: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Apply random convolution method to create final subcombination
+        matrix.
+
+        Args:
+            base: Base matrix of shape (m, 2, 2, ..., 2).
+
+        Returns:
+            np.ndarray: A multidimensional array of shape
+            (n, 2, 2, ..., 2) containing random payoff values.
+        """
+        matrix = np.empty(((self.N,) + (2,) * (self.K + 1)), dtype=np.float64)
+        for i in range(self.N):
+            matrix[i] = base[np.random.randint(0, base.shape[0])]
+        return matrix
+
+    def _get_symmetric_convolution_matrix(
+        self,
+        base: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Apply symmetric convolution method to create final subcombination
+        matrix.
+
+        Args:
+            base: Base matrix of shape (m, 2, 2, ..., 2).
+
+        Returns:
+            np.ndarray: A multidimensional array of shape
+            (n, 2, 2, ..., 2) containing random payoff values.
+        """
+        assert self.M == 2, "Symmetric convolution requires m=2"
+        assert self.N % 2 == 0, "Symmetric convolution requires n to be even"
+        ensemble = ensemble_builder(
+            assembler=assembler_v1_sym,
+            n=16,
+            max_len=self.N,
+            vectors=[[0], [1]],
+        )
+        filtered_ensemble = [
+            vector for vector in ensemble if len(vector) == self.N
+        ]
+        convolution_vector = random.choice(filtered_ensemble)
+        matrix = np.empty(((self.N,) + (2,) * (self.K + 1)), dtype=np.float64)
+        for i in range(self.N):
+            matrix[i] = base[convolution_vector[i]]
+        return matrix
+
+    def _get_node_interaction_matrix(self) -> np.ndarray:
+        """
+        Generate an interaction matrix for the NK model.
+        This matrix defines which nodes in our landscape interact
+        with each other (are considered neighbors).
+        The n rows of this matrix correspond to the n nodes in
+        our landscape.
+        The k+1 columns of this matrix correspond to the k+1
+        possible neighbors for each node.
+
+        Returns:
+            np.ndarray: An N x (K + 1) matrix where each row has
+            K + 1 indices that are the neighbors of the node.
+
+        Note:
+            This function ensures that each node interacts with K
+            other distinct nodes, creating a sparse interaction
+            matrix for the NK model landscape.
+        """
+        # Initialize interaction matrix with ones on the diagonal
+        interaction_matrix = np.zeros((self.N, self.K + 1), dtype=np.int64)
+
+        # Choose K non-zero entries for each row randomly
+        for row_index in range(self.N):
+            # Get indices of all columns except the diagonal element
+            indices = list(range(self.N))
+            indices.remove(row_index)
+
+            # Choose K non-zero entries for the row randomly
+            chosen_indices = np.random.choice(
+                indices, size=self.K, replace=False
+            )
+            chosen_indices = np.append(chosen_indices, row_index)
+            np.random.shuffle(chosen_indices)
+
+            interaction_matrix[row_index] = chosen_indices
+
+        return interaction_matrix
+
+    def _get_ring_node_interaction_matrix(self) -> np.ndarray:
+        """
+        Generate a ring interaction matrix for the NK model.
+        """
+        interaction_matrix = np.zeros((self.N, self.K + 1), dtype=np.int64)
+
+        for row_index in range(self.N):
+            indices = np.arange(2 * self.N)
+            indices = np.array(
+                [i for i in indices if i - row_index in range(0, self.K + 1)]
+            )
+            indices = indices % self.N
+            interaction_matrix[row_index] = indices
+        return interaction_matrix
+
+    def _get_nearest_neighbor_interaction_matrix(self) -> np.ndarray:
+        """
+        Generate a nearest neighbor interaction matrix for the NK model.
+        """
+        interaction_matrix = np.zeros((self.N, self.K + 1), dtype=np.int64)
+        for row_index in range(self.N):
+            indices = np.array(
+                range(
+                    row_index - floor((self.K + 1) / 2),
+                    row_index + ceil((self.K + 1) / 2),
+                )
+            )
+            indices = indices % self.N
+            interaction_matrix[row_index] = indices
+        return interaction_matrix