Coverage for src/amisc/system.py: 87%
943 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-11-05 19:26 +0000
1"""The `System` object is a framework for multidisciplinary models. It manages multiple single discipline component
2models and the connections between them. It provides a top-level interface for constructing and evaluating surrogates.
4Features:
6- Manages multidisciplinary models in a graph data structure, supports feedforward and feedback connections
7- Feedback connections are solved with a fixed-point iteration (FPI) nonlinear solver with anderson acceleration
8- Top-level interface for training and using surrogates of each component model
9- Adaptive experimental design for choosing training data efficiently
10- Convenient testing, plotting, and performance metrics provided to assess quality of surrogates
11- Detailed logging and traceback information
12- Supports parallel or vectorized execution of component models
13- Abstract and flexible interfacing with component models
14- Easy serialization and deserialization to/from YAML files
15- Supports approximating field quantities via compression
17Includes:
19- `TrainHistory` — a history of training iterations for the system surrogate
20- `System` — the top-level object for managing multidisciplinary models
21"""
22# ruff: noqa: E702
23from __future__ import annotations
25import copy
26import datetime
27import functools
28import logging
29import os
30import pickle
31import random
32import string
33import time
34from collections import ChainMap, UserList, deque
35from concurrent.futures import ALL_COMPLETED, Executor, wait
36from datetime import timezone
37from pathlib import Path
38from typing import Annotated, Callable, ClassVar, Literal, Optional
40import matplotlib.pyplot as plt
41import networkx as nx
42import numpy as np
43import yaml
44from pydantic import BaseModel, ConfigDict, Field, field_validator
46from amisc.component import Component, IndexSet, MiscTree
47from amisc.serialize import Serializable, _builtin
48from amisc.typing import LATENT_STR_ID, Dataset, MultiIndex, TrainIteration
49from amisc.utils import (
50 _combine_latent_arrays,
51 constrained_lls,
52 format_inputs,
53 format_outputs,
54 get_logger,
55 relative_error,
56 to_model_dataset,
57 to_surrogate_dataset,
58)
59from amisc.variable import VariableList
61__all__ = ['TrainHistory', 'System']
64class TrainHistory(UserList, Serializable):
65 """Stores the training history of a system surrogate as a list of `TrainIteration` objects."""
67 def __init__(self, data: list = None):
68 data = data or []
69 super().__init__(self._validate_data(data))
71 def serialize(self) -> list[dict]:
72 """Return a list of each result in the history serialized to a `dict`."""
73 ret_list = []
74 for res in self:
75 new_res = res.copy()
76 new_res['alpha'] = str(res['alpha'])
77 new_res['beta'] = str(res['beta'])
78 ret_list.append(new_res)
79 return ret_list
81 @classmethod
82 def deserialize(cls, serialized_data: list[dict]) -> TrainHistory:
83 """Deserialize a list of `dict` objects into a `TrainHistory` object."""
84 return TrainHistory(serialized_data)
86 @classmethod
87 def _validate_data(cls, data: list[dict]) -> list[TrainIteration]:
88 return [cls._validate_item(item) for item in data]
90 @classmethod
91 def _validate_item(cls, item: dict):
92 """Format a `TrainIteration` `dict` item before appending to the history."""
93 item.setdefault('test_error', None)
94 item['alpha'] = MultiIndex(item['alpha'])
95 item['beta'] = MultiIndex(item['beta'])
96 item['num_evals'] = int(item['num_evals'])
97 item['added_cost'] = float(item['added_cost'])
98 item['added_error'] = float(item['added_error'])
99 return item
101 def append(self, item: dict):
102 super().append(self._validate_item(item))
104 def __add__(self, other):
105 other_list = other.data if isinstance(other, TrainHistory) else other
106 return TrainHistory(data=self.data + other_list)
108 def extend(self, items):
109 super().extend([self._validate_item(item) for item in items])
111 def insert(self, index, item):
112 super().insert(index, self._validate_item(item))
114 def __setitem__(self, key, value):
115 super().__setitem__(key, self._validate_item(value))
117 def __eq__(self, other):
118 """Two `TrainHistory` objects are equal if they have the same length and all items are equal, excluding nans."""
119 if not isinstance(other, TrainHistory):
120 return False
121 if len(self) != len(other):
122 return False
123 for item_self, item_other in zip(self, other):
124 for key in item_self:
125 if key in item_other:
126 val_self = item_self[key]
127 val_other = item_other[key]
128 if isinstance(val_self, float) and isinstance(val_other, float):
129 if not (np.isnan(val_self) and np.isnan(val_other)) and val_self != val_other:
130 return False
131 elif isinstance(val_self, dict) and isinstance(val_other, dict):
132 for v, err in val_self.items():
133 if v in val_other:
134 err_other = val_other[v]
135 if not (np.isnan(err) and np.isnan(err_other)) and err != err_other:
136 return False
137 else:
138 return False
139 elif val_self != val_other:
140 return False
141 else:
142 return False
143 return True
146class _Converged:
147 """Helper class to track which samples have converged during `System.predict()`."""
148 def __init__(self, num_samples):
149 self.num_samples = num_samples
150 self.valid_idx = np.full(num_samples, True) # All samples are valid by default
151 self.converged_idx = np.full(num_samples, False) # For FPI convergence
153 def reset_convergence(self):
154 self.converged_idx = np.full(self.num_samples, False)
156 @property
157 def curr_idx(self):
158 return np.logical_and(self.valid_idx, ~self.converged_idx)
161def _merge_shapes(target_shape, arr):
162 """Helper to merge an array into the target shape."""
163 shape1, shape2 = target_shape, arr.shape
164 if len(shape2) > len(shape1):
165 shape1, shape2 = shape2, shape1
166 result = []
167 for i in range(len(shape1)):
168 if i < len(shape2):
169 if shape1[i] == 1:
170 result.append(shape2[i])
171 elif shape2[i] == 1:
172 result.append(shape1[i])
173 else:
174 result.append(shape1[i])
175 else:
176 result.append(1)
177 arr = arr.reshape(tuple(result))
178 return np.broadcast_to(arr, target_shape).copy()
181class System(BaseModel, Serializable):
182 """
183 Multidisciplinary (MD) surrogate framework top-level class. Construct a `System` from a list of
184 `Component` models.
186 !!! Example
187 ```python
188 def f1(x):
189 y = x ** 2
190 return y
191 def f2(y):
192 z = y + 1
193 return z
195 system = System(f1, f2)
196 ```
198 A `System` object can saved/loaded from `.yml` files using the `!System` yaml tag.
200 :ivar name: the name of the system
201 :ivar components: list of `Component` models that make up the MD system
202 :ivar train_history: history of training iterations for the system surrogate (filled in during training)
204 :ivar _root_dir: root directory where all surrogate build products are saved to file
205 :ivar _logger: logger object for the system
206 """
207 yaml_tag: ClassVar[str] = u'!System'
208 model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True, validate_default=True,
209 extra='allow')
211 name: Annotated[str, Field(default_factory=lambda: "System_" + "".join(random.choices(string.digits, k=3)))]
212 components: Callable | Component | list[Callable | Component]
213 train_history: list[dict] | TrainHistory = TrainHistory()
214 amisc_version: str = None
216 _root_dir: Optional[str]
217 _logger: Optional[logging.Logger] = None
219 def __init__(self, /, *args, components=None, root_dir=None, **kwargs):
220 """Construct a `System` object from a list of `Component` models in `*args` or `components`. If
221 a `root_dir` is provided, then a new directory will be created under `root_dir` with the name
222 `amisc_{timestamp}`. This directory will be used to save all build products and log files.
224 :param components: list of `Component` models that make up the MD system
225 :param root_dir: root directory where all surrogate build products are saved to file (optional)
226 """
227 if components is None:
228 components = []
229 for a in args:
230 if isinstance(a, Component) or callable(a):
231 components.append(a)
232 else:
233 try:
234 components.extend(a)
235 except TypeError as e:
236 raise ValueError(f"Invalid component: {a}") from e
238 import amisc
239 amisc_version = kwargs.pop('amisc_version', amisc.__version__)
240 super().__init__(components=components, amisc_version=amisc_version, **kwargs)
241 self.root_dir = root_dir
243 def __repr__(self):
244 s = f'---- {self.name} ----\n'
245 s += f'amisc version: {self.amisc_version}\n'
246 s += f'Refinement level: {self.refine_level}\n'
247 s += f'Components: {", ".join([comp.name for comp in self.components])}\n'
248 s += f'Inputs: {", ".join([var.name for var in self.inputs()])}\n'
249 s += f'Outputs: {", ".join([var.name for var in self.outputs()])}'
250 return s
252 def __str__(self):
253 return self.__repr__()
255 @field_validator('components')
256 @classmethod
257 def _validate_components(cls, comps) -> list[Component]:
258 if not isinstance(comps, list):
259 comps = [comps]
260 comps = [Component.deserialize(c) for c in comps]
262 # Merge all variables to avoid name conflicts
263 merged_vars = VariableList.merge(*[comp.inputs for comp in comps], *[comp.outputs for comp in comps])
264 for comp in comps:
265 comp.inputs.update({var.name: var for var in merged_vars.values() if var in comp.inputs})
266 comp.outputs.update({var.name: var for var in merged_vars.values() if var in comp.outputs})
268 return comps
270 @field_validator('train_history')
271 @classmethod
272 def _validate_train_history(cls, history) -> TrainHistory:
273 if isinstance(history, TrainHistory):
274 return history
275 else:
276 return TrainHistory.deserialize(history)
278 def graph(self) -> nx.DiGraph:
279 """Build a directed graph of the system components based on their input-output relationships."""
280 graph = nx.DiGraph()
281 model_deps = {}
282 for comp in self.components:
283 graph.add_node(comp.name)
284 for output in comp.outputs:
285 model_deps[output] = comp.name
286 for comp in self.components:
287 for in_var in comp.inputs:
288 if in_var in model_deps:
289 graph.add_edge(model_deps[in_var], comp.name)
291 return graph
293 def _save_on_error(func):
294 """Gracefully exit and save the `System` object on any errors."""
295 @functools.wraps(func)
296 def wrap(self, *args, **kwargs):
297 try:
298 return func(self, *args, **kwargs)
299 except:
300 if self.root_dir is not None:
301 self.save_to_file(f'{self.name}_error.yml')
302 self.logger.critical(f'An error occurred during execution of "{func.__name__}". Saving '
303 f'System object to {self.name}_error.yml', exc_info=True)
304 self.logger.info(f'Final system surrogate on exit: \n {self}')
305 raise
306 return wrap
307 _save_on_error = staticmethod(_save_on_error)
309 def insert_components(self, components: list | Callable | Component):
310 """Insert new components into the system."""
311 components = components if isinstance(components, list) else [components]
312 self.components = self.components + components
314 def swap_component(self, old_component: str | Component, new_component: Callable | Component):
315 """Replace an old component with a new component."""
316 old_name = old_component if isinstance(old_component, str) else old_component.name
317 comps = [comp if comp.name != old_name else new_component for comp in self.components]
318 self.components = comps
320 def remove_component(self, component: str | Component):
321 """Remove a component from the system."""
322 comp_name = component if isinstance(component, str) else component.name
323 self.components = [comp for comp in self.components if comp.name != comp_name]
325 def inputs(self) -> VariableList:
326 """Collect all inputs from each component in the `System` and combine them into a
327 single [`VariableList`][amisc.variable.VariableList] object, excluding variables that are also outputs of
328 any component.
330 :returns: A [`VariableList`][amisc.variable.VariableList] containing all inputs from the components.
331 """
332 all_inputs = ChainMap(*[comp.inputs for comp in self.components])
333 return VariableList({k: all_inputs[k] for k in all_inputs.keys() - self.outputs().keys()})
335 def outputs(self) -> VariableList:
336 """Collect all outputs from each component in the `System` and combine them into a
337 single [`VariableList`][amisc.variable.VariableList] object.
339 :returns: A [`VariableList`][amisc.variable.VariableList] containing all outputs from the components.
340 """
341 return VariableList({k: v for k, v in ChainMap(*[comp.outputs for comp in self.components]).items()})
343 def coupling_variables(self) -> VariableList:
344 """Collect all coupling variables from each component in the `System` and combine them into a
345 single [`VariableList`][amisc.variable.VariableList] object.
347 :returns: A [`VariableList`][amisc.variable.VariableList] containing all coupling variables from the components.
348 """
349 all_outputs = self.outputs()
350 return VariableList({k: all_outputs[k] for k in (all_outputs.keys() &
351 ChainMap(*[comp.inputs for comp in self.components]).keys())})
353 def variables(self):
354 """Iterator over all variables in the system (inputs and outputs)."""
355 yield from ChainMap(self.inputs(), self.outputs()).values()
357 @property
358 def refine_level(self) -> int:
359 """The total number of training iterations."""
360 return len(self.train_history)
362 @property
363 def logger(self) -> logging.Logger:
364 return self._logger
366 @logger.setter
367 def logger(self, logger: logging.Logger):
368 self._logger = logger
369 for comp in self.components:
370 comp.logger = logger
372 @staticmethod
373 def timestamp() -> str:
374 """Return a UTC timestamp string in the isoformat `YYYY-MM-DDTHH.MM.SS`."""
375 return datetime.datetime.now(tz=timezone.utc).isoformat().split('.')[0].replace(':', '.')
377 @property
378 def root_dir(self):
379 """Return the root directory of the surrogate (if available), otherwise `None`."""
380 return Path(self._root_dir) if self._root_dir is not None else None
382 @root_dir.setter
383 def root_dir(self, root_dir: str | Path):
384 """Set the root directory for all build products. If `root_dir` is `None`, then no products will be saved.
385 Otherwise, log files, model outputs, surrogate files, etc. will be saved under this directory.
387 !!! Note "`amisc` root directory"
388 If `root_dir` is not `None`, then a new directory will be created under `root_dir` with the name
389 `amisc_{timestamp}`. This directory will be used to save all build products. If `root_dir` matches the
390 `amisc_*` format, then it will be used directly.
392 :param root_dir: the root directory for all build products
393 """
394 if root_dir is not None:
395 parts = Path(root_dir).resolve().parts
396 if parts[-1].startswith('amisc_'):
397 self._root_dir = Path(root_dir).resolve().as_posix()
398 if not self.root_dir.is_dir():
399 os.mkdir(self.root_dir)
400 else:
401 root_dir = Path(root_dir) / ('amisc_' + self.timestamp())
402 os.mkdir(root_dir)
403 self._root_dir = Path(root_dir).resolve().as_posix()
405 log_file = None
406 if not (pth := self.root_dir / 'surrogates').is_dir():
407 os.mkdir(pth)
408 if not (pth := self.root_dir / 'components').is_dir():
409 os.mkdir(pth)
410 for comp in self.components:
411 if comp.model_kwarg_requested('output_path'):
412 if not (comp_pth := pth / comp.name).is_dir():
413 os.mkdir(comp_pth)
414 for f in os.listdir(self.root_dir):
415 if f.endswith('.log'):
416 log_file = (self.root_dir / f).resolve().as_posix()
417 break
418 if log_file is None:
419 log_file = (self.root_dir / f'amisc_{self.timestamp()}.log').resolve().as_posix()
420 self.set_logger(log_file=log_file)
422 else:
423 self._root_dir = None
424 self.set_logger(log_file=None)
426 def set_logger(self, log_file: str | Path | bool = None, stdout: bool = None, logger: logging.Logger = None,
427 level: int = logging.INFO):
428 """Set a new `logging.Logger` object.
430 :param log_file: log to this file if str or Path (defaults to whatever is currently set or empty);
431 set `False` to remove file logging or set `True` to create a default log file in the root dir
432 :param stdout: whether to connect the logger to console (defaults to whatever is currently set or `False`)
433 :param logger: the logging object to use (this will override the `log_file` and `stdout` arguments if set);
434 if `None`, then a new logger is created according to `log_file` and `stdout`
435 :param level: the logging level to set the logger to (defaults to `logging.INFO`)
436 """
437 # Decide whether to use stdout
438 if stdout is None:
439 stdout = False
440 if self._logger is not None:
441 for handler in self._logger.handlers:
442 if isinstance(handler, logging.StreamHandler):
443 stdout = True
444 break
446 # Decide what log_file to use (if any)
447 if log_file is True:
448 log_file = pth / f'amisc_{self.timestamp()}.log' if (pth := self.root_dir) is not None else (
449 f'amisc_{self.timestamp()}.log')
450 elif log_file is None:
451 if self._logger is not None:
452 for handler in self._logger.handlers:
453 if isinstance(handler, logging.FileHandler):
454 log_file = handler.baseFilename
455 break
456 elif log_file is False:
457 log_file = None
459 self._logger = logger or get_logger(self.name, log_file=log_file, stdout=stdout, level=level)
461 for comp in self.components:
462 comp.set_logger(log_file=log_file, stdout=stdout, logger=logger, level=level)
464 def sample_inputs(self, size: tuple | int, comp: str = 'System', use_pdf: bool = False,
465 nominal: dict[str: float] = None, constants: set[str] = None) -> Dataset:
466 """Return samples of the inputs according to provided options. Will return samples in the
467 normalized/compressed space of the surrogate. See [`to_model_dataset`][amisc.utils.to_model_dataset] to
468 convert the samples to be usable by the true model directly.
470 :param size: tuple or integer specifying shape or number of samples to obtain
471 :param comp: which component to sample inputs for (defaults to full system exogenous inputs)
472 :param use_pdf: whether to sample from each variable's pdf, defaults to random samples over input domain instead
473 :param nominal: `dict(var_id=value)` of nominal values for params with relative uncertainty, also can use
474 to specify constant values for a variable listed in `constants`. Specify nominal values
475 as the true model would expect them (i.e. not normalized or compressed) -- they will be
476 returned in normalized form.
477 :param constants: set of param types to hold constant while sampling (i.e. calibration, design, etc.),
478 can also put a `var.name` string in here to specify a single variable to hold constant
479 :returns: `dict` of `(*size,)` samples for each input variable
480 """
481 size = (size, ) if isinstance(size, int) else size
482 nominal = nominal or dict()
483 constants = constants or set()
484 inputs = self.inputs() if comp == 'System' else self[comp].inputs
485 samples = {}
486 for var in inputs:
487 # Sample from latent variable domains for field quantities
488 if var.compression is not None:
489 if var.category in constants or var in constants:
490 nom = nominal.get(var, None)
491 if nom is not None:
492 nom = np.broadcast_to(np.atleast_1d(nom), size + (var.compression.latent_size(),)).copy()
493 for i, d in enumerate(var.get_domain()):
494 samples[f'{var.name}{LATENT_STR_ID}{i}'] = nom[..., i] if nom is not None else np.mean(d)
495 else:
496 latent = var.sample_domain(size)
497 for i in range(latent.shape[-1]):
498 samples[f'{var.name}{LATENT_STR_ID}{i}'] = latent[..., i]
500 # Sample scalars normally
501 else:
502 # Set a constant value for this variable
503 if var.category in constants or var in constants:
504 if nom := nominal.get(var):
505 samples[var.name] = var.normalize(nom)
506 else:
507 samples[var.name] = var.normalize(var.get_nominal())
509 # Sample from this variable's pdf or randomly within its domain bounds (reject if outside bounds)
510 else:
511 lb, ub = var.get_domain()
512 x_sample = var.sample(size, nominal=nominal.get(var, None)) if use_pdf else var.sample_domain(size)
513 good_idx = (x_sample < ub) & (x_sample > lb)
514 num_reject = np.sum(~good_idx)
516 while num_reject > 0:
517 new_sample = var.sample((num_reject,), nominal=nominal.get(var, None)) \
518 if use_pdf else var.sample_domain((num_reject,))
519 x_sample[~good_idx] = new_sample
520 good_idx = (x_sample < ub) & (x_sample > lb)
521 num_reject = np.sum(~good_idx)
523 samples[var.name] = var.normalize(x_sample)
525 return samples
527 def simulate_fit(self):
528 """Loop back through training history and simulate each iteration. Will yield the internal data structures
529 of each `Component` surrogate after each iteration of training (without needing to call `fit()` or any
530 of the underlying models). This might be useful, for example, for computing the surrogate predictions on
531 a new test set or viewing cumulative training costs.
533 !!! Example
534 Say you have a new test set: `(new_xtest, new_ytest)`, and you want to compute the accuracy of the
535 surrogate fit at each iteration of the training history:
537 ```python
538 for train_iter, active_sets, candidate_sets, misc_coeff_train, misc_coeff_test in system.simulate_fit():
539 # Do something with the surrogate data structures
540 new_ysurr = system.predict(new_xtest, index_set=active_sets, misc_coeff=misc_coeff_train)
541 train_error = relative_error(new_ysurr, new_ytest)
542 ```
544 :return: a generator of the active index sets, candidate index sets, and MISC coefficients
545 of each component model at each iteration of the training history
546 """
547 # "Simulated" data structures for each component
548 active_sets = {comp.name: IndexSet() for comp in self.components} # active index sets for each component
549 candidate_sets = {comp.name: IndexSet() for comp in self.components} # candidate sets for each component
550 misc_coeff_train = {comp.name: MiscTree() for comp in self.components} # MISC coeff for active sets
551 misc_coeff_test = {comp.name: MiscTree() for comp in self.components} # MISC coeff for active + candidate sets
553 for train_result in self.train_history:
554 # The selected refinement component and indices
555 comp_star = train_result['component']
556 alpha_star = train_result['alpha']
557 beta_star = train_result['beta']
558 comp = self[comp_star]
560 # Get forward neighbors for the selected index
561 neighbors = comp._neighbors(alpha_star, beta_star, active_set=active_sets[comp_star], forward=True)
563 # "Activate" the index in the simulated data structure
564 s = set()
565 s.add((alpha_star, beta_star))
566 comp.update_misc_coeff(IndexSet(s), index_set=active_sets[comp_star],
567 misc_coeff=misc_coeff_train[comp_star])
569 if (alpha_star, beta_star) in candidate_sets[comp_star]:
570 candidate_sets[comp_star].remove((alpha_star, beta_star))
571 else:
572 # Only for initial index which didn't come from the candidate set
573 comp.update_misc_coeff(IndexSet(s), index_set=active_sets[comp_star].union(candidate_sets[comp_star]),
574 misc_coeff=misc_coeff_test[comp_star])
575 active_sets[comp_star].update(s)
577 comp.update_misc_coeff(neighbors, index_set=active_sets[comp_star].union(candidate_sets[comp_star]),
578 misc_coeff=misc_coeff_test[comp_star]) # neighbors will only ever pass here once
579 candidate_sets[comp_star].update(neighbors)
581 # Caller can now do whatever they want as if the system surrogate were at this training iteration
582 # See the "index_set" and "misc_coeff" overrides for `System.predict()` for example
583 yield train_result, active_sets, candidate_sets, misc_coeff_train, misc_coeff_test
586 def add_output(self):
587 """Add an output variable retroactively to a component surrogate. User should provide a callable that
588 takes a save path and extracts the model output data for given training point/location.
589 """
590 # TODO
591 # Loop back through the surrogate training history
592 # Simulate activate_index and extract the model output from file rather than calling the model
593 # Update all interpolator states
594 raise NotImplementedError
596 @_save_on_error
597 def fit(self, targets: list = None,
598 num_refine: int = 100,
599 max_iter: int = 20,
600 max_tol: float = 1e-3,
601 runtime_hr: float = 1.,
602 save_interval: int = 0,
603 estimate_bounds: bool = False,
604 update_bounds: bool = True,
605 test_set: tuple | str | Path = None,
606 start_test_check: int = None,
607 plot_interval: int = 1,
608 executor: Executor = None):
609 """Train the system surrogate adaptively by iterative refinement until an end condition is met.
611 :param targets: list of system output variables to focus refinement on, use all outputs if not specified
612 :param num_refine: number of input samples to compute error indicators on
613 :param max_iter: the maximum number of refinement steps to take
614 :param max_tol: the max allowable value in relative L2 error to achieve
615 :param runtime_hr: the threshold wall clock time (hr) at which to stop further refinement (will go
616 until all models finish the current iteration)
617 :param save_interval: number of refinement steps between each progress save, none if 0; `System.root_dir`
618 must be specified to save to file
619 :param estimate_bounds: whether to estimate bounds for the coupling variables; will only try to estimate from
620 the `test_set` if provided (defaults to `True`). Otherwise, you should manually
621 provide domains for all coupling variables.
622 :param update_bounds: whether to continuously update coupling variable bounds during refinement
623 :param test_set: `tuple` of `(xtest, ytest)` to show convergence of surrogate to the true model. The test set
624 inputs and outputs are specified as `dicts` of `np.ndarrays` with keys corresponding to the
625 variable names. Can also pass a path to a `.pkl` file that has the test set data as
626 {'test_set': (xtest, ytest)}.
627 :param start_test_check: the iteration to start checking the test set error (defaults to twice the number
628 of components); surrogate evaluation isn't useful during initialization so you
629 should at least allow one iteration per component before checking test set error
630 :param plot_interval: how often to plot the error indicator and test set error (defaults to every iteration);
631 will only plot and save to file if a root directory is set
632 :param executor: a `concurrent.futures.Executor` object to parallelize model evaluations (optional, but
633 recommended for expensive models)
634 """
635 start_test_check = start_test_check or 2 * len(self.components)
636 targets = targets or self.outputs()
637 xtest, ytest = self._get_test_set(test_set)
638 max_iter = self.refine_level + max_iter
639 curr_error = np.inf
640 t_start = time.time()
642 # Estimate bounds from test set if provided (override current bounds if they are set)
643 if estimate_bounds:
644 if ytest is not None:
645 y_samples = to_surrogate_dataset(ytest, self.outputs(), del_fields=True)[0] # normalize/compress
646 _combine_latent_arrays(y_samples)
647 coupling_vars = {k: v for k, v in self.coupling_variables().items() if k in y_samples}
648 y_min, y_max = {}, {}
649 for var in coupling_vars.values():
650 y_min[var] = np.nanmin(y_samples[var], axis=0)
651 y_max[var] = np.nanmax(y_samples[var], axis=0)
652 if var.compression is not None:
653 new_domain = list(zip(y_min[var].tolist(), y_max[var].tolist()))
654 var.update_domain(new_domain, override=True)
655 else:
656 new_domain = (float(y_min[var]), float(y_max[var]))
657 var.update_domain(var.denormalize(new_domain), override=True)
658 del y_samples
659 else:
660 self.logger.warning('Could not estimate bounds for coupling variables: no test set provided. '
661 'Make sure you manually provide (good) coupling variable domains.')
663 # Track convergence progress on the error indicator and test set (plot to file)
664 if self.root_dir is not None:
665 err_record = [res['added_error'] for res in self.train_history]
666 err_fig, err_ax = plt.subplots(figsize=(6, 5), layout='tight')
668 if xtest is not None and ytest is not None:
669 num_plot = min(len(targets), 3)
670 test_record = np.full((self.refine_level, num_plot), np.nan)
671 t_fig, t_ax = plt.subplots(1, num_plot, figsize=(3.5 * num_plot, 4), layout='tight', squeeze=False,
672 sharey='row')
673 for j, res in enumerate(self.train_history):
674 for i, var in enumerate(targets[:num_plot]):
675 if (perf := res.get('test_error')) is not None:
676 test_record[j, i] = perf[var]
678 while True:
679 # Adaptive refinement step
680 train_result = self.refine(targets=targets, num_refine=num_refine, update_bounds=update_bounds,
681 executor=executor)
682 if train_result['component'] is None:
683 self._print_title_str('Termination criteria reached: No candidates left to refine')
684 break
686 curr_error = train_result['added_error']
688 # Plot progress of error indicator
689 if self.root_dir is not None:
690 err_record.append(curr_error)
692 if plot_interval > 0 and self.refine_level % plot_interval == 0:
693 err_ax.clear(); err_ax.set_yscale('log'); err_ax.grid()
694 err_ax.plot(err_record, '-k')
695 err_ax.set_xlabel('Iteration'); err_ax.set_ylabel('Relative error indicator')
696 err_fig.savefig(str(Path(self.root_dir) / 'error_indicator.pdf'), format='pdf', bbox_inches='tight')
698 # Save performance on a test set
699 if xtest is not None and ytest is not None:
700 perf = self.test_set_performance(xtest, ytest) if self.refine_level >= start_test_check else (
701 {str(var): np.nan for var in ytest}) # don't compute if components are uninitialized
702 train_result['test_error'] = perf.copy()
704 if self.root_dir is not None:
705 test_record = np.vstack((test_record, np.array([perf[var] for var in targets[:num_plot]])))
707 if plot_interval > 0 and self.refine_level % plot_interval == 0:
708 for i in range(num_plot):
709 t_ax[0, i].clear(); t_ax[0, i].set_yscale('log'); t_ax[0, i].grid()
710 t_ax[0, i].plot(test_record[:, i], '-k')
711 t_ax[0, i].set_title(self.outputs()[targets[i]].get_tex(units=True))
712 t_ax[0, i].set_xlabel('Iteration')
713 t_ax[0, i].set_ylabel('Test set relative error' if i==0 else '')
714 t_fig.savefig(str(Path(self.root_dir) / 'test_set_error.pdf'),format='pdf',bbox_inches='tight')
716 self.train_history.append(train_result)
717 if self.root_dir is not None and save_interval > 0 and self.refine_level % save_interval == 0:
718 iter_name = f'{self.name}_iter{self.refine_level}'
719 if not (pth := self.root_dir / 'surrogates' / iter_name).is_dir():
720 os.mkdir(pth)
721 self.save_to_file(f'{iter_name}.yml', save_dir=pth) # Save to an iteration-specific directory
723 # Check all end conditions
724 if self.refine_level >= max_iter:
725 self._print_title_str(f'Termination criteria reached: Max iteration {self.refine_level}/{max_iter}')
726 break
727 if curr_error < max_tol:
728 self._print_title_str(f'Termination criteria reached: relative error {curr_error} < tol {max_tol}')
729 break
730 if ((time.time() - t_start) / 3600.0) >= runtime_hr:
731 actual = datetime.timedelta(seconds=time.time() - t_start)
732 target = datetime.timedelta(seconds=runtime_hr * 3600)
733 self._print_title_str(f'Termination criteria reached: runtime {str(actual)} > {str(target)}')
734 break
736 if self.root_dir is not None:
737 iter_name = f'{self.name}_iter{self.refine_level}'
738 if not (pth := self.root_dir / 'surrogates' / iter_name).is_dir():
739 os.mkdir(pth)
740 self.save_to_file(f'{iter_name}.yml', save_dir=pth)
742 if xtest is not None and ytest is not None:
743 self._save_test_set((xtest, ytest))
745 self.logger.info(f'Final system surrogate: \n {self}')
747 def test_set_performance(self, xtest: Dataset, ytest: Dataset, index_set='train') -> Dataset:
748 """Compute the relative L2 error on a test set for the given target output variables.
750 :param xtest: `dict` of test set input samples (unnormalized)
751 :param ytest: `dict` of test set output samples (unnormalized)
752 :param index_set: index set to use for prediction (defaults to 'train')
753 :returns: `dict` of relative L2 errors for each target output variable
754 """
755 xtest = to_surrogate_dataset(xtest, self.inputs(), del_fields=True)[0]
756 ysurr = self.predict(xtest, index_set=index_set, targets=list(ytest.keys()))
757 ysurr = to_model_dataset(ysurr, self.outputs(), del_latent=True)[0]
758 return {str(var): float(relative_error(ysurr[var], ytest[var])) for var in ytest}
760 def refine(self, targets: list = None, num_refine: int = 100, update_bounds: bool = True,
761 executor: Executor = None) -> TrainIteration:
762 """Perform a single adaptive refinement step on the system surrogate.
764 :param targets: list of system output variables to focus refinement on, use all outputs if not specified
765 :param num_refine: number of input samples to compute error indicators on
766 :param update_bounds: whether to continuously update coupling variable bounds during refinement
767 :param executor: a `concurrent.futures.Executor` object to parallelize model evaluations
768 :returns: `dict` of the refinement results indicating the chosen component and candidate index
769 """
770 self._print_title_str(f'Refining system surrogate: iteration {self.refine_level + 1}')
771 targets = targets or self.outputs()
773 # Check for uninitialized components and refine those first
774 for comp in self.components:
775 if len(comp.active_set) == 0 and comp.has_surrogate:
776 alpha_star = (0,) * len(comp.model_fidelity)
777 beta_star = (0,) * len(comp.max_beta)
778 self.logger.info(f"Initializing component {comp.name}: adding {(alpha_star, beta_star)} to active set")
779 model_dir = (pth / 'components' / comp.name) if (pth := self.root_dir) is not None else None
780 comp.activate_index(alpha_star, beta_star, model_dir=model_dir, executor=executor)
781 cost_star = max(1., comp.get_cost(alpha_star, beta_star)) # Cpu time (s)
782 err_star = np.nan
783 num_evals = round(cost_star / comp.model_costs.get(alpha_star, 1.))
784 return {'component': comp.name, 'alpha': alpha_star, 'beta': beta_star, 'num_evals': int(num_evals),
785 'added_cost': float(cost_star), 'added_error': float(err_star)}
787 # Compute entire integrated-surrogate on a random test set for global system QoI error estimation
788 x_samples = self.sample_inputs(num_refine)
789 y_curr = self.predict(x_samples, index_set='train', targets=targets)
790 _combine_latent_arrays(y_curr)
791 coupling_vars = {k: v for k, v in self.coupling_variables().items() if k in y_curr}
793 y_min, y_max = None, None
794 if update_bounds:
795 y_min = {var: np.nanmin(y_curr[var], axis=0, keepdims=True) for var in coupling_vars} # (1, ydim)
796 y_max = {var: np.nanmax(y_curr[var], axis=0, keepdims=True) for var in coupling_vars} # (1, ydim)
798 # Find the candidate surrogate with the largest error indicator
799 error_max, error_indicator = -np.inf, -np.inf
800 comp_star, alpha_star, beta_star, err_star, cost_star = None, None, None, -np.inf, 0
801 for comp in self.components:
802 if not comp.has_surrogate: # Skip analytic models that don't need a surrogate
803 continue
805 self.logger.info(f"Estimating error for component '{comp.name}'...")
807 if len(comp.candidate_set) > 0:
808 candidates = list(comp.candidate_set)
809 if executor is None:
810 ret = [self.predict(x_samples, targets=targets, index_set={comp.name: {(alpha, beta)}},
811 incremental={comp.name: True})
812 for alpha, beta in candidates]
813 else:
814 temp_buffer = self._remove_unpickleable()
815 futures = [executor.submit(self.predict, x_samples, targets=targets,
816 index_set={comp.name: {(alpha, beta)}}, incremental={comp.name: True})
817 for alpha, beta in candidates]
818 wait(futures, timeout=None, return_when=ALL_COMPLETED)
819 ret = [f.result() for f in futures]
820 self._restore_unpickleable(temp_buffer)
822 for i, y_cand in enumerate(ret):
823 alpha, beta = candidates[i]
824 _combine_latent_arrays(y_cand)
825 error = {}
826 for var, arr in y_cand.items():
827 if var in targets:
828 error[var] = relative_error(arr, y_curr[var])
830 if update_bounds and var in coupling_vars:
831 y_min[var] = np.nanmin(np.concatenate((y_min[var], arr), axis=0), axis=0, keepdims=True)
832 y_max[var] = np.nanmax(np.concatenate((y_max[var], arr), axis=0), axis=0, keepdims=True)
834 delta_error = np.nanmax([np.nanmax(error[var]) for var in error]) # Max error over all target QoIs
835 delta_work = max(1., comp.get_cost(alpha, beta)) # Cpu time (s)
836 error_indicator = delta_error / delta_work
838 self.logger.info(f"Candidate multi-index: {(alpha, beta)}. Relative error: {delta_error}. "
839 f"Error indicator: {error_indicator}.")
841 if error_indicator > error_max:
842 error_max = error_indicator
843 comp_star, alpha_star, beta_star, err_star, cost_star = (
844 comp.name, alpha, beta, delta_error, delta_work)
845 else:
846 self.logger.info(f"Component '{comp.name}' has no available candidates left!")
848 # Update all coupling variable ranges
849 if update_bounds:
850 for var in coupling_vars.values():
851 if np.all(~np.isnan(y_min[var])) and np.all(~np.isnan(y_max[var])):
852 if var.compression is not None:
853 new_domain = list(zip(np.squeeze(y_min[var], axis=0).tolist(),
854 np.squeeze(y_max[var], axis=0).tolist()))
855 var.update_domain(new_domain)
856 else:
857 new_domain = (y_min[var][0], y_max[var][0])
858 var.update_domain(var.denormalize(new_domain)) # bds will be in norm space from predict() call
860 # Add the chosen multi-index to the chosen component
861 if comp_star is not None:
862 self.logger.info(f"Candidate multi-index {(alpha_star, beta_star)} chosen for component '{comp_star}'.")
863 model_dir = (pth / 'components' / comp_star) if (pth := self.root_dir) is not None else None
864 self[comp_star].activate_index(alpha_star, beta_star, model_dir=model_dir, executor=executor)
865 num_evals = round(cost_star / self[comp_star].model_costs.get(alpha_star, 1.))
866 else:
867 self.logger.info(f"No candidates left for refinement, iteration: {self.refine_level}")
868 num_evals = 0
870 # Return the results of the refinement step
871 return {'component': comp_star, 'alpha': alpha_star, 'beta': beta_star, 'num_evals': int(num_evals),
872 'added_cost': float(cost_star), 'added_error': float(err_star)}
874 def predict(self, x: dict | Dataset,
875 max_fpi_iter: int = 100,
876 anderson_mem: int = 10,
877 fpi_tol: float = 1e-10,
878 use_model: str | tuple | dict = None,
879 model_dir: str | Path = None,
880 verbose: bool = False,
881 index_set: dict[str: IndexSet | Literal['train', 'test']] = 'test',
882 misc_coeff: dict[str: MiscTree] = None,
883 normalized_inputs: bool = True,
884 incremental: dict[str, bool] = False,
885 targets: list[str] = None,
886 executor: Executor = None,
887 var_shape: dict[str, tuple] = None) -> Dataset:
888 """Evaluate the system surrogate at inputs `x`. Return `y = system(x)`.
890 !!! Warning "Computing the true model with feedback loops"
891 You can use this function to predict outputs for your MD system using the full-order models rather than the
892 surrogate, by specifying `use_model`. This is convenient because the `System` manages all the
893 coupled information flow between models automatically. However, it is *highly* recommended to not use
894 the full model if your system contains feedback loops. The FPI nonlinear solver would be infeasible using
895 anything more computationally demanding than the surrogate.
897 :param x: `dict` of input samples for each variable in the system
898 :param max_fpi_iter: the limit on convergence for the fixed-point iteration routine
899 :param anderson_mem: hyperparameter for tuning the convergence of FPI with anderson acceleration
900 :param fpi_tol: tolerance limit for convergence of fixed-point iteration
901 :param use_model: 'best'=highest-fidelity, 'worst'=lowest-fidelity, tuple=specific fidelity, None=surrogate,
902 specify a `dict` of the above to assign different model fidelities for diff components
903 :param model_dir: directory to save model outputs if `use_model` is specified
904 :param verbose: whether to print out iteration progress during execution
905 :param index_set: `dict(comp=[indices])` to override the active set for a component, defaults to using the
906 `test` set for every component. Can also specify `train` for any component or a valid
907 `IndexSet` object. If `incremental` is specified, will be overwritten with `train`.
908 :param misc_coeff: `dict(comp=MiscTree)` to override the default coefficients for a component, passes through
909 along with `index_set` and `incremental` to `comp.predict()`.
910 :param normalized_inputs: true if the passed inputs are compressed/normalized for surrogate evaluation
911 (default), such as inputs returned by `sample_inputs`. Set to `False` if you are
912 passing inputs as the true models would expect them instead (i.e. not normalized).
913 :param incremental: whether to add `index_set` to the current active set for each component (temporarily);
914 this will set `index_set='train'` for all other components (since incremental will
915 augment the "training" active sets, not the "testing" candidate sets)
916 :param targets: list of output variables to return, defaults to returning all system outputs
917 :param executor: a `concurrent.futures.Executor` object to parallelize model evaluations
918 :param var_shape: (Optional) `dict` of shapes for field quantity inputs in `x` -- you would only specify this
919 if passing field qtys directly to the models (i.e. not using `sample_inputs`)
920 :returns: `dict` of output variables - the surrogate approximation of the system outputs (or the true model)
921 """
922 # Format inputs and allocate space
923 var_shape = var_shape or {}
924 x, loop_shape = format_inputs(x, var_shape=var_shape) # {'x': (N, *var_shape)}
925 y = {}
926 all_inputs = ChainMap(x, y)
927 N = int(np.prod(loop_shape))
928 t1 = 0
929 output_dir = None
930 norm_status = {var: normalized_inputs for var in x} # keep track of whether inputs are normalized or not
931 graph = self.graph()
933 # Keep track of what outputs are computed
934 is_computed = {}
935 for var in (targets or self.outputs()):
936 if (v := self.outputs().get(var, None)) is not None:
937 if v.compression is not None:
938 for field in v.compression.fields:
939 is_computed[field] = False
940 else:
941 is_computed[var] = False
943 def _set_default(struct: dict, default=None):
944 """Helper to set a default value for each component key in a `dict`. Ensures all components have a value."""
945 if struct is not None:
946 if not isinstance(struct, dict):
947 struct = {node: struct for node in graph.nodes} # use same for each component
948 else:
949 struct = {node: default for node in graph.nodes}
950 return {node: struct.get(node, default) for node in graph.nodes}
952 # Ensure use_model, index_set, and incremental are specified for each component model
953 use_model = _set_default(use_model, None)
954 incremental = _set_default(incremental, False) # default to train if incremental anywhere
955 index_set = _set_default(index_set, 'train' if any([incremental[node] for node in graph.nodes]) else 'test')
956 misc_coeff = _set_default(misc_coeff, None)
958 samples = _Converged(N) # track convergence of samples
960 def _gather_comp_inputs(comp, coupling=None):
961 """Helper to gather inputs for a component, making sure they are normalized correctly. Any coupling
962 variables passed in will be used in preference over `all_inputs`.
963 """
964 # Will access but not modify: all_inputs, use_model, norm_status
965 kwds = {}
966 comp_input = {}
967 coupling = coupling or {}
969 # Take coupling variables as a priority
970 comp_input.update({var: np.copy(arr[samples.curr_idx, ...]) for var, arr in
971 coupling.items() if str(var).split(LATENT_STR_ID)[0] in comp.inputs})
972 # Gather all other inputs
973 for var, arr in all_inputs.items():
974 var_id = str(var).split(LATENT_STR_ID)[0]
975 if var_id in comp.inputs and var not in coupling:
976 comp_input[var] = np.copy(arr[samples.curr_idx, ...])
978 # Gather extra fields (will never be in coupling since field couplings should always be latent coeff)
979 field_coords = {f'{var}_coords': comp.model_kwargs.get(f'{var}_coords', None) for var in comp.inputs}
980 for var in comp.inputs:
981 if var not in comp_input and var.compression is not None:
982 for field in var.compression.fields:
983 if field in all_inputs:
984 comp_input[field] = np.copy(all_inputs[field][samples.curr_idx, ...])
986 call_model = use_model.get(comp.name, None) is not None
988 # Make sure we format all inputs for model evaluation (i.e. denormalize)
989 if call_model:
990 norm_inputs = {var: arr for var, arr in comp_input.items() if norm_status[var]}
991 if len(norm_inputs) > 0:
992 denorm_inputs, field_coords = to_model_dataset(norm_inputs, comp.inputs, del_latent=True,
993 **field_coords)
994 for var in norm_inputs:
995 del comp_input[var]
996 kwds.update(field_coords)
997 comp_input.update(denorm_inputs)
999 # Otherwise, make sure we format inputs for surrogate evaluation (i.e. normalize)
1000 else:
1001 denorm_inputs = {var: arr for var, arr in comp_input.items() if not norm_status[var]}
1002 if len(denorm_inputs) > 0:
1003 norm_inputs, _ = to_surrogate_dataset(denorm_inputs, comp.inputs, del_fields=True, **field_coords)
1005 for var in denorm_inputs:
1006 del comp_input[var]
1007 comp_input.update(norm_inputs)
1009 return comp_input, kwds, call_model
1011 # Convert system into DAG by grouping strongly-connected-components
1012 dag = nx.condensation(graph)
1014 # Compute component models in topological order
1015 for supernode in nx.topological_sort(dag):
1016 if np.all(list(is_computed.values())):
1017 break # Exit early if all selected return qois are computed
1019 scc = [n for n in dag.nodes[supernode]['members']]
1020 samples.reset_convergence()
1022 # Compute single component feedforward output (no FPI needed)
1023 if len(scc) == 1:
1024 if verbose:
1025 self.logger.info(f"Running component '{scc[0]}'...")
1026 t1 = time.time()
1028 # Gather inputs
1029 comp = self[scc[0]]
1030 comp_input, kwds, call_model = _gather_comp_inputs(comp)
1032 # Compute outputs
1033 if model_dir is not None:
1034 output_dir = Path(model_dir) / scc[0]
1035 if not output_dir.exists():
1036 os.mkdir(output_dir)
1037 comp_output = comp.predict(comp_input, use_model=use_model.get(scc[0]), model_dir=output_dir,
1038 index_set=index_set.get(scc[0]), incremental=incremental.get(scc[0]),
1039 misc_coeff=misc_coeff.get(scc[0]), executor=executor, **kwds)
1040 for var, arr in comp_output.items():
1041 output_shape = arr.shape[1:]
1042 if y.get(var) is None:
1043 y.setdefault(var, np.full((N, *output_shape), np.nan))
1044 y[var][samples.curr_idx, ...] = arr
1045 samples.valid_idx = np.logical_and(samples.valid_idx, ~np.any(np.isnan(y[var]),
1046 axis=tuple(range(1, y[var].ndim))))
1047 is_computed[str(var).split(LATENT_STR_ID)[0]] = True
1048 norm_status[var] = not call_model
1050 if verbose:
1051 self.logger.info(f"Component '{scc[0]}' completed. Runtime: {time.time() - t1} s")
1053 # Handle FPI for SCCs with more than one component
1054 else:
1055 # Set the initial guess for all coupling vars (middle of domain)
1056 scc_inputs = ChainMap(*[self[comp].inputs for comp in scc])
1057 scc_outputs = ChainMap(*[self[comp].outputs for comp in scc])
1058 coupling_vars = [scc_inputs.get(var) for var in (scc_inputs.keys() - x.keys()) if var in scc_outputs]
1059 coupling_prev = {}
1060 for var in coupling_vars:
1061 domain = var.get_domain()
1062 if isinstance(domain, list): # Latent coefficients are the coupling variables
1063 for i, d in enumerate(domain):
1064 lb, ub = d
1065 coupling_prev[f'{var.name}{LATENT_STR_ID}{i}'] = np.broadcast_to((lb + ub) / 2, (N,)).copy()
1066 norm_status[f'{var.name}{LATENT_STR_ID}{i}'] = True
1067 else:
1068 lb, ub = var.normalize(domain)
1069 shape = (N,) + (1,) * len(var_shape.get(var, ()))
1070 coupling_prev[var] = np.broadcast_to((lb + ub) / 2, shape).copy()
1071 norm_status[var] = True
1073 residual_hist = deque(maxlen=anderson_mem)
1074 coupling_hist = deque(maxlen=anderson_mem)
1076 def _end_conditions_met():
1077 """Helper to compute residual, update history, and check end conditions."""
1078 residual = {}
1079 converged_idx = np.full(N, True)
1080 for var in coupling_prev:
1081 residual[var] = y[var] - coupling_prev[var]
1082 var_conv = np.all(np.abs(residual[var]) <= fpi_tol, axis=tuple(range(1, residual[var].ndim)))
1083 converged_idx = np.logical_and(converged_idx, var_conv)
1084 samples.valid_idx = np.logical_and(samples.valid_idx, ~np.isnan(coupling_prev[var]))
1085 samples.converged_idx = np.logical_or(samples.converged_idx, converged_idx)
1087 for var in coupling_prev:
1088 coupling_prev[var][samples.curr_idx, ...] = y[var][samples.curr_idx, ...]
1089 residual_hist.append(copy.deepcopy(residual))
1090 coupling_hist.append(copy.deepcopy(coupling_prev))
1092 if int(np.sum(samples.curr_idx)) == 0:
1093 if verbose:
1094 self.logger.info(f'FPI converged for SCC {scc} in {k} iterations with tol '
1095 f'{fpi_tol}. Final time: {time.time() - t1} s')
1096 return True
1098 max_error = np.max([np.max(np.abs(res[samples.curr_idx, ...])) for res in residual.values()])
1099 if verbose:
1100 self.logger.info(f'FPI iter: {k}. Max residual: {max_error}. Time: {time.time() - t1} s')
1102 if k >= max_fpi_iter:
1103 self.logger.warning(f'FPI did not converge in {max_fpi_iter} iterations for SCC {scc}: '
1104 f'{max_error} > tol {fpi_tol}. Some samples will be returned as NaN.')
1105 for var in coupling_prev:
1106 y[var][~samples.converged_idx, ...] = np.nan
1107 samples.valid_idx = np.logical_and(samples.valid_idx, samples.converged_idx)
1108 return True
1109 else:
1110 return False
1112 # Main FPI loop
1113 if verbose:
1114 self.logger.info(f"Initializing FPI for SCC {scc} ...")
1115 t1 = time.time()
1116 k = 0
1117 while True:
1118 for node in scc:
1119 # Gather inputs from exogenous and coupling sources
1120 comp = self[node]
1121 comp_input, kwds, call_model = _gather_comp_inputs(comp, coupling=coupling_prev)
1123 # Compute outputs (just don't do this FPI with expensive real models, please..)
1124 comp_output = comp.predict(comp_input, use_model=use_model.get(node), model_dir=None,
1125 index_set=index_set.get(node), incremental=incremental.get(node),
1126 misc_coeff=misc_coeff.get(node), executor=executor, **kwds)
1127 for var, arr in comp_output.items():
1128 output_shape = arr.shape[1:]
1129 if y.get(var) is not None:
1130 if output_shape != y.get(var).shape[1:]:
1131 y[var] = _merge_shapes((N, *output_shape), y[var])
1132 else:
1133 y.setdefault(var, np.full((N, *output_shape), np.nan))
1134 norm_status[var] = not call_model
1135 is_computed[var] = True
1136 y[var][samples.curr_idx, ...] = arr
1138 # Compute residual and check end conditions
1139 if _end_conditions_met():
1140 break
1142 # Skip anderson acceleration on first iteration
1143 if k == 0:
1144 k += 1
1145 continue
1147 # Iterate with anderson acceleration (only iterate on samples that are not yet converged)
1148 N_curr = int(np.sum(samples.curr_idx))
1149 mk = len(residual_hist) # Max of anderson mem
1150 var_shapes = []
1151 xdims = []
1152 for var in coupling_prev:
1153 shape = coupling_prev[var].shape[1:]
1154 var_shapes.append(shape)
1155 xdims.append(int(np.prod(shape)))
1156 N_couple = int(np.sum(xdims))
1157 res_snap = np.empty((N_curr, N_couple, mk)) # Shortened snapshot of residual history
1158 coupling_snap = np.empty((N_curr, N_couple, mk)) # Shortened snapshot of coupling history
1159 for i, (coupling_iter, residual_iter) in enumerate(zip(coupling_hist, residual_hist)):
1160 start_idx = 0
1161 for j, var in enumerate(coupling_prev):
1162 end_idx = start_idx + xdims[j]
1163 coupling_snap[:, start_idx:end_idx, i] = coupling_iter[var][samples.curr_idx, ...].reshape((N_curr, -1)) # noqa: E501
1164 res_snap[:, start_idx:end_idx, i] = residual_iter[var][samples.curr_idx, ...].reshape((N_curr, -1)) # noqa: E501
1165 start_idx = end_idx
1166 C = np.ones((N_curr, 1, mk))
1167 b = np.zeros((N_curr, N_couple, 1))
1168 d = np.ones((N_curr, 1, 1))
1169 alpha = np.expand_dims(constrained_lls(res_snap, b, C, d), axis=-3) # (..., 1, mk, 1)
1170 coupling_new = np.squeeze(coupling_snap[:, :, np.newaxis, :] @ alpha, axis=(-1, -2))
1171 start_idx = 0
1172 for j, var in enumerate(coupling_prev):
1173 end_idx = start_idx + xdims[j]
1174 coupling_prev[var][samples.curr_idx, ...] = coupling_new[:, start_idx:end_idx].reshape((N_curr, *var_shapes[j])) # noqa: E501
1175 start_idx = end_idx
1176 k += 1
1178 # Return all component outputs; samples that didn't converge during FPI are left as np.nan
1179 return format_outputs(y, loop_shape)
1181 def __call__(self, *args, **kwargs):
1182 """Convenience wrapper to allow calling as `ret = System(x)`."""
1183 return self.predict(*args, **kwargs)
1185 def __eq__(self, other):
1186 if not isinstance(other, System):
1187 return False
1188 return (self.components == other.components and
1189 self.name == other.name and
1190 self.train_history == other.train_history)
1192 def __getitem__(self, component: str) -> Component:
1193 """Convenience method to get a `Component` object from the `System`.
1195 :param component: the name of the component to get
1196 :returns: the `Component` object
1197 """
1198 return self.get_component(component)
1200 def get_component(self, comp_name: str) -> Component:
1201 """Return the `Component` object for this component.
1203 :param comp_name: name of the component to return
1204 :raises KeyError: if the component does not exist
1205 :returns: the `Component` object
1206 """
1207 if comp_name.lower() == 'system':
1208 return self
1209 else:
1210 for comp in self.components:
1211 if comp.name == comp_name:
1212 return comp
1213 raise KeyError(f"Component '{comp_name}' not found in system.")
1215 def _print_title_str(self, title_str: str):
1216 """Log an important message."""
1217 self.logger.info('-' * int(len(title_str)/2) + title_str + '-' * int(len(title_str)/2))
1219 def _remove_unpickleable(self) -> dict:
1220 """Remove and return unpickleable attributes before pickling (just the logger)."""
1221 stdout = False
1222 log_file = None
1223 if self._logger is not None:
1224 for handler in self._logger.handlers:
1225 if isinstance(handler, logging.StreamHandler):
1226 stdout = True
1227 break
1228 for handler in self._logger.handlers:
1229 if isinstance(handler, logging.FileHandler):
1230 log_file = handler.baseFilename
1231 break
1233 buffer = {'log_stdout': stdout, 'log_file': log_file}
1234 self.logger = None
1235 return buffer
1237 def _restore_unpickleable(self, buffer: dict):
1238 """Restore the unpickleable attributes from `buffer` after unpickling."""
1239 self.set_logger(log_file=buffer.get('log_file', None), stdout=buffer.get('log_stdout', None))
1241 def _get_test_set(self, test_set: str | Path | tuple = None) -> tuple:
1242 """Try to load a test set from the root directory if it exists."""
1243 if isinstance(test_set, tuple):
1244 return test_set # (xtest, ytest)
1245 else:
1246 ret = (None, None)
1247 if test_set is not None:
1248 test_set = Path(test_set)
1249 elif self.root_dir is not None:
1250 test_set = self.root_dir / 'test_set.pkl'
1252 if test_set is not None:
1253 if test_set.exists():
1254 with open(test_set, 'rb') as fd:
1255 data = pickle.load(fd)
1256 ret = data['test_set']
1258 return ret
1260 def _save_test_set(self, test_set: tuple = None):
1261 """Save the test set to the root directory if possible."""
1262 if self.root_dir is not None and test_set is not None:
1263 test_file = self.root_dir / 'test_set.pkl'
1264 if not test_file.exists():
1265 with open(test_file, 'wb') as fd:
1266 pickle.dump({'test_set': test_set}, fd)
1268 def save_to_file(self, filename: str, save_dir: str | Path = None, dumper=None):
1269 """Save surrogate to file. Defaults to `root/surrogates/filename.yml` with the default yaml encoder.
1271 :param filename: the name of the save file
1272 :param save_dir: the directory to save the file to (defaults to `root/surrogates` or `cwd()`)
1273 :param dumper: the encoder to use (defaults to the `amisc` yaml encoder)
1274 """
1275 from amisc import YamlLoader
1276 encoder = dumper or YamlLoader
1277 save_dir = save_dir or self.root_dir or Path.cwd()
1278 if Path(save_dir) == self.root_dir:
1279 save_dir = self.root_dir / 'surrogates'
1280 encoder.dump(self, Path(save_dir) / filename)
1282 @staticmethod
1283 def load_from_file(filename: str | Path, root_dir: str | Path = None, loader=None):
1284 """Load surrogate from file. Defaults to yaml loading. Tries to infer `amisc` directory structure.
1286 :param filename: the name of the load file
1287 :param root_dir: set this as the surrogate's root directory (will try to load from `amisc_` fmt by default)
1288 :param loader: the encoder to use (defaults to the `amisc` yaml encoder)
1289 """
1290 from amisc import YamlLoader
1291 encoder = loader or YamlLoader
1292 system = encoder.load(filename)
1293 root_dir = root_dir or system.root_dir
1295 # Try to infer amisc_root/surrogates/iter/filename structure
1296 if root_dir is None:
1297 parts = Path(filename).resolve().parts
1298 if len(parts) > 2:
1299 if parts[-3].startswith('amisc_'):
1300 root_dir = Path(filename).resolve().parent.parent
1301 elif len(parts) > 3:
1302 if parts[-4].startswith('amisc_'):
1303 root_dir = Path(filename).resolve().parent.parent.parent
1304 system.root_dir = root_dir
1305 return system
1307 def clear(self):
1308 """Clear all surrogate model data and reset the system."""
1309 for comp in self.components:
1310 comp.clear()
1311 self.train_history.clear()
1313 def plot_slice(self, inputs: list[str] = None,
1314 outputs: list[str] = None,
1315 num_steps: int = 20,
1316 show_surr: bool = True,
1317 show_model: list = None,
1318 model_dir: str | Path = None,
1319 executor: Executor = None,
1320 nominal: dict[str: float] = None,
1321 random_walk: bool = False,
1322 from_file: str | Path = None,
1323 subplot_size_in: float = 3.):
1324 """Helper function to plot 1d slices of the surrogate and/or model outputs over the inputs. A single
1325 "slice" works by smoothly stepping from the lower bound of an input to its upper bound, while holding all other
1326 inputs constant at their nominal values (or smoothly varying them if `random_walk=True`).
1327 This function is useful for visualizing the behavior of the system surrogate and/or model(s) over a
1328 single input variable at a time.
1330 :param inputs: list of input variables to take 1d slices of (defaults to first 3 in `System.inputs`)
1331 :param outputs: list of model output variables to plot 1d slices of (defaults to first 3 in `System.outputs`)
1332 :param num_steps: the number of points to take in the 1d slice for each input variable; this amounts to a total
1333 of `num_steps*len(inputs)` model/surrogate evaluations
1334 :param show_surr: whether to show the surrogate prediction
1335 :param show_model: also plot model predictions, `list` of ['best', 'worst', tuple(alpha), etc.]
1336 :param model_dir: base directory to save model outputs (if specified)
1337 :param executor: a `concurrent.futures.Executor` object to parallelize model or surrogate evaluations
1338 :param nominal: `dict` of `var->nominal` to use as constant values for all non-sliced variables (use
1339 unnormalized values only; use `var_LATENT0` to specify nominal latent values)
1340 :param random_walk: whether to slice in a random d-dimensional direction instead of holding all non-slice
1341 variables const at `nominal`
1342 :param from_file: path to a `.pkl` file to load a saved slice from disk
1343 :param subplot_size_in: side length size of each square subplot in inches
1344 :returns: `fig, ax` with `len(inputs)` by `len(outputs)` subplots
1345 """
1346 # Manage loading important quantities from file (if provided)
1347 input_slices, output_slices_model, output_slices_surr = None, None, None
1348 if from_file is not None:
1349 with open(Path(from_file), 'rb') as fd:
1350 slice_data = pickle.load(fd)
1351 inputs = slice_data['inputs'] # Must use same input slices as save file
1352 show_model = slice_data['show_model'] # Must use same model data as save file
1353 outputs = slice_data.get('outputs') if outputs is None else outputs
1354 input_slices = slice_data['input_slices']
1355 model_dir = None # Don't run or save any models if loading from file
1357 # Set default values (take up to the first 3 inputs by default)
1358 all_inputs = self.inputs()
1359 all_outputs = self.outputs()
1360 rand_id = ''.join(random.choices(string.ascii_uppercase + string.digits, k=4))
1361 if model_dir is not None:
1362 os.mkdir(Path(model_dir) / f'sweep_{rand_id}')
1363 if nominal is None:
1364 nominal = dict()
1365 inputs = all_inputs[:3] if inputs is None else inputs
1366 outputs = all_outputs[:3] if outputs is None else outputs
1368 if show_model is not None and not isinstance(show_model, list):
1369 show_model = [show_model]
1371 # Handle field quantities (directly use latent variables or only the first one)
1372 for i, var in enumerate(list(inputs)):
1373 if LATENT_STR_ID not in str(var) and all_inputs[var].compression is not None:
1374 inputs[i] = f'{var}{LATENT_STR_ID}0'
1375 for i, var in enumerate(list(outputs)):
1376 if LATENT_STR_ID not in str(var) and all_outputs[var].compression is not None:
1377 outputs[i] = f'{var}{LATENT_STR_ID}0'
1379 bds = all_inputs.get_domains()
1380 xlabels = [all_inputs[var].get_tex(units=True) if LATENT_STR_ID not in str(var) else
1381 all_inputs[str(var).split(LATENT_STR_ID)[0]].get_tex(units=True) +
1382 f' (latent {str(var).split(LATENT_STR_ID)[1]})' for var in inputs]
1384 ylabels = [all_outputs[var].get_tex(units=True) if LATENT_STR_ID not in str(var) else
1385 all_outputs[str(var).split(LATENT_STR_ID)[0]].get_tex(units=True) +
1386 f' (latent {str(var).split(LATENT_STR_ID)[1]})' for var in outputs]
1388 # Construct slices of model inputs (if not provided)
1389 if input_slices is None:
1390 input_slices = {} # Each input variable with shape (num_steps, num_slice)
1391 for i in range(len(inputs)):
1392 if random_walk:
1393 # Make a random straight-line walk across d-cube
1394 r0 = self.sample_inputs((1,), use_pdf=False)
1395 rf = self.sample_inputs((1,), use_pdf=False)
1397 for var, bd in bds.items():
1398 if var == inputs[i]:
1399 r0[var] = np.atleast_1d(bd[0]) # Start slice at this lower bound
1400 rf[var] = np.atleast_1d(bd[1]) # Slice up to this upper bound
1402 step_size = (rf[var] - r0[var]) / (num_steps - 1)
1403 arr = r0[var] + step_size * np.arange(num_steps)
1405 input_slices[var] = arr[..., np.newaxis] if input_slices.get(var) is None else (
1406 np.concatenate((input_slices[var], arr[..., np.newaxis]), axis=-1))
1407 else:
1408 # Otherwise, only slice one variable
1409 for var, bd in bds.items():
1410 nom = nominal.get(var, np.mean(bd)) if LATENT_STR_ID in str(var) else (
1411 all_inputs[var].normalize(nominal.get(var, all_inputs[var].get_nominal())))
1412 arr = np.linspace(bd[0], bd[1], num_steps) if var == inputs[i] else np.full(num_steps, nom)
1414 input_slices[var] = arr[..., np.newaxis] if input_slices.get(var) is None else (
1415 np.concatenate((input_slices[var], arr[..., np.newaxis]), axis=-1))
1417 # Walk through each model that is requested by show_model
1418 if show_model is not None:
1419 if from_file is not None:
1420 output_slices_model = slice_data['output_slices_model']
1421 else:
1422 output_slices_model = list()
1423 for model in show_model:
1424 output_dir = None
1425 if model_dir is not None:
1426 output_dir = (Path(model_dir) / f'sweep_{rand_id}' /
1427 str(model).replace('{', '').replace('}', '').replace(':', '=').replace("'", ''))
1428 os.mkdir(output_dir)
1429 output_slices_model.append(self.predict(input_slices, use_model=model, model_dir=output_dir,
1430 executor=executor))
1431 if show_surr:
1432 output_slices_surr = self.predict(input_slices, executor=executor) \
1433 if from_file is None else slice_data['output_slices_surr']
1435 # Make len(outputs) by len(inputs) grid of subplots
1436 fig, axs = plt.subplots(len(outputs), len(inputs), sharex='col', sharey='row', squeeze=False)
1437 for i, output_var in enumerate(outputs):
1438 for j, input_var in enumerate(inputs):
1439 ax = axs[i, j]
1440 x = input_slices[input_var][:, j]
1442 if show_model is not None:
1443 c = np.array([[0, 0, 0, 1], [0.5, 0.5, 0.5, 1]]) if len(show_model) <= 2 else (
1444 plt.get_cmap('jet')(np.linspace(0, 1, len(show_model))))
1445 for k in range(len(show_model)):
1446 model_str = (str(show_model[k]).replace('{', '').replace('}', '')
1447 .replace(':', '=').replace("'", ''))
1448 model_ret = to_surrogate_dataset(output_slices_model[k], all_outputs)[0]
1449 y_model = model_ret[output_var][:, j]
1450 label = {'best': 'High-fidelity' if len(show_model) > 1 else 'Model',
1451 'worst': 'Low-fidelity'}.get(model_str, model_str)
1452 ax.plot(x, y_model, ls='-', c=c[k, :], label=label)
1454 if show_surr:
1455 y_surr = output_slices_surr[output_var][:, j]
1456 ax.plot(x, y_surr, '--r', label='Surrogate')
1458 ax.set_xlabel(xlabels[j] if i == len(outputs) - 1 else '')
1459 ax.set_ylabel(ylabels[i] if j == 0 else '')
1460 if i == 0 and j == len(inputs) - 1:
1461 ax.legend()
1462 fig.set_size_inches(subplot_size_in * len(inputs), subplot_size_in * len(outputs))
1463 fig.tight_layout()
1465 # Save results (unless we were already loading from a save file)
1466 if from_file is None and self.root_dir is not None:
1467 fname = f'in={",".join([str(v) for v in inputs])}_out={",".join([str(v) for v in outputs])}'
1468 fname = f'sweep_rand{rand_id}_' + fname if random_walk else f'sweep_nom{rand_id}_' + fname
1469 fdir = Path(self.root_dir) if model_dir is None else Path(model_dir) / f'sweep_{rand_id}'
1470 fig.savefig(fdir / f'{fname}.pdf', bbox_inches='tight', format='pdf')
1471 save_dict = {'inputs': inputs, 'outputs': outputs, 'show_model': show_model, 'show_surr': show_surr,
1472 'nominal': nominal, 'random_walk': random_walk, 'input_slices': input_slices,
1473 'output_slices_model': output_slices_model, 'output_slices_surr': output_slices_surr}
1474 with open(fdir / f'{fname}.pkl', 'wb') as fd:
1475 pickle.dump(save_dict, fd)
1477 return fig, axs
1479 def get_allocation(self, idx: int = None):
1480 """Get a breakdown of cost allocation up to a certain iteration number during training (starting at 1).
1482 :param idx: the iteration number to get allocation results for (defaults to last refinement step)
1483 :returns: `cost_alloc, eval_alloc, cost_cum, eval_cum` - the cost allocation per model/fidelity, the
1484 number of model evaluations per model/fidelity, the cumulative training cost, and the cumulative
1485 number of model evaluations
1486 """
1487 if idx is None:
1488 idx = self.refine_level
1489 if idx > self.refine_level:
1490 raise ValueError(f'Specified index: {idx} is greater than the max training level of {self.refine_level}')
1492 eval_alloc = dict() # Model evaluation counts per node and model fidelity
1493 cost_alloc = dict() # Cost allocation (cpu time in s) per node and model fidelity
1494 cost_cum = [] # Cumulative cost allocation during training
1495 eval_cum = [] # Cumulative number of model evaluations during training
1497 prev_cands = {comp.name: IndexSet() for comp in self.components} # empty candidate sets
1499 # Add cumulative training costs
1500 for train_res, active_sets, cand_sets, misc_coeff_train, misc_coeff_test in self.simulate_fit():
1501 comp = train_res['component']
1502 alpha = train_res['alpha']
1503 beta = train_res['beta']
1505 eval_alloc.setdefault(comp, dict())
1506 cost_alloc.setdefault(comp, dict())
1508 new_cands = cand_sets[comp].union({(alpha, beta)}) - prev_cands[comp] # newly computed candidates
1510 iter_cost = 0.
1511 iter_eval = 0
1512 for alpha_new, beta_new in new_cands:
1513 eval_alloc[comp].setdefault(alpha_new, 0.)
1514 cost_alloc[comp].setdefault(alpha_new, 0.)
1516 added_cost = self[comp].get_cost(alpha_new, beta_new)
1517 model_cost = self[comp].model_costs.get(alpha_new, 1.)
1518 added_eval = round(added_cost / model_cost)
1520 iter_cost += added_cost
1521 iter_eval += added_eval
1523 eval_alloc[comp][alpha_new] += added_eval
1524 cost_alloc[comp][alpha_new] += added_cost
1526 cost_cum.append(iter_cost)
1527 eval_cum.append(iter_eval)
1528 prev_cands[comp] = cand_sets[comp].union({(alpha, beta)})
1530 return cost_alloc, eval_alloc, np.cumsum(cost_cum), np.cumsum(eval_cum)
1532 def plot_allocation(self, cmap: str = 'Blues', text_bar_width: float = 0.06, arrow_bar_width: float = 0.02):
1533 """Plot bar charts showing cost allocation during training.
1535 !!! Warning "Beta feature"
1536 This has pretty good default settings, but it might look terrible for your use. Mostly provided here as
1537 a template for making cost allocation bar charts. Please feel free to copy and edit in your own code.
1539 :param cmap: the colormap string identifier for `plt`
1540 :param text_bar_width: the minimum total cost fraction above which a bar will print centered model fidelity text
1541 :param arrow_bar_width: the minimum total cost fraction above which a bar will try to print text with an arrow;
1542 below this amount, the bar is too skinny and won't print any text
1543 :returns: `fig, ax`, Figure and Axes objects
1544 """
1545 # Get total cost (including offline overhead)
1546 cost_alloc, eval_alloc, cost_cum, _ = self.get_allocation()
1547 total_cost = cost_cum[-1]
1549 # Remove nodes with cost=0 from alloc dicts (i.e. analytical models)
1550 remove_nodes = []
1551 for node, alpha_dict in cost_alloc.items():
1552 if len(alpha_dict) == 0:
1553 remove_nodes.append(node)
1554 for node in remove_nodes:
1555 del cost_alloc[node]
1556 del eval_alloc[node]
1558 # Bar chart showing cost allocation breakdown for MF system at final iteration
1559 fig, ax = plt.subplots(figsize=(6, 5), layout='tight')
1560 width = 0.7
1561 x = np.arange(len(cost_alloc))
1562 xlabels = list(cost_alloc.keys()) # One bar for each component
1563 cmap = plt.get_cmap(cmap)
1565 for j, (node, alpha_dict) in enumerate(cost_alloc.items()):
1566 bottom = 0
1567 c_intervals = np.linspace(0, 1, len(alpha_dict))
1568 bars = [(alpha, cost, cost / total_cost) for alpha, cost in alpha_dict.items()]
1569 bars = sorted(bars, key=lambda ele: ele[2], reverse=True)
1570 for i, (alpha, cost, frac) in enumerate(bars):
1571 p = ax.bar(x[j], frac, width, color=cmap(c_intervals[i]), linewidth=1,
1572 edgecolor=[0, 0, 0], bottom=bottom)
1573 bottom += frac
1574 if frac > text_bar_width:
1575 ax.bar_label(p, labels=[f'{alpha}, {round(eval_alloc[node][alpha])}'], label_type='center')
1576 elif frac > arrow_bar_width:
1577 xy = (x[j] + width / 2, bottom - frac / 2) # Label smaller bars with a text off to the side
1578 ax.annotate(f'{alpha}, {round(eval_alloc[node][alpha])}', xy, xytext=(xy[0] + 0.2, xy[1]),
1579 arrowprops={'arrowstyle': '->', 'linewidth': 1})
1580 else:
1581 pass # Don't label really small bars
1582 ax.set_xlabel('')
1583 ax.set_ylabel('Fraction of total cost')
1584 ax.set_xticks(x, xlabels)
1585 ax.set_xlim(left=-1, right=x[-1] + 1)
1587 if self.root_dir is not None:
1588 fig.savefig(Path(self.root_dir) / 'mf_allocation.pdf', bbox_inches='tight', format='pdf')
1590 return fig, ax
1592 def serialize(self, keep_components=False, serialize_args=None, serialize_kwargs=None) -> dict:
1593 """Convert to a `dict` with only standard Python types for fields.
1595 :param keep_components: whether to serialize the components as well (defaults to False)
1596 :param serialize_args: `dict` of arguments to pass to each component's serialize method
1597 :param serialize_kwargs: `dict` of keyword arguments to pass to each component's serialize method
1598 :returns: a `dict` representation of the `System` object
1599 """
1600 serialize_args = serialize_args or dict()
1601 serialize_kwargs = serialize_kwargs or dict()
1602 d = {}
1603 for key, value in self.__dict__.items():
1604 if value is not None and not key.startswith('_'):
1605 if key == 'components' and not keep_components:
1606 d[key] = [comp.serialize(keep_yaml_objects=False,
1607 serialize_args=serialize_args.get(comp.name),
1608 serialize_kwargs=serialize_kwargs.get(comp.name))
1609 for comp in value]
1610 elif key == 'train_history':
1611 if len(value) > 0:
1612 d[key] = value.serialize()
1613 else:
1614 if not isinstance(value, _builtin):
1615 self.logger.warning(f"Attribute '{key}' of type '{type(value)}' may not be a builtin "
1616 f"Python type. This may cause issues when saving/loading from file.")
1617 d[key] = value
1618 return d
1620 @classmethod
1621 def deserialize(cls, serialized_data: dict) -> System:
1622 """Construct a `System` object from serialized data."""
1623 return cls(**serialized_data)
1625 @staticmethod
1626 def _yaml_representer(dumper: yaml.Dumper, system: System):
1627 """Serialize a `System` object to a YAML representation."""
1628 return dumper.represent_mapping(System.yaml_tag, system.serialize(keep_components=True))
1630 @staticmethod
1631 def _yaml_constructor(loader: yaml.Loader, node):
1632 """Convert the `!SystemSurrogate` tag in yaml to a [`System`][amisc.system.System] object."""
1633 if isinstance(node, yaml.SequenceNode):
1634 return [ele if isinstance(ele, System) else System.deserialize(ele)
1635 for ele in loader.construct_sequence(node, deep=True)]
1636 elif isinstance(node, yaml.MappingNode):
1637 return System.deserialize(loader.construct_mapping(node, deep=True))
1638 else:
1639 raise NotImplementedError(f'The "{System.yaml_tag}" yaml tag can only be used on a yaml sequence or '
1640 f'mapping, not a "{type(node)}".')