Upload experiment/parameter_space.py with huggingface_hub

experiment/parameter_space.py

@@ -0,0 +1,209 @@
+"""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)