experiment/parameter_space.py

"""Parameter space definitions for experiment design."""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Union

import torch
from torch import Tensor
import numpy as np


@dataclass
class ContinuousParameter:
    """A continuous real-valued parameter."""

    name: str
    lower: float
    upper: float
    log_scale: bool = False  # Use log-scale for parameters spanning orders of magnitude
    units: str = ""

    def sample(self, n: int = 1, dtype=torch.float64) -> Tensor:
        if self.log_scale:
            log_samples = torch.rand(n, dtype=dtype) * (
                np.log(self.upper) - np.log(self.lower)
            ) + np.log(self.lower)
            return log_samples.exp()
        return torch.rand(n, dtype=dtype) * (self.upper - self.lower) + self.lower


@dataclass
class IntegerParameter:
    """An integer-valued parameter."""

    name: str
    lower: int
    upper: int
    units: str = ""

    def sample(self, n: int = 1, dtype=torch.float64) -> Tensor:
        return torch.randint(self.lower, self.upper + 1, (n,)).to(dtype=dtype)


@dataclass
class CategoricalParameter:
    """A categorical parameter with discrete choices."""

    name: str
    categories: List[str]
    units: str = ""

    def sample(self, n: int = 1, dtype=torch.float64) -> Tensor:
        indices = torch.randint(0, len(self.categories), (n,))
        return indices.to(dtype=dtype)

    def encode(self, category: str) -> int:
        return self.categories.index(category)

    def decode(self, index: int) -> str:
        return self.categories[index]


class ParameterSpace:
    """Defines the experimental parameter space for optimization.

    Supports continuous, integer, and categorical parameters with
    optional linear constraints between parameters.
    """

    def __init__(self):
        self._parameters: Dict[str, Union[ContinuousParameter, IntegerParameter, CategoricalParameter]] = {}
        self._order: List[str] = []
        self._constraints: List[Dict] = []

    def add_continuous(
        self,
        name: str,
        lower: float,
        upper: float,
        log_scale: bool = False,
        units: str = "",
    ) -> "ParameterSpace":
        """Add a continuous parameter."""
        self._parameters[name] = ContinuousParameter(name, lower, upper, log_scale, units)
        self._order.append(name)
        return self

    def add_integer(
        self, name: str, lower: int, upper: int, units: str = ""
    ) -> "ParameterSpace":
        """Add an integer parameter."""
        self._parameters[name] = IntegerParameter(name, lower, upper, units)
        self._order.append(name)
        return self

    def add_categorical(
        self, name: str, categories: List[str], units: str = ""
    ) -> "ParameterSpace":
        """Add a categorical parameter."""
        self._parameters[name] = CategoricalParameter(name, categories, units)
        self._order.append(name)
        return self

    def add_sum_constraint(
        self, parameter_names: List[str], target_sum: float = 1.0
    ) -> "ParameterSpace":
        """Add a constraint that parameters must sum to a target value.

        Useful for mixture/composition experiments.
        """
        self._constraints.append({
            "type": "sum",
            "parameters": parameter_names,
            "target": target_sum,
        })
        return self

    def add_linear_constraint(
        self,
        parameter_names: List[str],
        coefficients: List[float],
        bound: float,
        constraint_type: str = "<=",
    ) -> "ParameterSpace":
        """Add a linear constraint: sum(coeff_i * param_i) <= bound."""
        self._constraints.append({
            "type": "linear",
            "parameters": parameter_names,
            "coefficients": coefficients,
            "bound": bound,
            "constraint_type": constraint_type,
        })
        return self

    @property
    def dimension(self) -> int:
        return len(self._parameters)

    @property
    def parameter_names(self) -> List[str]:
        return self._order

    @property
    def bounds(self) -> Tensor:
        """Get bounds as a (2, d) tensor for BoTorch."""
        lowers, uppers = [], []
        for name in self._order:
            p = self._parameters[name]
            if isinstance(p, ContinuousParameter):
                lowers.append(p.lower)
                uppers.append(p.upper)
            elif isinstance(p, IntegerParameter):
                lowers.append(float(p.lower))
                uppers.append(float(p.upper))
            elif isinstance(p, CategoricalParameter):
                lowers.append(0.0)
                uppers.append(float(len(p.categories) - 1))
        return torch.tensor([lowers, uppers], dtype=torch.float64)

    def sample_random(self, n: int = 1, dtype=torch.float64) -> Tensor:
        """Generate random samples from the parameter space."""
        samples = []
        for name in self._order:
            samples.append(self._parameters[name].sample(n, dtype))
        return torch.stack(samples, dim=-1)

    def sample_latin_hypercube(self, n: int, dtype=torch.float64) -> Tensor:
        """Generate Latin Hypercube samples for space-filling initial design."""
        d = self.dimension
        # Create LHS grid
        intervals = torch.linspace(0, 1, n + 1)
        samples = torch.zeros(n, d, dtype=dtype)

        for j in range(d):
            # Random permutation within each dimension
            perm = torch.randperm(n)
            for i in range(n):
                low = intervals[perm[i]]
                high = intervals[perm[i] + 1]
                samples[i, j] = low + (high - low) * torch.rand(1, dtype=dtype)

        # Scale to parameter bounds
        bounds = self.bounds
        samples = samples * (bounds[1] - bounds[0]) + bounds[0]
        return samples

    def to_dict(self, X: Tensor) -> List[Dict]:
        """Convert a tensor of parameter values to list of dicts."""
        results = []
        for i in range(len(X)):
            d = {}
            for j, name in enumerate(self._order):
                p = self._parameters[name]
                if isinstance(p, CategoricalParameter):
                    d[name] = p.decode(int(X[i, j].item()))
                else:
                    d[name] = X[i, j].item()
            results.append(d)
        return results

    def from_dict(self, params: Dict[str, float], dtype=torch.float64) -> Tensor:
        """Convert a parameter dict to a tensor row."""
        values = []
        for name in self._order:
            p = self._parameters[name]
            if isinstance(p, CategoricalParameter):
                values.append(float(p.encode(params[name])))
            else:
                values.append(float(params[name]))
        return torch.tensor(values, dtype=dtype)