Coverage for src/amisc/component.py: 88%
853 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-11-05 19:26 +0000
1"""A `Component` is an `amisc` wrapper around a single discipline model. It manages surrogate construction and
2a hierarchy of modeling fidelities.
4!!! Info "Multi-indices in the MISC approximation"
5 A multi-index is a tuple of natural numbers, each specifying a level of fidelity. You will frequently see two
6 multi-indices: `alpha` and `beta`. The `alpha` (or $\\alpha$) indices specify physical model fidelity and get
7 passed to the model as an additional argument (e.g. things like discretization level, time step size, etc.).
8 The `beta` (or $\\beta$) indices specify surrogate refinement level, so typically an indication of the amount of
9 training data used or the complexity of the surrogate model. We divide $\\beta$ into `data_fidelity` and
10 `surrogate_fidelity` for specifying training data and surrogate model complexity, respectively.
12Includes:
14- `ModelKwargs` — a dataclass for storing model keyword arguments
15- `StringKwargs` — a dataclass for storing model keyword arguments as a string
16- `IndexSet` — a dataclass that maintains a list of multi-indices
17- `MiscTree` — a dataclass that maintains MISC data in a `dict` tree, indexed by `alpha` and `beta`
18- `Component` — a class that manages a single discipline model and its surrogate hierarchy
19"""
20from __future__ import annotations
22import ast
23import copy
24import inspect
25import itertools
26import logging
27import random
28import string
29import traceback
30import typing
31import warnings
32from collections import UserDict, deque
33from concurrent.futures import ALL_COMPLETED, Executor, wait
34from pathlib import Path
35from typing import Any, Callable, ClassVar, Iterable, Literal, Optional
37import numpy as np
38import yaml
39from pydantic import BaseModel, ConfigDict, ValidationInfo, field_validator
40from typing_extensions import TypedDict
42from amisc.interpolator import Interpolator, InterpolatorState, Lagrange
43from amisc.serialize import PickleSerializable, Serializable, StringSerializable, YamlSerializable
44from amisc.training import SparseGrid, TrainingData
45from amisc.typing import LATENT_STR_ID, Dataset, MultiIndex
46from amisc.utils import (
47 _get_yaml_path,
48 _inspect_assignment,
49 _inspect_function,
50 format_inputs,
51 format_outputs,
52 get_logger,
53 search_for_file,
54 to_model_dataset,
55 to_surrogate_dataset,
56)
57from amisc.variable import Variable, VariableList
59__all__ = ["ModelKwargs", "StringKwargs", "IndexSet", "MiscTree", "Component"]
60_VariableLike = list[Variable | dict | str] | str | Variable | dict | VariableList # Generic type for Variables
63class ModelKwargs(UserDict, Serializable):
64 """Default dataclass for storing model keyword arguments in a `dict`. If you have kwargs that require
65 more complicated serialization/specification than a plain `dict`, then you can subclass from here.
66 """
68 def serialize(self):
69 return self.data
71 @classmethod
72 def deserialize(cls, serialized_data):
73 return ModelKwargs(**serialized_data)
75 @classmethod
76 def from_dict(cls, config: dict) -> ModelKwargs:
77 """Create a `ModelKwargs` object from a `dict` configuration."""
78 method = config.pop('method', 'default_kwargs').lower()
79 match method:
80 case 'default_kwargs':
81 return ModelKwargs(**config)
82 case 'string_kwargs':
83 return StringKwargs(**config)
84 case other:
85 config['method'] = other
86 return ModelKwargs(**config) # Pass the method through
89class StringKwargs(StringSerializable, ModelKwargs):
90 """Dataclass for storing model keyword arguments as a string."""
91 def __repr__(self):
92 return str(self.data)
94 def __str__(self):
95 def format_value(value):
96 if isinstance(value, str):
97 return f'"{value}"'
98 else:
99 return str(value)
101 kw_str = ", ".join([f"{key}={format_value(value)}" for key, value in self.items()])
102 return f"ModelKwargs({kw_str})"
105class IndexSet(set, Serializable):
106 """Dataclass that maintains a list of multi-indices. Overrides basic `set` functionality to ensure
107 elements are formatted correctly as `(alpha, beta)`; that is, as a tuple of `alpha` and
108 `beta`, which are themselves instances of a [`MultiIndex`][amisc.typing.MultiIndex] tuple.
110 !!! Example "An example index set"
111 $\\mathcal{I} = [(\\alpha, \\beta)_1 , (\\alpha, \\beta)_2, (\\alpha, \\beta)_3 , ...]$ would be specified
112 as `I = [((0, 0), (0, 0, 0)) , ((0, 1), (0, 1, 0)), ...]`.
113 """
114 def __init__(self, s=()):
115 s = [self._validate_element(ele) for ele in s]
116 super().__init__(s)
118 def __str__(self):
119 return str(list(self))
121 def __repr__(self):
122 return self.__str__()
124 def add(self, __element):
125 super().add(self._validate_element(__element))
127 def update(self, __elements):
128 super().update([self._validate_element(ele) for ele in __elements])
130 @classmethod
131 def _validate_element(cls, element):
132 """Validate that the element is a tuple of two multi-indices."""
133 alpha, beta = ast.literal_eval(element) if isinstance(element, str) else tuple(element)
134 return MultiIndex(alpha), MultiIndex(beta)
136 @classmethod
137 def _wrap_methods(cls, names):
138 """Make sure set operations return an `IndexSet` object."""
139 def wrap_method_closure(name):
140 def inner(self, *args):
141 result = getattr(super(cls, self), name)(*args)
142 if isinstance(result, set):
143 result = cls(result)
144 return result
145 inner.fn_name = name
146 setattr(cls, name, inner)
148 for name in names:
149 wrap_method_closure(name)
151 def serialize(self) -> list[str]:
152 """Return a list of each multi-index in the set serialized to a string."""
153 return [str(ele) for ele in self]
155 @classmethod
156 def deserialize(cls, serialized_data: list[str]) -> IndexSet:
157 """Deserialize a list of tuples to an `IndexSet`."""
158 return cls(serialized_data)
161IndexSet._wrap_methods(['__ror__', 'difference_update', '__isub__', 'symmetric_difference', '__rsub__', '__and__',
162 '__rand__', 'intersection', 'difference', '__iand__', 'union', '__ixor__',
163 'symmetric_difference_update', '__or__', 'copy', '__rxor__', 'intersection_update', '__xor__',
164 '__ior__', '__sub__'
165 ])
168class MiscTree(UserDict, Serializable):
169 """Dataclass that maintains MISC data in a `dict` tree, indexed by `alpha` and `beta`. Overrides
170 basic `dict` functionality to ensure elements are formatted correctly as `(alpha, beta) -> data`.
171 Used to store MISC coefficients, model costs, and interpolator states.
173 The underlying data structure is: `dict[MultiIndex, dict[MultiIndex, float | InterpolatorState]]`.
174 """
175 SERIALIZER_KEY = 'state_serializer'
177 def __init__(self, data: dict = None, **kwargs):
178 data_dict = data or {}
179 if isinstance(data_dict, MiscTree):
180 data_dict = data_dict.data
181 data_dict.update(kwargs)
182 super().__init__(self._validate_data(data_dict))
184 def serialize(self, *args, keep_yaml_objects=False, **kwargs) -> dict:
185 """Serialize `alpha, beta` indices to string and return a `dict` of internal data.
187 :param args: extra serialization arguments for internal `InterpolatorState`
188 :param keep_yaml_objects: whether to keep `YamlSerializable` instances in the serialization
189 :param kwargs: extra serialization keyword arguments for internal `InterpolatorState`
190 """
191 ret_dict = {}
192 if state_serializer := self.state_serializer(self.data):
193 ret_dict[self.SERIALIZER_KEY] = state_serializer.obj if keep_yaml_objects else state_serializer.serialize()
194 for alpha, beta, data in self:
195 ret_dict.setdefault(str(alpha), dict())
196 serialized_data = data.serialize(*args, **kwargs) if isinstance(data, InterpolatorState) else float(data)
197 ret_dict[str(alpha)][str(beta)] = serialized_data
198 return ret_dict
200 @classmethod
201 def deserialize(cls, serialized_data: dict) -> MiscTree:
202 """Deserialize a `dict` to a `MiscTree`.
204 :param serialized_data: the data to deserialize to a `MiscTree` object
205 """
206 return cls(serialized_data)
208 @classmethod
209 def state_serializer(cls, data: dict) -> YamlSerializable | None:
210 """Infer and return the interpolator state serializer from the `MiscTree` data (if possible). If no
211 `InterpolatorState` instance could be found, return `None`.
212 """
213 serializer = data.get(cls.SERIALIZER_KEY, None) # if `data` is serialized
214 if serializer is None: # Otherwise search for an InterpolatorState
215 for alpha, beta_dict in data.items():
216 if alpha == cls.SERIALIZER_KEY:
217 continue
218 for beta, value in beta_dict.items():
219 if isinstance(value, InterpolatorState):
220 serializer = type(value)
221 break
222 if serializer is not None:
223 break
224 return cls._validate_state_serializer(serializer)
226 @classmethod
227 def _validate_state_serializer(cls, state_serializer: Optional[str | type[Serializable] | YamlSerializable]
228 ) -> YamlSerializable | None:
229 if state_serializer is None:
230 return None
231 elif isinstance(state_serializer, YamlSerializable):
232 return state_serializer
233 elif isinstance(state_serializer, str):
234 return YamlSerializable.deserialize(state_serializer) # Load the serializer type from string
235 else:
236 return YamlSerializable(obj=state_serializer)
238 @classmethod
239 def _validate_data(cls, serialized_data: dict) -> dict:
240 state_serializer = cls.state_serializer(serialized_data)
241 ret_dict = {}
242 for alpha, beta_dict in serialized_data.items():
243 if alpha == cls.SERIALIZER_KEY:
244 continue
245 alpha_tup = MultiIndex(alpha)
246 ret_dict.setdefault(alpha_tup, dict())
247 for beta, data in beta_dict.items():
248 beta_tup = MultiIndex(beta)
249 if isinstance(data, InterpolatorState):
250 pass
251 elif state_serializer is not None:
252 data = state_serializer.obj.deserialize(data)
253 else:
254 data = float(data)
255 assert isinstance(data, InterpolatorState | float)
256 ret_dict[alpha_tup][beta_tup] = data
257 return ret_dict
259 @staticmethod
260 def _is_alpha_beta_access(key):
261 """Check that the key is of the format `(alpha, beta).`"""
262 return (isinstance(key, tuple) and len(key) == 2 and isinstance(key[0], str | tuple)
263 and isinstance(key[1], str | tuple))
265 def get(self, key, default=None) -> float | InterpolatorState:
266 try:
267 return self.__getitem__(key)
268 except Exception:
269 return default
271 def update(self, data_dict: dict = None, **kwargs):
272 """Force `dict.update()` through the validator."""
273 data_dict = data_dict or dict()
274 data_dict.update(kwargs)
275 super().update(self._validate_data(data_dict))
277 def __setitem__(self, key: tuple | MultiIndex, value: float | InterpolatorState):
278 """Allows `misc_tree[alpha, beta] = value` usage."""
279 if self._is_alpha_beta_access(key):
280 alpha, beta = MultiIndex(key[0]), MultiIndex(key[1])
281 self.data.setdefault(alpha, dict())
282 self.data[alpha][beta] = value
283 else:
284 super().__setitem__(MultiIndex(key), value)
286 def __getitem__(self, key: tuple | MultiIndex) -> float | InterpolatorState:
287 """Allows `value = misc_tree[alpha, beta]` usage."""
288 if self._is_alpha_beta_access(key):
289 alpha, beta = MultiIndex(key[0]), MultiIndex(key[1])
290 return self.data[alpha][beta]
291 else:
292 return super().__getitem__(MultiIndex(key))
294 def clear(self):
295 """Clear the `MiscTree` data."""
296 for key in list(self.data.keys()):
297 del self.data[key]
299 def __eq__(self, other):
300 if isinstance(other, MiscTree):
301 try:
302 for alpha, beta, data in self:
303 if other[alpha, beta] != data:
304 return False
305 return True
306 except KeyError:
307 return False
308 else:
309 return False
311 def __iter__(self) -> Iterable[tuple[tuple, tuple, float | InterpolatorState]]:
312 for alpha, beta_dict in self.data.items():
313 if alpha == self.SERIALIZER_KEY:
314 continue
315 for beta, data in beta_dict.items():
316 yield alpha, beta, data
319class ComponentSerializers(TypedDict, total=False):
320 """Type hint for the `Component` class data serializers.
322 :ivar model_kwargs: the model kwarg object class
323 :ivar interpolator: the interpolator object class
324 :ivar training_data: the training data object class
325 """
326 model_kwargs: str | type[Serializable] | YamlSerializable
327 interpolator: str | type[Serializable] | YamlSerializable
328 training_data: str | type[Serializable] | YamlSerializable
331class Component(BaseModel, Serializable):
332 """A `Component` wrapper around a single discipline model. It manages MISC surrogate construction and a hierarchy of
333 modeling fidelities.
335 A `Component` can be constructed by specifying a model, input and output variables, and additional configurations
336 such as the maximum fidelity levels, the interpolator type, and the training data type. If `model_fidelity`,
337 `data_fidelity`, and `surrogate_fidelity` are all left empty, then the `Component` will not use a surrogate model,
338 instead calling the underlying model directly. The `Component` can be serialized to a YAML file and deserialized
339 back into a Python object.
341 !!! Example "A simple `Component`"
342 ```python
343 from amisc import Component, Variable
345 x = Variable(domain=(0, 1))
346 y = Variable()
347 model = lambda x: {'y': x['x']**2}
348 comp = Component(model=model, inputs=[x], outputs=[y])
349 ```
351 Each fidelity index in $\\alpha$ increases in refinement from $0$ up to `model_fidelity`. Each fidelity index
352 in $\\beta$ increases from $0$ up to `(data_fidelity, surrogate_fidelity)`. From the `Component's` perspective,
353 the concatenation of $(\\alpha, \\beta)$ fully specifies a single fidelity "level". The `Component`
354 forms an approximation of the model by summing up over many of these concatenated sets of $(\\alpha, \\beta)$.
356 :ivar name: the name of the `Component`
357 :ivar model: the model or function that is to be approximated, callable as `y = f(x)`
358 :ivar inputs: the input variables to the model
359 :ivar outputs: the output variables from the model
360 :ivar model_kwargs: extra keyword arguments to pass to the model
361 :ivar model_fidelity: the maximum level of refinement for each fidelity index in $\\alpha$ for model fidelity
362 :ivar data_fidelity: the maximum level of refinement for each fidelity index in $\\beta$ for training data
363 :ivar surrogate_fidelity: the max level of refinement for each fidelity index in $\\beta$ for the surrogate
364 :ivar interpolator: the interpolator to use as the underlying surrogate model
365 :ivar vectorized: whether the model supports vectorized input/output (i.e. datasets with arbitrary shape `(...,)`)
366 :ivar call_unpacked: whether the model expects unpacked input arguments (i.e. `func(x1, x2, ...)`)
367 :ivar ret_unpacked: whether the model returns unpacked output arguments (i.e. `func() -> (y1, y2, ...)`)
369 :ivar active_set: the current active set of multi-indices in the MISC approximation
370 :ivar candidate_set: all neighboring multi-indices that are candidates for inclusion in `active_set`
371 :ivar misc_states: the interpolator states for each multi-index in the MISC approximation
372 :ivar misc_costs: the computational cost associated with each multi-index in the MISC approximation
373 :ivar misc_coeff_train: the combination technique coefficients for the active set multi-indices
374 :ivar misc_coeff_test: the combination technique coefficients for the active and candidate set multi-indices
375 :ivar model_costs: the average single fidelity model costs for each $\\alpha$
376 :ivar training_data: the training data storage structure for the surrogate model
378 :ivar serializers: the custom serializers for the `[model_kwargs, interpolator, training_data]`
379 `Component` attributes -- these should be the _types_ of the serializer objects, which will
380 be inferred from the data passed in if not explicitly set
381 :ivar _logger: the logger for the `Component`
382 """
383 yaml_tag: ClassVar[str] = u'!Component'
384 model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True, validate_default=True,
385 protected_namespaces=(), extra='allow')
386 # Configuration
387 serializers: Optional[ComponentSerializers] = None
388 name: Optional[str] = None
389 model: str | Callable[[dict | Dataset, ...], dict | Dataset]
390 model_kwargs: str | dict | ModelKwargs = {}
391 inputs: _VariableLike
392 outputs: _VariableLike
393 model_fidelity: str | tuple = MultiIndex()
394 data_fidelity: str | tuple = MultiIndex()
395 surrogate_fidelity: str | tuple = MultiIndex()
396 interpolator: Any | Interpolator = Lagrange()
397 vectorized: bool = False
398 call_unpacked: Optional[bool] = None # If the model expects inputs/outputs like `func(x1, x2, ...)->(y1, y2, ...)
399 ret_unpacked: Optional[bool] = None
401 # Data storage/states for a MISC component
402 active_set: list | set | IndexSet = IndexSet() # set of active (alpha, beta) multi-indices
403 candidate_set: list | set | IndexSet = IndexSet() # set of candidate (alpha, beta) multi-indices
404 misc_states: dict | MiscTree = MiscTree() # (alpha, beta) -> Interpolator state
405 misc_costs: dict | MiscTree = MiscTree() # (alpha, beta) -> Added computational cost for this mult-index
406 misc_coeff_train: dict | MiscTree = MiscTree() # (alpha, beta) -> c_[alpha, beta] (active set only)
407 misc_coeff_test: dict | MiscTree = MiscTree() # (alpha, beta) -> c_[alpha, beta] (including candidate set)
408 model_costs: dict = dict() # Average single fidelity model costs (for each alpha)
409 training_data: Any | TrainingData = SparseGrid() # Stores surrogate training data
411 # Internal
412 _logger: Optional[logging.Logger] = None
414 def __init__(self, /, model, *args, inputs=None, outputs=None, name=None, **kwargs):
415 if name is None:
416 name = _inspect_assignment('Component') # try to assign the name from inspection
417 name = name or model.__name__ or "Component_" + "".join(random.choices(string.digits, k=3))
419 # Determine how the model expects to be called and gather inputs/outputs
420 _ = self._validate_model_signature(model, args, inputs, outputs, kwargs.get('call_unpacked', None),
421 kwargs.get('ret_unpacked', None))
422 model, inputs, outputs, call_unpacked, ret_unpacked = _
423 kwargs['call_unpacked'] = call_unpacked
424 kwargs['ret_unpacked'] = ret_unpacked
426 # Gather all model kwargs (anything else passed in for kwargs is assumed to be a model kwarg)
427 model_kwargs = kwargs.get('model_kwargs', {})
428 for key in kwargs.keys() - self.model_fields.keys():
429 model_kwargs[key] = kwargs.pop(key)
430 kwargs['model_kwargs'] = model_kwargs
432 # Gather data serializers from type checks (if not passed in as a kwarg)
433 serializers = kwargs.get('serializers', {}) # directly passing serializers will override type checks
434 for key in ComponentSerializers.__annotations__.keys():
435 field = kwargs.get(key, None)
436 if isinstance(field, dict):
437 field_super = next(filter(lambda x: issubclass(x, Serializable),
438 typing.get_args(self.model_fields[key].annotation)), None)
439 field = field_super.from_dict(field) if field_super is not None else field
440 if not serializers.get(key, None):
441 serializers[key] = type(field) if isinstance(field, Serializable) else (
442 type(self.model_fields[key].default))
443 kwargs['serializers'] = serializers
445 super().__init__(model=model, inputs=inputs, outputs=outputs, name=name, **kwargs) # Runs pydantic validation
447 # Set internal properties
448 assert self.is_downward_closed(self.active_set.union(self.candidate_set))
449 self.set_logger()
451 @classmethod
452 def _validate_model_signature(cls, model, args=(), inputs=None, outputs=None,
453 call_unpacked=None, ret_unpacked=None):
454 """Parse model signature and decide how the model expects to be called based on what input/output information
455 is provided or inspected from the model signature.
456 """
457 if inputs is not None:
458 inputs = cls._validate_variables(inputs)
459 if outputs is not None:
460 outputs = cls._validate_variables(outputs)
461 model = cls._validate_model(model)
463 # Default to `dict` (i.e. packed) model call/return signatures
464 if call_unpacked is None:
465 call_unpacked = False
466 if ret_unpacked is None:
467 ret_unpacked = False
468 inputs_inspect, outputs_inspect = _inspect_function(model)
469 call_unpacked = call_unpacked or (len(inputs_inspect) > 1) # Assume multiple inputs require unpacking
470 ret_unpacked = ret_unpacked or (len(outputs_inspect) > 1) # Assume multiple outputs require unpacking
472 # Extract inputs/outputs from args
473 arg_inputs = ()
474 arg_outputs = ()
475 if len(args) > 0:
476 if call_unpacked:
477 if isinstance(args[0], dict | str | Variable):
478 arg_inputs = args[:len(inputs_inspect)]
479 arg_outputs = args[len(inputs_inspect):]
480 else:
481 arg_inputs = args[0]
482 arg_outputs = args[1:]
483 else:
484 arg_inputs = args[0] # Assume first arg is a single or list of inputs
485 arg_outputs = args[1:] # Assume rest are outputs
487 # Resolve inputs
488 inputs = inputs or []
489 inputs = VariableList.merge(inputs, arg_inputs)
490 if len(inputs) == 0:
491 inputs = inputs_inspect
492 call_unpacked = True
493 if len(inputs) == 0:
494 raise ValueError("Could not infer input variables from model signature. Either your model does not "
495 "accept input arguments or an error occurred during inspection.\nPlease provide the "
496 "inputs directly as `Component(inputs=[...])` or fix the model signature.")
497 if call_unpacked:
498 if not all([var == inputs_inspect[i] for i, var in enumerate(inputs)]):
499 warnings.warn(f"Mismatch between provided inputs: {inputs.values()} and inputs inferred from "
500 f"model signature: {inputs_inspect}. This may cause unexpected results.")
501 else:
502 if len(inputs_inspect) > 1:
503 warnings.warn(f"Model signature expects multiple input arguments: {inputs_inspect}. "
504 f"Please set `call_unpacked=True` to use this model signature for multiple "
505 f"inputs.\nOtherwise, move all inputs into a single `dict` argument and all "
506 f"extra arguments into the `model_kwargs` field.")
508 # Can't assume unpacked for single input/output, so warn user if they may be trying to do so
509 if len(inputs) == 1 and len(inputs_inspect) == 1 and str(inputs[0]) == str(inputs_inspect[0]):
510 warnings.warn(f"Single input argument: {inputs[0]} provided to model with input signature: "
511 f"{inputs_inspect}.\nIf you intended to use a single input argument, set "
512 f"`call_unpacked=True` to use this model signature.\nOtherwise, the first input will "
513 f"be passed to your model as a `dict`.\nIf you are expecting a `dict` input already, "
514 f"change the name of the input to not exactly "
515 f"match {inputs_inspect} in order to silence this warning.")
516 # Resolve outputs
517 outputs = outputs or []
518 outputs = VariableList.merge(outputs, *arg_outputs)
519 if len(outputs) == 0:
520 outputs = outputs_inspect
521 ret_unpacked = True
522 if len(outputs) == 0:
523 raise ValueError("Could not infer output variables from model inspection. Either your model does not "
524 "return outputs or an error occurred during inspection.\nPlease provide the "
525 "outputs directly as `Component(outputs=[...])` or fix the model return values.")
526 if ret_unpacked:
527 if not all([var == outputs_inspect[i] for i, var in enumerate(outputs)]):
528 warnings.warn(f"Mismatch between provided outputs: {outputs.values()} and outputs inferred "
529 f"from model: {outputs_inspect}. This may cause unexpected results.")
530 else:
531 if len(outputs_inspect) > 1:
532 warnings.warn(f"Model expects multiple return values: {outputs_inspect}. Please set "
533 f"`ret_unpacked=True` to use this model signature for multiple outputs.\n"
534 f"Otherwise, move all outputs into a single `dict` return value.")
536 if len(outputs) == 1 and len(outputs_inspect) == 1 and str(outputs[0]) == str(outputs_inspect[0]):
537 warnings.warn(f"Single output: {outputs[0]} provided to model with single expected return: "
538 f"{outputs_inspect}.\nIf you intended to output a single return value, set "
539 f"`ret_unpacked=True` to use this model signature.\nOtherwise, the output should "
540 f"be returned from your model as a `dict`.\nIf you are returning a `dict` already, "
541 f"then change its name to not exactly match {outputs_inspect} in order to silence "
542 f"this warning.")
543 return model, inputs, outputs, call_unpacked, ret_unpacked
545 def __repr__(self):
546 s = f'---- {self.name} ----\n'
547 s += f'Inputs: {self.inputs}\n'
548 s += f'Outputs: {self.outputs}\n'
549 s += f'Model: {self.model}'
550 return s
552 def __str__(self):
553 return self.__repr__()
555 @field_validator('serializers')
556 @classmethod
557 def _validate_serializers(cls, serializers: ComponentSerializers) -> ComponentSerializers:
558 """Make sure custom serializer object types are themselves serializable as `YamlSerializable`."""
559 for key, serializer in serializers.items():
560 if serializer is None:
561 serializers[key] = None
562 elif isinstance(serializer, YamlSerializable):
563 serializers[key] = serializer
564 elif isinstance(serializer, str):
565 serializers[key] = YamlSerializable.deserialize(serializer)
566 else:
567 serializers[key] = YamlSerializable(obj=serializer)
568 return serializers
570 @field_validator('model')
571 @classmethod
572 def _validate_model(cls, model: str | Callable) -> Callable:
573 """Expects model as a callable or a yaml !!python/name string representation."""
574 if isinstance(model, str):
575 return YamlSerializable.deserialize(model).obj
576 else:
577 return model
579 @field_validator('inputs', 'outputs')
580 @classmethod
581 def _validate_variables(cls, variables: _VariableLike) -> VariableList:
582 if isinstance(variables, VariableList):
583 return variables
584 else:
585 return VariableList.deserialize(variables)
587 @field_validator('model_fidelity', 'data_fidelity', 'surrogate_fidelity')
588 @classmethod
589 def _validate_indices(cls, multi_index) -> MultiIndex:
590 return MultiIndex(multi_index)
592 @field_validator('active_set', 'candidate_set')
593 @classmethod
594 def _validate_index_set(cls, index_set) -> IndexSet:
595 return IndexSet.deserialize(index_set)
597 @field_validator('misc_states', 'misc_costs', 'misc_coeff_train', 'misc_coeff_test')
598 @classmethod
599 def _validate_misc_tree(cls, misc_tree) -> MiscTree:
600 return MiscTree.deserialize(misc_tree)
602 @field_validator('model_costs')
603 @classmethod
604 def _validate_model_costs(cls, model_costs: dict) -> dict:
605 return {MultiIndex(key): float(value) for key, value in model_costs.items()}
607 @field_validator('model_kwargs', 'interpolator', 'training_data')
608 @classmethod
609 def _validate_arbitrary_serializable(cls, data: Any, info: ValidationInfo) -> Any:
610 """Use the stored custom serialization classes to deserialize arbitrary objects."""
611 serializer = info.data.get('serializers').get(info.field_name).obj
612 if isinstance(data, Serializable):
613 return data
614 else:
615 return serializer.deserialize(data)
617 @property
618 def xdim(self) -> int:
619 return len(self.inputs)
621 @property
622 def ydim(self) -> int:
623 return len(self.outputs)
625 @property
626 def max_alpha(self) -> MultiIndex:
627 """The maximum model fidelity multi-index (alias for `model_fidelity`)."""
628 return self.model_fidelity
630 @property
631 def max_beta(self) -> MultiIndex:
632 """The maximum surrogate fidelity multi-index is a combination of training and interpolator indices."""
633 return self.data_fidelity + self.surrogate_fidelity
635 @property
636 def has_surrogate(self) -> bool:
637 """The component has no surrogate model if there are no fidelity indices."""
638 return (len(self.max_alpha) + len(self.max_beta)) > 0
640 @property
641 def logger(self) -> logging.Logger:
642 return self._logger
644 @logger.setter
645 def logger(self, logger: logging.Logger):
646 self._logger = logger
648 def __eq__(self, other):
649 if isinstance(other, Component):
650 return (self.model.__code__.co_code == other.model.__code__.co_code and self.inputs == other.inputs
651 and self.outputs == other.outputs and self.name == other.name
652 and self.model_kwargs.data == other.model_kwargs.data
653 and self.model_fidelity == other.model_fidelity and self.max_beta == other.max_beta and
654 self.interpolator == other.interpolator
655 and self.active_set == other.active_set and self.candidate_set == other.candidate_set
656 and self.misc_states == other.misc_states and self.misc_costs == other.misc_costs
657 )
658 else:
659 return False
661 def _neighbors(self, alpha: MultiIndex, beta: MultiIndex, active_set: IndexSet = None, forward: bool = True):
662 """Get all possible forward or backward multi-index neighbors (distance of one unit vector away).
664 :param alpha: the model fidelity index
665 :param beta: the surrogate fidelity index
666 :param active_set: the set of active multi-indices
667 :param forward: whether to get forward or backward neighbors
668 :returns: a set of multi-indices that are neighbors of the input multi-index pair `(alpha, beta)`
669 """
670 active_set = active_set or self.active_set
671 ind = np.array(alpha + beta)
672 max_ind = np.array(self.max_alpha + self.max_beta)
673 new_candidates = IndexSet()
674 for i in range(len(ind)):
675 ind_new = ind.copy()
676 ind_new[i] += 1 if forward else -1
678 # Don't add if we surpass a refinement limit or lower bound
679 if np.any(ind_new > max_ind) or np.any(ind_new < 0):
680 continue
682 # Add the new index if it maintains downward-closedness
683 down_closed = True
684 for j in range(len(ind)):
685 ind_check = ind_new.copy()
686 ind_check[j] -= 1
687 if ind_check[j] >= 0:
688 tup_check = (MultiIndex(ind_check[:len(alpha)]), MultiIndex(ind_check[len(alpha):]))
689 if tup_check not in active_set and tup_check != (alpha, beta):
690 down_closed = False
691 break
692 if down_closed:
693 new_candidates.add((ind_new[:len(alpha)], ind_new[len(alpha):]))
695 return new_candidates
697 def _surrogate_outputs(self):
698 """Helper function to get the names of the surrogate outputs (including latent variables)."""
699 y_vars = []
700 for var in self.outputs:
701 if var.compression is not None:
702 for i in range(var.compression.latent_size()):
703 y_vars.append(f'{var.name}{LATENT_STR_ID}{i}')
704 else:
705 y_vars.append(var.name)
706 return y_vars
708 def _match_index_set(self, index_set, misc_coeff):
709 """Helper function to grab the correct data structures for the given index set and MISC coefficients."""
710 if misc_coeff is None:
711 match index_set:
712 case 'train':
713 misc_coeff = self.misc_coeff_train
714 case 'test':
715 misc_coeff = self.misc_coeff_test
716 case other:
717 raise ValueError(f"Index set must be 'train' or 'test' if you do not provide `misc_coeff`. "
718 f"{other} not recognized.")
719 if isinstance(index_set, str):
720 match index_set:
721 case 'train':
722 index_set = self.active_set
723 case 'test':
724 index_set = self.active_set.union(self.candidate_set)
725 case other:
726 raise ValueError(f"Index set must be 'train' or 'test'. {other} not recognized.")
728 return index_set, misc_coeff
730 def get_training_data(self, alpha: Literal['best', 'worst'] | MultiIndex = 'best',
731 beta: Literal['best', 'worst'] | MultiIndex = 'best',
732 y_vars: list = None) -> tuple[Dataset, Dataset]:
733 """Get all training data for a given multi-index pair `(alpha, beta)`.
735 :param alpha: the model fidelity index (defaults to the maximum available model fidelity)
736 :param beta: the surrogate fidelity index (defaults to the maximum available surrogate fidelity)
737 :param y_vars: the training data to return (defaults to all stored data)
738 :returns: `(xtrain, ytrain)` - the training data for the given multi-indices
739 """
740 # Find the best alpha
741 if alpha == 'best':
742 alpha_best = ()
743 for a, _ in self.active_set.union(self.candidate_set):
744 if sum(a) > sum(alpha_best):
745 alpha_best = a
746 alpha = alpha_best
747 elif alpha == 'worst':
748 alpha = (0,) * len(self.max_alpha)
750 # Find the best beta for the given alpha
751 if beta == 'best':
752 beta_best = ()
753 for a, b in self.active_set.union(self.candidate_set):
754 if a == alpha and sum(b) > sum(beta_best):
755 beta_best = b
756 beta = beta_best
757 elif beta == 'worst':
758 beta = (0,) * len(self.max_beta)
760 try:
761 return self.training_data.get(alpha, beta, y_vars=y_vars, skip_nan=True)
762 except Exception as e:
763 self.logger.error(f"Error getting training data for alpha={alpha}, beta={beta}.")
764 raise e
766 def call_model(self, inputs: dict | Dataset,
767 model_fidelity: Literal['best', 'worst'] | tuple | list = None,
768 output_path: str | Path = None,
769 executor: Executor = None,
770 **kwds) -> Dataset:
771 """Wrapper function for calling the underlying component model.
773 This function formats the input data, calls the model, and processes the output data.
774 It supports vectorized calls, parallel execution using an executor, or serial execution. These options are
775 checked in that order, with the first available method used. Must set `Component.vectorized=True` if the
776 model supports input arrays of the form `(N,)` or even arbitrary shape `(...,)`.
778 !!! Warning "Parallel Execution"
779 The underlying model must be defined in a global module scope if `pickle` is the serialization method for
780 the provided `Executor`.
782 !!! Note "Additional return values"
783 The model can return additional items that are not part of `Component.outputs`. These items are returned
784 as object arrays in the output `dict`. Two special return values are `model_cost` and `output_path`.
785 Returning `model_cost` will store the computational cost of a single model evaluation (which is used by
786 `amisc` adaptive surrogate training). Returning `output_path` will store the output file name if the model
787 wrote any files to disk.
789 !!! Note "Handling errors"
790 If the underlying component model raises an exception, the error is stored in `output_dict['errors']` with
791 the index of the input data that caused the error. The output data for that index is set to `np.nan`
792 for each output variable.
794 :param inputs: The input data for the model, formatted as a `dict` with a key for each input variable and
795 a corresponding value that is an array of the input data. If specified as a plain list, then the
796 order is assumed the same as `Component.inputs`.
797 :param model_fidelity: Fidelity indices to tune the model fidelity (model must request this
798 in its keyword arguments).
799 :param output_path: Directory to save model output files (model must request this in its keyword arguments).
800 :param executor: Executor for parallel execution if the model is not vectorized (optional).
801 :param kwds: Additional keyword arguments to pass to the model (model must request these in its keyword args).
802 :returns: The output data from the model, formatted as a `dict` with a key for each output variable and a
803 corresponding value that is an array of the output data.
804 """
805 # Format inputs to a common loop shape (fail if missing any)
806 if len(inputs) == 0:
807 return {} # your fault
808 if isinstance(inputs, list | np.ndarray):
809 inputs = np.atleast_1d(inputs)
810 inputs = {var.name: inputs[..., i] for i, var in enumerate(self.inputs)}
812 var_shape = {}
813 for var in self.inputs:
814 s = None
815 if (arr := kwds.get(f'{var.name}_coords')) is not None:
816 s = arr.shape if len(arr.shape) == 1 else arr.shape[:-1] # skip the coordinate dim (last axis)
817 if var.compression is not None:
818 for field in var.compression.fields:
819 var_shape[field] = s
820 else:
821 var_shape[var.name] = s
822 inputs, loop_shape = format_inputs(inputs, var_shape=var_shape)
824 N = int(np.prod(loop_shape))
825 list_alpha = isinstance(model_fidelity, list | np.ndarray)
826 alpha_requested = self.model_kwarg_requested('model_fidelity')
827 for var in self.inputs:
828 if var.compression is not None:
829 for field in var.compression.fields:
830 if field not in inputs:
831 raise ValueError(f"Missing field '{field}' for input variable '{var}'.")
832 elif var.name not in inputs:
833 raise ValueError(f"Missing input variable '{var.name}'.")
835 # Pass extra requested items to the model kwargs
836 kwargs = copy.deepcopy(self.model_kwargs.data)
837 if self.model_kwarg_requested('output_path'):
838 kwargs['output_path'] = output_path
839 if self.model_kwarg_requested('input_vars'):
840 kwargs['input_vars'] = self.inputs
841 if self.model_kwarg_requested('output_vars'):
842 kwargs['output_vars'] = self.outputs
843 if alpha_requested:
844 if not list_alpha:
845 model_fidelity = [model_fidelity] * N
846 for i in range(N):
847 if model_fidelity[i] == 'best':
848 model_fidelity[i] = self.max_alpha
849 elif model_fidelity[i] == 'worst':
850 model_fidelity[i] = (0,) * len(self.model_fidelity)
852 for k, v in kwds.items():
853 if self.model_kwarg_requested(k):
854 kwargs[k] = v
856 # Compute model (vectorized, executor parallel, or serial)
857 errors = {}
858 if self.vectorized:
859 if alpha_requested:
860 kwargs['model_fidelity'] = np.atleast_1d(model_fidelity).reshape((N, -1))
861 output_dict = self.model(*[inputs[var.name] for var in self.inputs], **kwargs) if self.call_unpacked \
862 else self.model(inputs, **kwargs)
863 if self.ret_unpacked:
864 output_dict = (output_dict,) if not isinstance(output_dict, tuple) else output_dict
865 output_dict = {out_var.name: output_dict[i] for i, out_var in enumerate(self.outputs)}
866 else:
867 if executor is None: # Serial
868 results = deque(maxlen=N)
869 for i in range(N):
870 try:
871 if alpha_requested:
872 kwargs['model_fidelity'] = model_fidelity[i]
873 ret = self.model(*[{k: v[i] for k, v in inputs.items()}[var.name] for var in self.inputs],
874 **kwargs) if self.call_unpacked else (
875 self.model({k: v[i] for k, v in inputs.items()}, **kwargs))
876 if self.ret_unpacked:
877 ret = (ret,) if not isinstance(ret, tuple) else ret
878 ret = {out_var.name: ret[i] for i, out_var in enumerate(self.outputs)}
879 results.append(ret)
880 except Exception:
881 results.append({'inputs': {k: v[i] for k, v in inputs.items()}, 'index': i,
882 'model_kwargs': kwargs.copy(), 'error': traceback.format_exc()})
883 else: # Parallel
884 results = deque(maxlen=N)
885 futures = []
886 for i in range(N):
887 if alpha_requested:
888 kwargs['model_fidelity'] = model_fidelity[i]
889 fs = executor.submit(self.model,
890 *[{k: v[i] for k, v in inputs.items()}[var.name] for var in self.inputs],
891 **kwargs) if self.call_unpacked else (
892 executor.submit(self.model, {k: v[i] for k, v in inputs.items()}, **kwargs))
893 futures.append(fs)
894 wait(futures, timeout=None, return_when=ALL_COMPLETED)
896 for i, fs in enumerate(futures):
897 try:
898 if alpha_requested:
899 kwargs['model_fidelity'] = model_fidelity[i]
900 ret = fs.result()
901 if self.ret_unpacked:
902 ret = (ret,) if not isinstance(ret, tuple) else ret
903 ret = {out_var.name: ret[i] for i, out_var in enumerate(self.outputs)}
904 results.append(ret)
905 except Exception:
906 results.append({'inputs': {k: v[i] for k, v in inputs.items()}, 'index': i,
907 'model_kwargs': kwargs.copy(), 'error': traceback.format_exc()})
909 # Collect parallel/serial results
910 output_dict = {}
911 for i in range(N):
912 res = results.popleft()
913 if 'error' in res:
914 errors[i] = res
915 else:
916 for key, val in res.items():
917 # Save numeric outputs
918 numeric_flag = False
919 for var in self.outputs:
920 if var.compression is not None: # field quantity return values
921 if key in var.compression.fields:
922 if output_dict.get(key) is None:
923 output_dict.setdefault(key, np.full((N, *np.atleast_1d(val).shape), np.nan))
924 output_dict[key][i, ...] = np.atleast_1d(val)
925 numeric_flag = True
926 break
927 elif key == var:
928 if output_dict.get(key) is None:
929 output_dict.setdefault(key, np.full((N, *np.atleast_1d(val).shape), np.nan))
930 output_dict[key][i, ...] = np.atleast_1d(val)
931 numeric_flag = True
932 break
933 # Otherwise, save other objects
934 if not numeric_flag:
935 if key == 'model_cost':
936 if output_dict.get(key) is None:
937 output_dict.setdefault(key, np.full((N,), np.nan))
938 output_dict[key][i] = val
939 else:
940 if output_dict.get(key) is None:
941 output_dict.setdefault(key, np.full((N,), None, dtype=object))
942 output_dict[key][i] = val
944 # Save average model costs for each alpha fidelity
945 if model_fidelity is not None and output_dict.get('model_cost') is not None:
946 alpha_costs = {}
947 for i, cost in enumerate(output_dict['model_cost']):
948 alpha_costs.setdefault(MultiIndex(model_fidelity[i]), [])
949 alpha_costs[MultiIndex(model_fidelity[i])].append(cost)
950 for a, costs in alpha_costs.items():
951 self.model_costs.setdefault(a, np.empty(0))
952 self.model_costs[a] = np.nanmean(np.hstack((costs, self.model_costs[a])))
954 # Reshape loop dimensions to match the original input shape
955 output_dict = format_outputs(output_dict, loop_shape)
957 for var in self.outputs:
958 if var.compression is not None:
959 for field in var.compression.fields:
960 if field not in output_dict:
961 self.logger.warning(f"Model return missing field '{field}' for output variable '{var}'. "
962 f"Returning NaNs...")
963 output_dict[field].setdefault(field, np.full((N,), np.nan))
964 elif var.name not in output_dict:
965 self.logger.warning(f"Model return missing output variable '{var.name}'. Returning NaNs...")
966 output_dict[var.name] = np.full((N,), np.nan)
968 # Return the output dictionary and any errors
969 if errors:
970 output_dict['errors'] = errors
971 return output_dict
973 def predict(self, inputs: dict | Dataset,
974 use_model: Literal['best', 'worst'] | tuple = None,
975 model_dir: str | Path = None,
976 index_set: Literal['train', 'test'] | IndexSet = 'test',
977 misc_coeff: MiscTree = None,
978 incremental: bool = False,
979 executor: Executor = None,
980 **kwds) -> Dataset:
981 """Evaluate the MISC surrogate approximation at new inputs `x`.
983 !!! Note "Using the underlying model"
984 By default this will predict the MISC surrogate approximation; all inputs are assumed to be in a compressed
985 and normalized form. If the component does not have a surrogate (i.e. it is analytical), then the inputs
986 will be converted to model form and the underlying model will be called in place. If you instead want to
987 override the surrogate, passing `use_model` will call the underlying model directly. In that case, the
988 inputs should be passed in already in model form (i.e. full fields, denormalized).
990 :param inputs: `dict` of input arrays for each variable input
991 :param use_model: 'best'=high-fidelity, 'worst'=low-fidelity, tuple=a specific `alpha`, None=surrogate (default)
992 :param model_dir: directory to save output files if `use_model` is specified, ignored otherwise
993 :param index_set: the active index set, defaults to `self.active_set` if `'train'` or both
994 `self.active_set + self.candidate_set` if `'test'`
995 :param misc_coeff: the data structure holding the MISC coefficients to use, which defaults to the
996 training or testing coefficients depending on the `index_set` parameter.
997 :param incremental: a special flag to use if the provided `index_set` is an incremental update to the active
998 index set. A temporary copy of the internal `misc_coeff` data structure will be updated
999 and used to incorporate the new indices.
1000 :param executor: executor for parallel execution if the model is not vectorized (optional), will use the
1001 executor for looping over MISC coefficients if evaluating the surrogate rather than the model
1002 :param kwds: additional keyword arguments to pass to the model (if using the underlying model)
1003 :returns: the surrogate approximation of the model (or the model return itself if `use_model`)
1004 """
1005 # Use raw model inputs/outputs
1006 if use_model is not None:
1007 outputs = self.call_model(inputs, model_fidelity=use_model, output_path=model_dir, executor=executor,**kwds)
1008 ret = {}
1009 for var in self.outputs:
1010 if var in outputs:
1011 ret[var.name] = outputs[var.name]
1012 elif var.compression is not None:
1013 for field in var.compression.fields:
1014 ret[field] = outputs[field]
1015 return ret
1017 # Convert inputs/outputs to/from model if no surrogate (i.e. analytical models)
1018 if not self.has_surrogate:
1019 field_coords = {f'{var}_coords': self.model_kwargs.get(f'{var}_coords', kwds.get(f'{var}_coords', None))
1020 for var in self.inputs}
1021 inputs, field_coords = to_model_dataset(inputs, self.inputs, del_latent=True, **field_coords)
1022 field_coords.update(kwds)
1023 outputs = self.call_model(inputs, model_fidelity=use_model or 'best', output_path=model_dir,
1024 executor=executor, **field_coords)
1025 outputs, surr_vars = to_surrogate_dataset(outputs, self.outputs, del_fields=True, **field_coords)
1026 return {str(var): outputs[var] for var in surr_vars}
1028 # Choose the correct index set and misc_coeff data structures
1029 if incremental:
1030 misc_coeff = copy.deepcopy(self.misc_coeff_train)
1031 self.update_misc_coeff(index_set, self.active_set, misc_coeff)
1032 index_set = self.active_set.union(index_set)
1033 else:
1034 index_set, misc_coeff = self._match_index_set(index_set, misc_coeff)
1036 # Format inputs for surrogate prediction (all scalars at this point, including latent coeffs)
1037 inputs, loop_shape = format_inputs(inputs) # {'x': (N,)}
1038 outputs = {}
1040 # Handle prediction with empty active set (return nan)
1041 if len(index_set) == 0:
1042 for var in self.outputs:
1043 outputs[var.name] = np.full(loop_shape, np.nan)
1044 return outputs
1046 y_vars = self._surrogate_outputs() # Only request this component's specified outputs (ignore all extras)
1048 # Combination technique MISC surrogate prediction
1049 results = []
1050 coeffs = []
1051 for alpha, beta in index_set:
1052 comb_coeff = misc_coeff[alpha, beta]
1053 if np.abs(comb_coeff) > 0:
1054 coeffs.append(comb_coeff)
1055 args = (self.misc_states.get((alpha, beta)),
1056 self.training_data.get(alpha, beta[:len(self.data_fidelity)], skip_nan=True, y_vars=y_vars))
1058 results.append(self.interpolator.predict(inputs, *args) if executor is None else
1059 executor.submit(self.interpolator.predict, inputs, *args))
1061 if executor is not None:
1062 wait(results, timeout=None, return_when=ALL_COMPLETED)
1063 results = [future.result() for future in results]
1065 for coeff, interp_pred in zip(coeffs, results):
1066 for var, arr in interp_pred.items():
1067 if outputs.get(var) is None:
1068 outputs[str(var)] = coeff * arr
1069 else:
1070 outputs[str(var)] += coeff * arr
1072 return format_outputs(outputs, loop_shape)
1074 def update_misc_coeff(self, new_indices: IndexSet, index_set: Literal['test', 'train'] | IndexSet = 'train',
1075 misc_coeff: MiscTree = None):
1076 """Update MISC coefficients incrementally resulting from the addition of new indices to an index set.
1078 !!! Warning "Incremental updates"
1079 This function is used to update the MISC coefficients stored in `misc_coeff` after adding new indices
1080 to the given `index_set`. If a custom `index_set` or `misc_coeff` are provided, the user is responsible
1081 for ensuring the data structures are consistent. Since this is an incremental update, this means all
1082 existing coefficients for every index in `index_set` should be precomputed and stored in `misc_coeff`.
1084 :param new_indices: a set of $(\\alpha, \\beta)$ tuples that are being added to the `index_set`
1085 :param index_set: the active index set, defaults to `self.active_set` if `'train'` or both
1086 `self.active_set + self.candidate_set` if `'test'`
1087 :param misc_coeff: the data structure holding the MISC coefficients to update, which defaults to the
1088 training or testing coefficients depending on the `index_set` parameter. This data structure
1089 is modified in place.
1090 """
1091 index_set, misc_coeff = self._match_index_set(index_set, misc_coeff)
1093 for new_alpha, new_beta in new_indices:
1094 new_ind = np.array(new_alpha + new_beta)
1096 # Update all existing/new coefficients if they are a distance of [0, 1] "below" the new index
1097 # Note that new indices can only be [0, 1] away from themselves -- not any other new indices
1098 for old_alpha, old_beta in itertools.chain(index_set, [(new_alpha, new_beta)]):
1099 old_ind = np.array(old_alpha + old_beta)
1100 diff = new_ind - old_ind
1101 if np.all(np.isin(diff, [0, 1])):
1102 if misc_coeff.get((old_alpha, old_beta)) is None:
1103 misc_coeff[old_alpha, old_beta] = 0
1104 misc_coeff[old_alpha, old_beta] += (-1) ** int(np.sum(np.abs(diff)))
1106 def activate_index(self, alpha: MultiIndex, beta: MultiIndex, model_dir: str | Path = None,
1107 executor: Executor = None):
1108 """Add a multi-index to the active set and all neighbors to the candidate set.
1110 !!! Warning
1111 The user of this function is responsible for ensuring that the index set maintains downward-closedness.
1112 That is, only activate indices that are neighbors of the current active set.
1114 :param alpha: A multi-index specifying model fidelity
1115 :param beta: A multi-index specifying surrogate fidelity
1116 :param model_dir: Directory to save model output files
1117 :param executor: Executor for parallel execution of model on training data if the model is not vectorized
1118 """
1119 if (alpha, beta) in self.active_set:
1120 self.logger.warning(f'Multi-index {(alpha, beta)} is already in the active index set. Ignoring...')
1121 return
1122 if (alpha, beta) not in self.candidate_set and (sum(alpha) + sum(beta)) > 0:
1123 # Can only activate the initial index (0, 0, ... 0) without it being in the candidate set
1124 self.logger.warning(f'Multi-index {(alpha, beta)} is not a neighbor of the active index set, so it '
1125 f'cannot be activated. Please only add multi-indices from the candidate set. '
1126 f'Ignoring...')
1127 return
1129 # Collect all neighbor candidate indices
1130 neighbors = self._neighbors(alpha, beta, forward=True)
1131 indices = list(itertools.chain([(alpha, beta)] if (alpha, beta) not in self.candidate_set else [], neighbors))
1133 # Refine and collect all new model inputs (i.e. training points) requested by the new candidates
1134 alpha_list = [] # keep track of model fidelities
1135 design_list = [] # keep track of training data coordinates/locations/indices
1136 model_inputs = {} # concatenate all model inputs
1137 field_coords = {f'{var}_coords': self.model_kwargs.get(f'{var}_coords', None) for var in self.inputs}
1138 domains = self.inputs.get_domains()
1139 weight_fcns = self.inputs.get_pdfs()
1140 for a, b in indices:
1141 design_coords, design_pts = self.training_data.refine(a, b[:len(self.data_fidelity)],
1142 domains, weight_fcns)
1143 design_pts, fc = to_model_dataset(design_pts, self.inputs, del_latent=True, **field_coords)
1145 # Remove duplicate (alpha, coords) pairs -- so you don't evaluate the model twice for the same input
1146 i = 0
1147 del_idx = []
1148 for other_design in design_list:
1149 for other_coord in other_design:
1150 for j, curr_coord in enumerate(design_coords):
1151 if curr_coord == other_coord and a == alpha_list[i] and j not in del_idx:
1152 del_idx.append(j)
1153 i += 1
1154 design_coords = [design_coords[j] for j in range(len(design_coords)) if j not in del_idx]
1155 design_pts = {var: np.delete(arr, del_idx, axis=0) for var, arr in design_pts.items()}
1157 alpha_list.extend([tuple(a)] * len(design_coords))
1158 design_list.append(design_coords)
1159 field_coords.update(fc)
1160 for var in design_pts:
1161 model_inputs[var] = design_pts[var] if model_inputs.get(var) is None else (
1162 np.concatenate((model_inputs[var], design_pts[var]), axis=0))
1164 # Evaluate model at designed training points
1165 if len(alpha_list) > 0:
1166 self.logger.info(f"Running {len(alpha_list)} total model evaluations for component "
1167 f"'{self.name}' new candidate indices: {indices}...")
1168 model_outputs = self.call_model(model_inputs, model_fidelity=alpha_list, output_path=model_dir,
1169 executor=executor, **field_coords)
1170 errors = model_outputs.pop('errors', {})
1172 # Unpack model outputs and update states
1173 start_idx = 0
1174 for i, (a, b) in enumerate(indices):
1175 num_train_pts = len(design_list[i])
1176 end_idx = start_idx + num_train_pts # Ensure loop dim of 1 gets its own axis (might have been squeezed)
1177 yi_dict = {var: arr[np.newaxis, ...] if len(alpha_list) == 1 and arr.shape[0] != 1 else
1178 arr[start_idx:end_idx, ...] for var, arr in model_outputs.items()}
1180 # Check for errors and store
1181 err_coords = []
1182 err_list = []
1183 for idx in list(errors.keys()):
1184 if idx < end_idx:
1185 err_info = errors.pop(idx)
1186 err_info['index'] = idx - start_idx
1187 err_coords.append(design_list[i][idx - start_idx])
1188 err_list.append(err_info)
1189 if len(err_list) > 0:
1190 self.logger.warning(f"Model errors occurred while adding candidate ({a}, {b}) for component "
1191 f"{self.name}. Leaving NaN values in training data...")
1192 self.training_data.set_errors(a, b[:len(self.data_fidelity)], err_coords, err_list)
1194 # Compress field quantities and normalize
1195 yi_dict, y_vars = to_surrogate_dataset(yi_dict, self.outputs, del_fields=False, **field_coords)
1197 # Store training data, computational cost, and new interpolator state
1198 self.training_data.set(a, b[:len(self.data_fidelity)], design_list[i], yi_dict)
1199 self.training_data.impute_missing_data(a, b[:len(self.data_fidelity)])
1200 self.misc_costs[a, b] = self.model_costs.get(a, 1.) * num_train_pts
1201 self.misc_states[a, b] = self.interpolator.refine(b[len(self.data_fidelity):],
1202 self.training_data.get(a, b[:len(self.data_fidelity)],
1203 y_vars=y_vars, skip_nan=True),
1204 self.misc_states.get((alpha, beta)),
1205 domains)
1206 start_idx = end_idx
1208 # Move to the active index set
1209 s = set()
1210 s.add((alpha, beta))
1211 self.update_misc_coeff(IndexSet(s), index_set='train')
1212 if (alpha, beta) in self.candidate_set:
1213 self.candidate_set.remove((alpha, beta))
1214 else:
1215 # Only for initial index which didn't come from the candidate set
1216 self.update_misc_coeff(IndexSet(s), index_set='test')
1217 self.active_set.update(s)
1219 self.update_misc_coeff(neighbors, index_set='test') # neighbors will only ever pass through here once
1220 self.candidate_set.update(neighbors)
1222 def gradient(self, inputs: dict | Dataset,
1223 index_set: Literal['train', 'test'] | IndexSet = 'test',
1224 misc_coeff: MiscTree = None,
1225 derivative: Literal['first', 'second'] = 'first',
1226 executor: Executor = None) -> Dataset:
1227 """Evaluate the Jacobian or Hessian of the MISC surrogate approximation at new `inputs`, i.e.
1228 the first or second derivatives, respectively.
1230 :param inputs: `dict` of input arrays for each variable input
1231 :param index_set: the active index set, defaults to `self.active_set` if `'train'` or both
1232 `self.active_set + self.candidate_set` if `'test'`
1233 :param misc_coeff: the data structure holding the MISC coefficients to use, which defaults to the
1234 training or testing coefficients depending on the `index_set` parameter.
1235 :param derivative: whether to compute the first or second derivative (i.e. Jacobian or Hessian)
1236 :param executor: executor for looping over MISC coefficients (optional)
1237 :returns: a `dict` of the Jacobian or Hessian of the surrogate approximation for each output variable
1238 """
1239 if not self.has_surrogate:
1240 self.logger.warning("No surrogate model available for gradient computation.")
1241 return None
1243 index_set, misc_coeff = self._match_index_set(index_set, misc_coeff)
1244 inputs, loop_shape = format_inputs(inputs) # {'x': (N,)}
1245 outputs = {}
1247 if len(index_set) == 0:
1248 for var in self.outputs:
1249 outputs[var] = np.full(loop_shape, np.nan)
1250 return outputs
1251 y_vars = self._surrogate_outputs()
1253 # Combination technique MISC gradient prediction
1254 results = []
1255 coeffs = []
1256 for alpha, beta in index_set:
1257 comb_coeff = misc_coeff[alpha, beta]
1258 if np.abs(comb_coeff) > 0:
1259 coeffs.append(comb_coeff)
1260 func = self.interpolator.gradient if derivative == 'first' else self.interpolator.hessian
1261 args = (self.misc_states.get((alpha, beta)),
1262 self.training_data.get(alpha, beta[:len(self.data_fidelity)], skip_nan=True, y_vars=y_vars))
1264 results.append(func(inputs, *args) if executor is None else executor.submit(func, inputs, *args))
1266 if executor is not None:
1267 wait(results, timeout=None, return_when=ALL_COMPLETED)
1268 results = [future.result() for future in results]
1270 for coeff, interp_pred in zip(coeffs, results):
1271 for var, arr in interp_pred.items():
1272 if outputs.get(var) is None:
1273 outputs[str(var)] = coeff * arr
1274 else:
1275 outputs[str(var)] += coeff * arr
1277 return format_outputs(outputs, loop_shape)
1279 def hessian(self, *args, **kwargs):
1280 """Alias for `Component.gradient(*args, derivative='second', **kwargs)`."""
1281 return self.gradient(*args, derivative='second', **kwargs)
1283 def model_kwarg_requested(self, kwarg_name: str) -> bool:
1284 """Return whether the underlying component model requested this `kwarg_name`. Special kwargs include:
1286 - `output_path` — a save directory created by `amisc` will be passed to the model for saving model output files.
1287 - `alpha` — a tuple or list of model fidelity indices will be passed to the model to adjust fidelity.
1288 - `input_vars` — a list of `Variable` objects will be passed to the model for input variable information.
1289 - `output_vars` — a list of `Variable` objects will be passed to the model for output variable information.
1291 :param kwarg_name: the argument to check for in the underlying component model's function signature kwargs
1292 :returns: whether the component model requests this `kwarg` argument
1293 """
1294 signature = inspect.signature(self.model)
1295 for param in signature.parameters.values():
1296 if param.name == kwarg_name and param.default != param.empty:
1297 return True
1298 return False
1300 def set_logger(self, log_file: str | Path = None, stdout: bool = None, logger: logging.Logger = None,
1301 level: int = logging.INFO):
1302 """Set a new `logging.Logger` object.
1304 :param log_file: log to file (if provided)
1305 :param stdout: whether to connect the logger to console (defaults to whatever is currently set or False)
1306 :param logger: the logging object to use (if None, then a new logger is created; this will override
1307 the `log_file` and `stdout` arguments if set)
1308 :param level: the logging level to set (default is `logging.INFO`)
1309 """
1310 if stdout is None:
1311 stdout = False
1312 if self._logger is not None:
1313 for handler in self._logger.handlers:
1314 if isinstance(handler, logging.StreamHandler):
1315 stdout = True
1316 break
1317 self._logger = logger or get_logger(self.name, log_file=log_file, stdout=stdout, level=level)
1319 def update_model(self, new_model: callable = None, model_kwargs: dict = None, **kwargs):
1320 """Update the underlying component model or its kwargs."""
1321 if new_model is not None:
1322 self.model = new_model
1323 new_kwargs = self.model_kwargs.data
1324 new_kwargs.update(model_kwargs or {})
1325 new_kwargs.update(kwargs)
1326 self.model_kwargs = new_kwargs
1328 def get_cost(self, alpha: MultiIndex, beta: MultiIndex) -> float:
1329 """Return the total cost (wall time s) required to add $(\\alpha, \\beta)$ to the MISC approximation.
1331 :param alpha: A multi-index specifying model fidelity
1332 :param beta: A multi-index specifying surrogate fidelity
1333 :returns: the total cost of adding this multi-index pair to the MISC approximation
1334 """
1335 try:
1336 return self.misc_costs[alpha, beta]
1337 except Exception:
1338 return 0.0
1340 @staticmethod
1341 def is_downward_closed(indices: IndexSet) -> bool:
1342 """Return if a list of $(\\alpha, \\beta)$ multi-indices is downward-closed.
1344 MISC approximations require a downward-closed set in order to use the combination-technique formula for the
1345 coefficients (as implemented by `Component.update_misc_coeff()`).
1347 !!! Example
1348 The list `[( (0,), (0,) ), ( (1,), (0,) ), ( (1,), (1,) )]` is downward-closed. You can visualize this as
1349 building a stack of cubes: in order to place a cube, all adjacent cubes must be present (does the logo
1350 make sense now?).
1352 :param indices: `IndexSet` of (`alpha`, `beta`) multi-indices
1353 :returns: whether the set of indices is downward-closed
1354 """
1355 # Iterate over every multi-index
1356 for alpha, beta in indices:
1357 # Every smaller multi-index must also be included in the indices list
1358 sub_sets = [np.arange(tuple(alpha + beta)[i] + 1) for i in range(len(alpha) + len(beta))]
1359 for ele in itertools.product(*sub_sets):
1360 tup = (MultiIndex(ele[:len(alpha)]), MultiIndex(ele[len(alpha):]))
1361 if tup not in indices:
1362 return False
1363 return True
1365 def clear(self):
1366 """Clear the component of all training data, index sets, and MISC states."""
1367 self.active_set.clear()
1368 self.candidate_set.clear()
1369 self.misc_states.clear()
1370 self.misc_costs.clear()
1371 self.misc_coeff_train.clear()
1372 self.misc_coeff_test.clear()
1373 self.model_costs.clear()
1374 self.training_data.clear()
1376 def serialize(self, keep_yaml_objects: bool = False, serialize_args: dict[str, tuple] = None,
1377 serialize_kwargs: dict[str: dict] = None) -> dict:
1378 """Convert to a `dict` with only standard Python types as fields and values.
1380 :param keep_yaml_objects: whether to keep `Variable` or other yaml serializable objects instead of
1381 also serializing them (default is False)
1382 :param serialize_args: additional arguments to pass to the `serialize` method of each `Component` attribute;
1383 specify as a `dict` of attribute names to tuple of arguments to pass
1384 :param serialize_kwargs: additional keyword arguments to pass to the `serialize` method of each
1385 `Component` attribute
1386 :returns: a `dict` representation of the `Component` object
1387 """
1388 serialize_args = serialize_args or dict()
1389 serialize_kwargs = serialize_kwargs or dict()
1390 d = {}
1391 for key, value in self.__dict__.items():
1392 if value is not None and not key.startswith('_'):
1393 if key == 'serializers':
1394 # Update the serializers
1395 serializers = self._validate_serializers({k: type(getattr(self, k)) for k in value.keys()})
1396 d[key] = {k: (v.obj if keep_yaml_objects else v.serialize()) for k, v in serializers.items()}
1397 elif key in ['inputs', 'outputs'] and not keep_yaml_objects:
1398 d[key] = value.serialize(**serialize_kwargs.get(key, {}))
1399 elif key == 'model' and not keep_yaml_objects:
1400 d[key] = YamlSerializable(obj=value).serialize()
1401 elif key in ['data_fidelity', 'surrogate_fidelity', 'model_fidelity']:
1402 if len(value) > 0:
1403 d[key] = str(value)
1404 elif key in ['active_set', 'candidate_set']:
1405 if len(value) > 0:
1406 d[key] = value.serialize()
1407 elif key in ['misc_costs', 'misc_coeff_train', 'misc_coeff_test', 'misc_states']:
1408 if len(value) > 0:
1409 d[key] = value.serialize(keep_yaml_objects=keep_yaml_objects)
1410 elif key in ['model_costs']:
1411 if len(value) > 0:
1412 d[key] = {str(k): float(v) for k, v in value.items()}
1413 elif key in ComponentSerializers.__annotations__.keys():
1414 if key in ['training_data', 'interpolator'] and not self.has_surrogate:
1415 continue
1416 else:
1417 d[key] = value.serialize(*serialize_args.get(key, ()), **serialize_kwargs.get(key, {}))
1418 else:
1419 d[key] = value
1420 return d
1422 @classmethod
1423 def deserialize(cls, serialized_data: dict, search_paths: list[str | Path] = None,
1424 search_keys: list[str] = None) -> Component:
1425 """Return a `Component` from `data`. Let pydantic handle field validation and conversion. If any component
1426 data has been saved to file and the save file doesn't exist, then the loader will search for the file
1427 in the current working directory and any additional search paths provided.
1429 :param serialized_data: the serialized data to construct the object from
1430 :param search_paths: paths to try and find any save files (i.e. if they moved since they were serialized),
1431 will always search in the current working directory by default
1432 :param search_keys: keys to search for save files in each component (default is all keys in
1433 [`ComponentSerializers`][amisc.component.ComponentSerializers], in addition to variable
1434 inputs and outputs)
1435 """
1436 if isinstance(serialized_data, Component):
1437 return serialized_data
1438 elif callable(serialized_data):
1439 # try to construct a component from a raw function (assume data fidelity is (2,) for each inspected input)
1440 return cls(serialized_data, data_fidelity=(2,) * len(_inspect_function(serialized_data)[0]))
1442 search_paths = search_paths or []
1443 search_keys = search_keys or []
1444 search_keys.extend(ComponentSerializers.__annotations__.keys())
1445 comp = serialized_data
1447 for key in search_keys:
1448 if (filename := comp.get(key, None)) is not None:
1449 comp[key] = search_for_file(filename, search_paths=search_paths) # will ret original str if not found
1451 for key in ['inputs', 'outputs']:
1452 for var in comp.get(key, []):
1453 if isinstance(var, dict):
1454 if (compression := var.get('compression', None)) is not None:
1455 var['compression'] = search_for_file(compression, search_paths=search_paths)
1457 return cls(**comp)
1459 @staticmethod
1460 def _yaml_representer(dumper: yaml.Dumper, comp: Component) -> yaml.MappingNode:
1461 """Convert a single `Component` object (`data`) to a yaml MappingNode (i.e. a `dict`)."""
1462 save_path, save_file = _get_yaml_path(dumper)
1463 serialize_kwargs = {}
1464 for key, serializer in comp.serializers.items():
1465 if issubclass(serializer.obj, PickleSerializable):
1466 filename = save_path / f'{save_file}_{comp.name}_{key}.pkl'
1467 serialize_kwargs[key] = {'save_path': save_path / filename}
1468 return dumper.represent_mapping(Component.yaml_tag, comp.serialize(serialize_kwargs=serialize_kwargs,
1469 keep_yaml_objects=True))
1471 @staticmethod
1472 def _yaml_constructor(loader: yaml.Loader, node):
1473 """Convert the `!Component` tag in yaml to a `Component` object."""
1474 # Add a file search path in the same directory as the yaml file being loaded from
1475 save_path, save_file = _get_yaml_path(loader)
1476 if isinstance(node, yaml.SequenceNode):
1477 return [ele if isinstance(ele, Component) else Component.deserialize(ele, search_paths=[save_path])
1478 for ele in loader.construct_sequence(node, deep=True)]
1479 elif isinstance(node, yaml.MappingNode):
1480 return Component.deserialize(loader.construct_mapping(node, deep=True), search_paths=[save_path])
1481 else:
1482 raise NotImplementedError(f'The "{Component.yaml_tag}" yaml tag can only be used on a yaml sequence or '
1483 f'mapping, not a "{type(node)}".')