Coverage for src/amisc/training.py: 89%
264 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-11-05 19:26 +0000
1"""Classes for storing and managing training data for surrogate models. The `TrainingData` interface also
2specifies how new training data should be sampled over the input space (i.e. experimental design).
4Includes:
6- `TrainingData` — an interface for storing surrogate training data.
7- `SparseGrid` — a class for storing training data in a sparse grid format.
8"""
9from __future__ import annotations
11import copy
12import itertools
13from abc import ABC, abstractmethod
14from dataclasses import dataclass, field
15from typing import Any, ClassVar
17import numpy as np
18from numpy.typing import ArrayLike
19from scipy.optimize import direct
21from amisc.serialize import PickleSerializable, Serializable
22from amisc.typing import LATENT_STR_ID, Dataset, MultiIndex
23from amisc.utils import _RidgeRegression
25__all__ = ['TrainingData', 'SparseGrid']
28class TrainingData(Serializable, ABC):
29 """Interface for storing and collecting surrogate training data. `TrainingData` objects should:
31 - `get` - retrieve the training data
32 - `set` - store the training data
33 - `refine` - generate new design points for the parent `Component` model
34 - `clear` - clear all training data
35 - `set_errors` - store error information (if desired)
36 - `impute_missing_data` - fill in missing values in the training data (if desired)
37 """
39 @abstractmethod
40 def get(self, alpha: MultiIndex, beta: MultiIndex, y_vars: list[str] = None,
41 skip_nan: bool = False) -> tuple[Dataset, Dataset]:
42 """Return the training data for a given multi-index pair.
44 :param alpha: the model fidelity indices
45 :param beta: the surrogate fidelity indices
46 :param y_vars: the keys of the outputs to return (if `None`, return all outputs)
47 :param skip_nan: skip any data points with remaining `nan` values if `skip_nan=True`
48 :returns: `dicts` of model inputs `x_train` and outputs `y_train`
49 """
50 raise NotImplementedError
52 @abstractmethod
53 def set(self, alpha: MultiIndex, beta: MultiIndex, coords: list[Any], yi_dict: Dataset):
54 """Store training data for a given multi-index pair.
56 :param alpha: the model fidelity indices
57 :param beta: the surrogate fidelity indices
58 :param coords: locations for storing the `yi` values in the underlying data structure
59 :param yi_dict: a `dict` of model output `yi` values, each entry should be the same length as `coords`
60 """
61 raise NotImplementedError
63 @abstractmethod
64 def set_errors(self, alpha: MultiIndex, beta: MultiIndex, coords: list[Any], errors: list[dict]):
65 """Store error information for a given multi-index pair (just pass if you don't care).
67 :param alpha: the model fidelity indices
68 :param beta: the surrogate fidelity indices
69 :param coords: locations for storing the error information in the underlying data structure
70 :param errors: a list of error dictionaries, should be the same length as `coords`
71 """
72 raise NotImplementedError
74 @abstractmethod
75 def impute_missing_data(self, alpha: MultiIndex, beta: MultiIndex):
76 """Impute missing values in the training data for a given multi-index pair (just pass if you don't care).
78 :param alpha: the model fidelity indices
79 :param beta: the surrogate fidelity indices
80 """
81 raise NotImplementedError
83 @abstractmethod
84 def refine(self, alpha: MultiIndex, beta: MultiIndex, input_domains: dict[str, tuple],
85 weight_fcns: dict[str, callable] = None) -> tuple[list[Any], Dataset]:
86 """Return new design/training points for a given multi-index pair and their coordinates/locations in the
87 `TrainingData` storage structure.
89 !!! Example
90 ```python
91 domains = {'x1': (0, 1), 'x2': (0, 1)}
92 alpha, beta = (0, 1), (1, 1)
93 coords, x_train = training_data.refine(alpha, beta, domains)
94 y_train = my_model(x_train)
95 training_data.set(alpha, beta, coords, y_train)
96 ```
98 The returned data coordinates `coords` should be any object that can be used to locate the corresponding
99 `x_train` training points in the `TrainingData` storage structure. These `coords` will be passed back to the
100 `set` function to store the training data at a later time (i.e. after model evaluation).
102 :param alpha: the model fidelity indices
103 :param beta: the surrogate fidelity indices
104 :param input_domains: a `dict` specifying domain bounds for each input variable
105 :param weight_fcns: a `dict` of weighting functions for each input variable
106 :returns: a list of new data coordinates `coords` and the corresponding training points `x_train`
107 """
108 raise NotImplementedError
110 @abstractmethod
111 def clear(self):
112 """Clear all training data."""
113 raise NotImplementedError
115 @classmethod
116 def from_dict(cls, config: dict) -> TrainingData:
117 """Create a `TrainingData` object from a `dict` configuration. Currently, only `method='sparse-grid'` is
118 supported for the `SparseGrid` class.
119 """
120 method = config.pop('method', 'sparse-grid').lower()
121 match method:
122 case 'sparse-grid':
123 return SparseGrid(**config)
124 case other:
125 raise NotImplementedError(f"Unknown training data method: {other}")
128@dataclass
129class SparseGrid(TrainingData, PickleSerializable):
130 """A class for storing training data in a sparse grid format. The `SparseGrid` class stores training points
131 by their coordinate location in a larger tensor-product grid, and obtains new training data by refining
132 a single 1d grid at a time.
134 !!! Note "MISC and sparse grids"
135 MISC itself can be thought of as an extension to the well-known sparse grid technique, so this class
136 readily integrates with the MISC implementation in `Component`. Sparse grids limit the curse
137 of dimensionality up to about `dim = 10-15` for the input space (which would otherwise be infeasible with a
138 normal full tensor-product grid of the same size).
140 !!! Info "About points in a sparse grid"
141 A sparse grid approximates a full tensor-product grid $(N_1, N_2, ..., N_d)$, where $N_i$ is the number of grid
142 points along dimension $i$, for a $d$-dimensional space. Each point is uniquely identified in the sparse grid
143 by a list of indices $(j_1, j_2, ..., j_d)$, where $j_i = 0 ... N_i$. We refer to this unique identifier as a
144 "grid coordinate". In the `SparseGrid` data structure, these coordinates are used along with the `alpha`
145 fidelity index to uniquely locate the training data for a given multi-index pair.
147 :ivar collocation_rule: the collocation rule to use for generating new grid points (only 'leja' is supported)
148 :ivar knots_per_level: the number of grid knots/points per level in the `beta` fidelity multi-index
149 :ivar expand_latent_method: method for expanding latent grids, either 'round-robin' or 'tensor-product'
150 :ivar opt_args: extra arguments for the global 1d `direct` optimizer
151 :ivar betas: a set of all `beta` multi-indices that have been seen so far
152 :ivar x_grids: a `dict` of grid points for each 1d input dimension
153 :ivar yi_map: a `dict` of model outputs for each grid coordinate
154 :ivar yi_nan_map: a `dict` of imputed model outputs for each grid coordinate where the model failed (or gave nan)
155 :ivar error_map: a `dict` of error information for each grid coordinate where the model failed
156 :ivar latent_size: the number of latent coefficients for each variable (0 if scalar)
157 """
158 MAX_IMPUTE_SIZE: ClassVar[int] = 10 # don't try to impute large arrays
160 collocation_rule: str = 'leja'
161 knots_per_level: int = 2
162 expand_latent_method: str = 'round-robin' # or 'tensor-product', for converting beta to latent grid sizes
163 opt_args: dict = field(default_factory=lambda: {'locally_biased': False, 'maxfun': 300}) # for leja optimizer
165 betas: set[MultiIndex] = field(default_factory=set)
166 x_grids: dict[str, ArrayLike] = field(default_factory=dict)
167 yi_map: dict[MultiIndex, dict[tuple[int, ...], dict[str, ArrayLike]]] = field(default_factory=dict)
168 yi_nan_map: dict[MultiIndex, dict[tuple[int, ...], dict[str, ArrayLike]]] = field(default_factory=dict)
169 error_map: dict[MultiIndex, dict[tuple[int, ...], dict[str, Any]]] = field(default_factory=dict)
170 latent_size: dict[str, int] = field(default_factory=dict) # keep track of latent grid sizes for each variable
172 def clear(self):
173 """Clear all training data."""
174 self.betas.clear()
175 self.x_grids.clear()
176 self.yi_map.clear()
177 self.yi_nan_map.clear()
178 self.error_map.clear()
179 self.latent_size.clear()
181 def get_by_coord(self, alpha: MultiIndex, coords: list, y_vars: list = None, skip_nan: bool = False):
182 """Get training data from the sparse grid for a given `alpha` and list of grid coordinates. Try to replace
183 `nan` values with imputed values. Skip any data points with remaining `nan` values if `skip_nan=True`.
185 :param alpha: the model fidelity indices
186 :param coords: a list of grid coordinates to locate the `yi` values in the sparse grid data structure
187 :param y_vars: the keys of the outputs to return (if `None`, return all outputs)
188 :param skip_nan: skip any data points with remaining `nan` values if `skip_nan=True`
189 :returns: `dicts` of model inputs `xi_dict` and outputs `yi_dict`
190 """
191 xi_dict = {}
192 yi_dict = {}
193 y_squeeze = {} # Keep track of whether to squeeze extra output dimensions (i.e. non field quantities)
194 for coord in coords:
195 try:
196 skip_coord = False
197 yi_curr = copy.deepcopy(self.yi_map[alpha][coord])
198 output_vars = self._numeric_outputs(yi_curr)
199 for var in output_vars:
200 if np.any(np.isnan(yi_curr[var])):
201 yi_curr[var] = self.yi_nan_map[alpha].get(coord, yi_curr)[var]
202 if skip_nan and np.any(np.isnan(yi_curr[var])):
203 skip_coord = True
204 break
206 if not skip_coord:
207 self._append_grid_points(coord, xi_dict)
208 yi_vars = y_vars if y_vars is not None else yi_curr.keys()
209 for var in yi_vars:
210 yi = np.atleast_1d(yi_curr[var])
211 y_squeeze[var] = self._is_singleton(yi) # squeeze for scalar quantities
212 yi = np.expand_dims(yi, axis=0)
213 yi_dict[var] = yi if yi_dict.get(var) is None else np.concatenate((yi_dict[var], yi), axis=0)
214 except KeyError as e:
215 raise KeyError(f"Can't access sparse grid data for alpha={alpha}, coord={coord}. "
216 f"Make sure the data has been set first.") from e
218 # Squeeze out extra dimension for scalar quantities
219 for var in yi_dict.keys():
220 if y_squeeze[var]:
221 yi_dict[var] = np.squeeze(yi_dict[var], axis=-1)
223 return xi_dict, yi_dict # Both with elements of shape (N, ...) for N grid points
225 def get(self, alpha: MultiIndex, beta: MultiIndex, y_vars: list[str] = None, skip_nan: bool = False):
226 """Get the training data from the sparse grid for a given `alpha` and `beta` pair."""
227 return self.get_by_coord(alpha, self._expand_grid_coords(beta), y_vars=y_vars, skip_nan=skip_nan)
229 def set_errors(self, alpha: MultiIndex, beta: MultiIndex, coords: list, errors: list[dict]):
230 """Store error information in the sparse-grid for a given multi-index pair."""
231 for coord, error in zip(coords, errors):
232 self.error_map[alpha][coord] = copy.deepcopy(error)
234 def set(self, alpha: MultiIndex, beta: MultiIndex, coords: list, yi_dict: dict[str, ArrayLike]):
235 """Store model output `yi_dict` values.
237 :param alpha: the model fidelity indices
238 :param beta: the surrogate fidelity indices
239 :param coords: a list of grid coordinates to locate the `yi` values in the sparse grid data structure
240 :param yi_dict: a `dict` of model output `yi` values
241 """
242 for i, coord in enumerate(coords): # First dim of yi is loop dim aligning with coords
243 new_yi = {}
244 for var, yi in yi_dict.items():
245 yi = np.atleast_1d(yi[i])
246 new_yi[var] = (float(yi[0]) if self._is_numeric(yi) else yi[0]) if self._is_singleton(yi) else yi.tolist() # noqa: E501
247 self.yi_map[alpha][coord] = copy.deepcopy(new_yi)
249 def impute_missing_data(self, alpha: MultiIndex, beta: MultiIndex):
250 """Impute missing values in the sparse grid for a given multi-index pair by linear regression imputation."""
251 imputer, xi_all, yi_all = None, None, None
252 for coord, yi_dict in self.yi_map[alpha].items():
253 # only impute (small-length) numeric quantities
254 output_vars = [var for var in self._numeric_outputs(yi_dict)
255 if len(np.ravel(yi_dict[var])) <= self.MAX_IMPUTE_SIZE]
257 if any([np.any(np.isnan(yi_dict[var])) for var in output_vars]):
258 if imputer is None:
259 # Grab all 'good' interpolation points and train a simple linear regression fit
260 xi_all, yi_all = self.get(alpha, beta, y_vars=output_vars, skip_nan=True)
261 if len(xi_all) == 0:
262 continue # possible if no good data has been set yet
264 N = next(iter(xi_all.values())).shape[0] # number of grid points
265 xi_mat = np.concatenate([xi_all[var][:, np.newaxis] if len(xi_all[var].shape) == 1 else
266 xi_all[var] for var in xi_all.keys()], axis=-1)
267 yi_mat = np.concatenate([yi_all[var][:, np.newaxis] if len(yi_all[var].shape) == 1 else
268 yi_all[var].reshape((N, -1)) for var in output_vars], axis=-1)
270 imputer = _RidgeRegression(alpha=1.0)
271 imputer.fit(xi_mat, yi_mat)
273 # Run the imputer for this coordinate
274 x_interp = self._append_grid_points(coord)
275 x_interp = np.concatenate([x_interp[var][:, np.newaxis] if len(x_interp[var].shape) == 1 else
276 x_interp[var] for var in x_interp.keys()], axis=-1)
277 y_interp = imputer.predict(x_interp)
279 # Unpack the imputed value
280 y_impute = {}
281 start_idx = 0
282 for var in output_vars:
283 var_shape = yi_all[var].shape[1:] or (1,)
284 end_idx = start_idx + int(np.prod(var_shape))
285 yi = np.atleast_1d(y_interp[0, start_idx:end_idx]).reshape(var_shape)
286 nan_idx = np.isnan(np.atleast_1d(yi_dict[var]))
287 yi[~nan_idx] = np.atleast_1d(yi_dict[var])[~nan_idx] # Only keep imputed values where yi is nan
288 y_impute[var] = float(yi[0]) if self._is_singleton(yi) else yi.tolist()
289 start_idx = end_idx
291 self.yi_nan_map[alpha][coord] = copy.deepcopy(y_impute)
293 def refine(self, alpha: MultiIndex, beta: MultiIndex, input_domains: dict, weight_fcns: dict = None):
294 """Refine the sparse grid for a given `alpha` and `beta` pair and given collocation rules. Return any new
295 grid points that do not have model evaluations saved yet.
297 !!! Note
298 The `beta` multi-index is used to determine the number of collocation points in each input dimension. The
299 length of `beta` should therefore match the number of variables in `x_vars`.
300 """
301 weight_fcns = weight_fcns or {}
303 # Initialize a sparse grid for beta=(0, 0, ..., 0)
304 if np.sum(beta) == 0:
305 if len(self.x_grids) == 0:
306 num_latent = {}
307 for var in input_domains:
308 if LATENT_STR_ID in var:
309 base_id = var.split(LATENT_STR_ID)[0]
310 num_latent[base_id] = 1 if base_id not in num_latent else num_latent[base_id] + 1
311 else:
312 num_latent[var] = 0
313 self.latent_size = num_latent
315 new_pt = {}
316 domains = iter(input_domains.items())
317 for grid_size in self.beta_to_knots(beta):
318 if isinstance(grid_size, int): # scalars
319 var, domain = next(domains)
320 new_pt[var] = self.collocation_1d(grid_size, domain, method=self.collocation_rule,
321 wt_fcn=weight_fcns.get(var, None),
322 opt_args=self.opt_args).tolist()
323 else: # latent coeffs
324 for s in grid_size:
325 var, domain = next(domains)
326 new_pt[var] = self.collocation_1d(s, domain, method=self.collocation_rule,
327 wt_fcn=weight_fcns.get(var, None),
328 opt_args=self.opt_args).tolist()
329 self.x_grids = new_pt
330 self.betas.add(beta)
331 self.yi_map.setdefault(alpha, dict())
332 self.yi_nan_map.setdefault(alpha, dict())
333 self.error_map.setdefault(alpha, dict())
334 new_coords = list(self._expand_grid_coords(beta))
335 return new_coords, self._append_grid_points(new_coords[0])
337 # Otherwise, refine the sparse grid
338 for beta_old in self.betas:
339 # Get the first lower neighbor in the sparse grid and refine the 1d grid if necessary
340 if self.is_one_level_refinement(beta_old, beta):
341 new_grid_size = self.beta_to_knots(beta)
342 inputs = zip(self.x_grids.keys(), self.x_grids.values(), input_domains.values())
344 for new_size in new_grid_size:
345 if isinstance(new_size, int): # scalar grid
346 var, grid, domain = next(inputs)
347 if len(grid) < new_size:
348 num_new_pts = new_size - len(grid)
349 self.x_grids[var] = self.collocation_1d(num_new_pts, domain, grid, opt_args=self.opt_args,
350 wt_fcn=weight_fcns.get(var, None),
351 method=self.collocation_rule).tolist()
352 else: # latent grid
353 for s_new in new_size:
354 var, grid, domain = next(inputs)
355 if len(grid) < s_new:
356 num_new_pts = s_new - len(grid)
357 self.x_grids[var] = self.collocation_1d(num_new_pts, domain, grid,
358 opt_args=self.opt_args,
359 wt_fcn=weight_fcns.get(var, None),
360 method=self.collocation_rule).tolist()
361 break
363 new_coords = []
364 new_pts = {}
365 for coord in self._expand_grid_coords(beta):
366 if coord not in self.yi_map[alpha]:
367 # If we have not computed this grid coordinate yet
368 self._append_grid_points(coord, new_pts)
369 new_coords.append(coord)
371 self.betas.add(beta)
372 return new_coords, new_pts
374 def _append_grid_points(self, coord: tuple, pts: dict = None):
375 """Extract the `x` grid point located at `coord` from `x_grids` and append to the `pts` dictionary."""
376 if pts is None:
377 pts = {}
378 grids = iter(self.x_grids.items())
379 for idx in coord:
380 if isinstance(idx, int): # scalar grid point
381 var, grid = next(grids)
382 new_pt = np.atleast_1d(grid[idx])
383 pts[var] = new_pt if pts.get(var) is None else np.concatenate((pts[var], new_pt), axis=0)
384 else: # latent coefficients
385 for i in idx:
386 var, grid = next(grids)
387 new_pt = np.atleast_1d(grid[i])
388 pts[var] = new_pt if pts.get(var) is None else np.concatenate((pts[var], new_pt), axis=0)
390 return pts
392 def _expand_grid_coords(self, beta: MultiIndex):
393 """Iterable over all grid coordinates for a given `beta`, accounting for scalars and latent coefficients."""
394 grid_size = self.beta_to_knots(beta)
395 grid_coords = []
396 for s in grid_size:
397 if isinstance(s, int): # scalars
398 grid_coords.append(range(s))
399 else: # latent coefficients
400 grid_coords.append(itertools.product(*[range(latent_size) for latent_size in s]))
402 yield from itertools.product(*grid_coords)
404 @staticmethod
405 def _is_singleton(arr: np.ndarray):
406 return len(arr.shape) == 1 and arr.shape[0] == 1
408 @staticmethod
409 def _is_numeric(arr: np.ndarray):
410 return np.issubdtype(arr.dtype, np.number)
412 @classmethod
413 def _numeric_outputs(cls, yi_dict: dict[str, ArrayLike]) -> list[str]:
414 """Return a list of the output variables that have numeric data."""
415 output_vars = []
416 for var in yi_dict.keys():
417 try:
418 if cls._is_numeric(np.atleast_1d(yi_dict[var])):
419 output_vars.append(var)
420 except Exception:
421 continue
422 return output_vars
424 @staticmethod
425 def is_one_level_refinement(beta_old: tuple, beta_new: tuple) -> bool:
426 """Check if a new `beta` multi-index is a one-level refinement from a previous `beta`.
428 !!! Example
429 Refining from `(0, 1, 2)` to the new multi-index `(1, 1, 2)` is a one-level refinement. But refining to
430 either `(2, 1, 2)` or `(1, 2, 2)` are not, since more than one refinement occurs at the same time.
432 :param beta_old: the starting multi-index
433 :param beta_new: the new refined multi-index
434 :returns: whether `beta_new` is a one-level refinement from `beta_old`
435 """
436 level_diff = np.array(beta_new, dtype=int) - np.array(beta_old, dtype=int)
437 ind = np.nonzero(level_diff)[0]
438 return ind.shape[0] == 1 and level_diff[ind] == 1
440 def beta_to_knots(self, beta: MultiIndex, knots_per_level: int = None, latent_size: dict = None,
441 expand_latent_method: str = None) -> tuple:
442 """Convert a `beta` multi-index to the number of knots per dimension in the sparse grid.
444 :param beta: refinement level indices
445 :param knots_per_level: level-to-grid-size multiplier, i.e. number of new points (or knots) for each beta level
446 :param latent_size: the number of latent coefficients for each variable (0 if scalar); number of variables and
447 order should match the `beta` multi-index
448 :param expand_latent_method: method for expanding latent grids, either 'round-robin' or 'tensor-product'
449 :returns: the number of knots/points per dimension for the sparse grid
450 """
451 knots_per_level = knots_per_level or self.knots_per_level
452 latent_size = latent_size or self.latent_size
453 expand_latent_method = expand_latent_method or self.expand_latent_method
455 grid_size = []
456 for i, (var, num_latent) in enumerate(latent_size.items()):
457 if num_latent > 0:
458 match expand_latent_method:
459 case 'round-robin':
460 if beta[i] == 0:
461 grid_size.append((1,) * num_latent) # initializes all latent grids to 1
462 else:
463 latent_refine_idx = (beta[i] - 1) % num_latent
464 latent_refine_num = ((beta[i] - 1) // num_latent) + 1
465 latent_beta = tuple([latent_refine_num] * (latent_refine_idx + 1) +
466 [latent_refine_num - 1] * (num_latent - latent_refine_idx - 1))
467 latent_grid = [knots_per_level * latent_beta[j] + 1 for j in range(num_latent)]
468 grid_size.append(tuple(latent_grid))
469 case 'tensor-product':
470 grid_size.append((knots_per_level * beta[i] + 1,) * num_latent)
471 case other:
472 raise NotImplementedError(f"Unknown method for expanding latent grids: {other}")
473 else:
474 grid_size.append(knots_per_level * beta[i] + 1)
476 return tuple(grid_size)
478 @staticmethod
479 def collocation_1d(N: int, z_bds: tuple, z_pts: np.ndarray = None,
480 wt_fcn: callable = None, method='leja', opt_args=None) -> np.ndarray:
481 """Find the next `N` points in the 1d sequence of `z_pts` using the provided collocation method.
483 :param N: number of new points to add to the sequence
484 :param z_bds: bounds on the 1d domain
485 :param z_pts: current univariate sequence `(Nz,)`, start at middle of `z_bds` if `None`
486 :param wt_fcn: weighting function, uses a constant weight if `None`, callable as `wt_fcn(z)`
487 :param method: collocation method to use, currently only 'leja' is supported
488 :param opt_args: extra arguments for the global 1d `direct` optimizer
489 :returns: the univariate sequence `z_pts` augmented by `N` new points
490 """
491 opt_args = opt_args or {}
492 if wt_fcn is None:
493 wt_fcn = lambda z: 1
494 if z_pts is None:
495 z_pts = (z_bds[1] + z_bds[0]) / 2
496 N = N - 1
497 z_pts = np.atleast_1d(z_pts)
499 match method:
500 case 'leja':
501 # Construct Leja sequence by maximizing the Leja objective sequentially
502 for i in range(N):
503 obj_fun = lambda z: -wt_fcn(np.array(z)) * np.prod(np.abs(z - z_pts))
504 res = direct(obj_fun, [z_bds], **opt_args) # Use global DIRECT optimization over 1d domain
505 z_star = res.x
506 z_pts = np.concatenate((z_pts, z_star))
507 case other:
508 raise NotImplementedError(f"Unknown collocation method: {other}")
510 return z_pts