File size: 7,080 Bytes
d5a9c75 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """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)
|