Coverage for src/amisc/variable.py: 90%
416 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-11-05 19:26 +0000
1"""Provides an object-oriented interface for model inputs/outputs, random variables, scalars, and field quantities.
3Includes:
5- `Variable` — an object that stores information about a variable and includes methods for sampling, pdf evaluation,
6 normalization, compression, loading from file, etc. Variables can mostly be treated as strings
7 that have some additional information and utilities attached to them.
8- `VariableList` — a container for `Variables` that provides dict-like access of `Variables` by `name` along with normal
9 indexing and slicing.
11The preferred serialization of `Variable` and `VariableList` is to/from yaml. This is done by default with the
12`!Variable` and `!VariableList` yaml tags.
13"""
14from __future__ import annotations
16import ast
17import random
18import string
19from collections import OrderedDict
20from pathlib import Path
21from typing import ClassVar, Optional, Union
23import numpy as np
24import yaml
25from numpy.typing import ArrayLike
26from pydantic import BaseModel, ConfigDict, ValidationInfo, field_validator
28from amisc.compression import Compression
29from amisc.distribution import Distribution, LogUniform, Normal, Uniform
30from amisc.serialize import Serializable
31from amisc.transform import Minmax, Transform, Zscore
32from amisc.typing import LATENT_STR_ID, CompressionData
33from amisc.utils import _get_yaml_path, _inspect_assignment, search_for_file
35__all__ = ['Variable', 'VariableList']
36_TransformLike = Union[str, Transform, list[str | Transform]] # something that can be converted to a Transform
39class Variable(BaseModel, Serializable):
40 """Object for storing information about variables and providing methods for pdf evaluation, sampling, etc.
41 All fields will undergo pydantic validation and conversion to the correct types.
43 A simple variable object can be created with `var = Variable()`. All initialization options are optional and will
44 be given good defaults. You should probably at the very least give a memorable `name` and a `domain`. Variables
45 can mostly be treated as strings with some extra information/utilities attached.
47 With the `pyyaml` library installed, all `Variable` objects can be saved or loaded directly from a `.yml` file by
48 using the `!Variable` yaml tag (which is loaded by default with `amisc`).
50 - Use `Variable.distribution` to specify PDFs, such as for random variables. See the `Distribution` classes.
51 - Use `Variable.norm` to specify a transformed-space that is more amenable to surrogate construction
52 (e.g. mapping to the range (0,1)). See the `Transform` classes.
53 - Use `Variable.compression` to specify high-dimensional, coordinate-based field quantities,
54 such as from the output of many simulation software programs. See the `Compression` classes.
55 - Use `Variable.category` as an additional layer for using Variable's in different ways (e.g. set a "calibration"
56 category for Bayesian inference).
58 !!! Example
59 ```python
60 # Random variable
61 temp = Variable(name='T', description='Temperature', units='K', distribution='Uniform(280, 320)')
62 samples = temp.sample(100)
63 pdf = temp.pdf(samples)
65 # Field quantity
66 vel = Variable(name='u', description='Velocity', units='m/s', compression={'fields': ['ux', 'uy', 'uz']})
67 vel_data = ... # from a simulation
68 reduced_vel = vel.compress(vel_data)
69 ```
71 !!! Warning
72 Changes to collection fields (like `Variable.norm`) should completely reassign the _whole_
73 collection to trigger the correct validation, rather than editing particular entries. For example, reassign
74 `norm=['log', 'linear(2, 2)']` rather than editing norm via `norm.append('linear(2, 2)')`.
76 :ivar name: an identifier for the variable, can compare variables directly with strings for indexing purposes
77 :ivar nominal: a typical value for this variable
78 :ivar description: a lengthier description of the variable
79 :ivar units: assumed units for the variable (if applicable)
80 :ivar category: an additional descriptor for how this variable is used, e.g. calibration, operating, design, etc.
81 :ivar tex: latex format for the variable, i.e. "$x_i$"
82 :ivar compression: specifies field quantities and links to relevant compression data
83 :ivar distribution: a string specifier of a probability distribution function (see the `Distribution` types)
84 :ivar domain: the explicit domain bounds of the variable (limits of where you expect to use it);
85 for field quantities, this is a list of domains for each latent dimension
86 :ivar norm: specifier of a map to a transformed-space for surrogate construction (see the `Transform` types)
87 """
88 yaml_tag: ClassVar[str] = u'!Variable'
89 model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True, validate_default=True)
91 name: Optional[str] = None
92 nominal: Optional[float] = None
93 description: Optional[str] = None
94 units: Optional[str] = None
95 category: Optional[str] = None
96 tex: Optional[str] = None
97 compression: Optional[str | dict | Compression] = None
98 distribution: Optional[str | Distribution] = None
99 domain: Optional[str | tuple[float, float] | list] = None
100 norm: Optional[_TransformLike] = None
102 def __init__(self, /, name=None, **kwargs):
103 # Try to set the variable name if instantiated as "x = Variable()"
104 if name is None:
105 name = _inspect_assignment('Variable')
106 name = name or "X_" + "".join(random.choices(string.digits, k=3))
107 super().__init__(name=name, **kwargs)
109 @field_validator('tex')
110 @classmethod
111 def _validate_tex(cls, tex: str) -> str | None:
112 if tex is None:
113 return tex
114 if not tex.startswith('$'):
115 tex = rf'${tex}'
116 if not tex[-1] == '$':
117 tex = rf'{tex}$'
118 return tex
120 @field_validator('compression')
121 @classmethod
122 def _validate_compression(cls, compression: str | dict | Compression, info: ValidationInfo) -> Compression | None:
123 if compression is None:
124 return compression
125 elif isinstance(compression, str):
126 return Compression.deserialize(compression)
127 elif isinstance(compression, dict):
128 compression['fields'] = compression.get('fields', None) or [info.data['name']]
129 return Compression.from_dict(compression)
130 else:
131 compression.fields = compression.fields or [info.data['name']]
132 return compression
134 @field_validator('distribution')
135 @classmethod
136 def _validate_dist(cls, dist: str | Distribution) -> Distribution | None:
137 if dist is None:
138 return dist
139 if isinstance(dist, Distribution):
140 return dist
141 elif isinstance(dist, str):
142 return Distribution.from_string(dist)
143 else:
144 raise ValueError(f'Cannot convert {dist} to a Distribution object.')
146 @field_validator('domain')
147 @classmethod
148 def _validate_domain(cls, domain: list | tuple | str, info: ValidationInfo) -> tuple | list | None:
149 """Try to extract the domain from the distribution if not provided, or convert from a string.
150 Returns a list of domains for each latent dimension if this is a field quantity with compression.
151 """
152 if domain is None:
153 if dist := info.data['distribution']:
154 domain = dist.domain()
155 elif compression := info.data['compression']:
156 domain = compression.estimate_latent_ranges()
157 elif isinstance(domain, str):
158 domain = tuple(ast.literal_eval(domain.strip()))
159 elif isinstance(domain, list):
160 domain = [tuple(ast.literal_eval(d.strip())) if isinstance(d, str) else d for d in domain]
162 if domain is None:
163 return domain
165 if isinstance(domain, list):
166 for d in domain:
167 assert isinstance(d, tuple) and len(d) == 2
168 assert d[1] > d[0], 'Domain must be specified as (lower_bound, upper_bound)'
169 else:
170 assert isinstance(domain, tuple) and len(domain) == 2
171 assert domain[1] > domain[0], 'Domain must be specified as (lower_bound, upper_bound)'
173 return domain
175 @field_validator('norm')
176 @classmethod
177 def _validate_norm(cls, norm: _TransformLike, info: ValidationInfo) -> list[Transform] | None:
178 if norm is None:
179 return norm
180 norm = Transform.from_string(norm)
182 # Set default values for minmax and zscore transforms
183 domain = info.data['domain']
184 normal_args = None
185 if dist := info.data['distribution']:
186 if isinstance(dist, Normal):
187 normal_args = dist.dist_args
188 for transform in norm:
189 if isinstance(transform, Minmax):
190 if domain and np.any(np.isnan(transform.transform_args[0:2])):
191 transform.update(lb=domain[0], ub=domain[1])
192 elif isinstance(transform, Zscore):
193 if normal_args and np.any(np.isnan(transform.transform_args)):
194 transform.update(mu=normal_args[0], std=normal_args[1])
196 return norm
198 def __getitem__(self, item):
199 return getattr(self, item)
201 def __setitem__(self, key, value):
202 setattr(self, key, value)
204 def __str__(self):
205 return self.name
207 def __repr__(self):
208 return self.__str__()
210 def __hash__(self):
211 """Allows variables to be used as keys in dictionaries and to be considered equal to their string
212 representations.
213 """
214 return hash(self.name)
216 def __eq__(self, other):
217 """Consider two `Variables` equal if they share the same string name
219 Also returns true when checking if this `Variable` is equal to a string by itself.
220 """
221 if isinstance(other, Variable):
222 return self.name == other.name
223 elif isinstance(other, str):
224 return self.name == other
225 else:
226 return False
228 def get_tex(self, units: bool = False, symbol: bool = True) -> str:
229 """Return a raw string that is well-formatted for plotting (with latex).
231 :param units: whether to include the units in the string
232 :param symbol: just latex symbol if true, otherwise the full description
233 :returns: the latex formatted string
234 """
235 s = (self.tex if symbol else self.description) or self.name
236 return r'{} [{}]'.format(s, self.units or '-') if units else r'{}'.format(s)
238 def get_nominal(self) -> float | list | None:
239 """Return the nominal value of the variable. Defaults to the mean for a normal distribution or the
240 center of the domain if `var.nominal` is not specified. Returns a list of nominal values for each latent
241 dimension if this is a field quantity with compression.
243 :returns: the nominal value(s)
244 """
245 nominal = self.nominal
246 if nominal is None:
247 if dist := self.distribution:
248 nominal = float(dist.nominal())
249 elif domain := self.get_domain():
250 nominal = [np.mean(d) for d in domain] if isinstance(domain, list) else float(np.mean(domain))
252 return nominal
254 def get_domain(self) -> tuple | list | None:
255 """Return a tuple of the defined domain of this variable. Returns a list of domains for each latent dimension
256 if this is a field quantity with compression.
258 :returns: the domain(s) of this variable
259 """
260 if self.domain is None:
261 return None
262 elif isinstance(self.domain, list):
263 return self.domain
264 elif self.compression is not None:
265 # Try to infer a list of domains from compression latent size
266 try:
267 return [self.domain] * self.compression.latent_size()
268 except Exception as e:
269 raise ValueError(f'Variables with `compression` data should return a list of domains, one '
270 f'for each latent coefficient. Could not infer domain for "{self.name}".') from e
271 else:
272 return self.domain
274 def sample_domain(self, shape: tuple | int) -> np.ndarray:
275 """Return an array of the given `shape` for uniform samples over the domain of this variable. Returns
276 samples for each latent dimension if this is a field quantity with compression.
278 !!! Note
279 The last dim of the returned samples will be the latent space size for field quantities.
281 :param shape: the shape of samples to return
282 :returns: the random samples over the domain of the variable
283 """
284 if isinstance(shape, int):
285 shape = (shape, )
286 if domain := self.get_domain():
287 if isinstance(domain, list):
288 lb = np.atleast_1d([d[0] for d in domain])
289 ub = np.atleast_1d([d[1] for d in domain])
290 return np.random.rand(*shape, 1) * (ub - lb) + lb
291 else:
292 return np.random.rand(*shape) * (domain[1] - domain[0]) + domain[0]
293 else:
294 raise RuntimeError(f'Variable "{self.name}" does not have a domain specified.')
296 def update_domain(self, domain: tuple[float, float] | list[tuple], override: bool = False):
297 """Update the domain of this variable by taking the minimum or maximum of the new domain with the current domain
298 for the lower and upper bounds, respectively. Will attempt to update the domain of each latent dimension
299 if this is a field quantity with compression. If the variable has a `Uniform` distribution, this will
300 update the distribution's bounds too.
302 :param domain: the new domain(s) to update with
303 :param override: will simply set the domain to the new values rather than update against the current domain;
304 (default `False`)
305 """
306 def _update_domain(domain, curr_domain):
307 lb, ub = domain
308 ret = (lb, ub) if override else (min(lb, curr_domain[0]) if curr_domain is not None else lb,
309 max(ub, curr_domain[1]) if curr_domain is not None else ub)
310 return tuple(map(float, ret))
312 curr_domain = self.get_domain()
313 if isinstance(domain, list):
314 if not isinstance(curr_domain, list):
315 curr_domain = [curr_domain] * len(domain)
316 self.domain = [_update_domain(d, curr_domain[i]) for i, d in enumerate(domain)]
317 elif isinstance(curr_domain, list):
318 if not isinstance(domain, list):
319 domain = [domain] * len(curr_domain)
320 self.domain = [_update_domain(d, curr_domain[i]) for i, d in enumerate(domain)]
321 else:
322 self.domain = _update_domain(domain, curr_domain)
323 if (dist := self.distribution) is not None and isinstance(dist, Uniform | LogUniform):
324 dist.dist_args = self.domain # keep Uniform dist in sync
326 def pdf(self, x: np.ndarray) -> np.ndarray:
327 """Compute the PDF of the Variable at the given `x` locations. Will just return one's if the variable
328 does not have a distribution.
330 :param x: locations to compute the PDF at
331 :returns: the PDF evaluations at `x`
332 """
333 if dist := self.distribution:
334 return dist.pdf(x)
335 else:
336 return np.ones(x.shape) # No pdf if no dist is specified
338 def sample(self, shape: tuple | int, nominal: float | np.ndarray = None) -> np.ndarray:
339 """Draw samples from this `Variable's` distribution. Just returns the nominal value of the given shape if
340 this `Variable` has no distribution.
342 :param shape: the shape of the returned samples
343 :param nominal: a nominal value to use if applicable (i.e. a center for relative, tolerance, or normal)
344 :returns: samples from the PDF of this `Variable's` distribution
345 """
346 if isinstance(shape, int):
347 shape = (shape, )
348 if nominal is None:
349 nominal = self.get_nominal()
351 if dist := self.distribution:
352 return dist.sample(shape, nominal)
353 else:
354 # Variable's with no distribution
355 if nominal is None:
356 raise ValueError(f'Cannot sample "{self.name}" with no dist or nominal value specified.')
357 elif isinstance(nominal, list | np.ndarray):
358 return np.ones(shape + (len(nominal),)) * np.atleast_1d(nominal) # For field quantities
359 else:
360 return np.ones(shape) * nominal
362 def normalize(self, values: ArrayLike, denorm: bool = False) -> ArrayLike | None:
363 """Normalize `values` based on this `Variable's` `norm` method(s). See `Transform` for available norm methods.
365 !!! Note
366 If this Variable's `self.norm` was specified as a list of norm methods, then each will be applied in
367 sequence in the original order (and in reverse for `denorm=True`). When `self.distribution` is involved in
368 the transforms (only for `minmax` and `zscore`), the `dist_args` will get normalized too at each
369 transform before applying the next transform.
371 :param values: the values to normalize (array-like)
372 :param denorm: whether to denormalize instead using the inverse of the original normalization method
373 :returns: the normalized (or unnormalized) values
374 """
375 if not self.norm or values is None:
376 return values
377 if dist := self.distribution:
378 normal_dist = isinstance(dist, Normal)
379 else:
380 normal_dist = False
382 def _normalize_single(values, transform, inverse, domain, dist_args):
383 """Do a single transform. Might need to override transform_args depending on the transform."""
384 transform_args = None
385 if isinstance(transform, Minmax) and domain:
386 transform_args = domain + transform.transform_args[2:] # Update minmax bounds
387 elif isinstance(transform, Zscore) and dist_args:
388 transform_args = dist_args # Update N(mu, std)
390 return transform.transform(values, inverse=inverse, transform_args=transform_args)
392 domain = self.get_domain() or ()
393 dist_args = self.distribution.dist_args if normal_dist else []
394 if isinstance(domain, list):
395 domain = () # For field quantities, domain is not used in normalization
397 if denorm:
398 # First, send domain and dist_args through the forward norm list (up until the last norm)
399 hyperparams = [np.hstack((domain, dist_args))]
400 for i, transform in enumerate(self.norm):
401 domain, dist_args = tuple(hyperparams[i][:2]), tuple(hyperparams[i][2:])
402 hyperparams.append(_normalize_single(hyperparams[i], transform, False, domain, dist_args))
404 # Now denormalize in reverse
405 hp_idx = -2
406 for transform in reversed(self.norm):
407 domain, dist_args = tuple(hyperparams[hp_idx][:2]), tuple(hyperparams[hp_idx][2:])
408 values = _normalize_single(values, transform, True, domain, dist_args)
409 hp_idx -= 1
410 else:
411 # Normalize values and hyperparams through the forward norm list
412 hyperparams = np.hstack((domain, dist_args))
413 for transform in self.norm:
414 domain, dist_args = tuple(hyperparams[:2]), tuple(hyperparams[2:])
415 values = _normalize_single(values, transform, denorm, domain, dist_args)
416 hyperparams = _normalize_single(hyperparams, transform, denorm, domain, dist_args)
418 return values
420 def denormalize(self, values):
421 """Alias for `normalize(denorm=True)`. See `normalize` for more details."""
422 return self.normalize(values, denorm=True)
424 def compress(self, values: CompressionData, coords: np.ndarray = None,
425 reconstruct: bool = False) -> CompressionData:
426 """Compress or reconstruct field quantity values using this `Variable's` compression info.
428 !!! Note "Specifying compression values"
429 If only one field quantity is associated with this variable, then
430 specify `values` as `dict(coords=..., name=...)` for this Variable's `name`. If `coords` is not specified,
431 then this assumes the locations are the same as the reconstruction data (and skips interpolation).
433 !!! Info "Compression workflow"
434 Generally, compression follows `interpolate -> normalize -> compress` to take raw values into the compressed
435 "latent" space. The interpolation step is required to make sure `values` align with the coordinates used
436 when building the compression map in the first place (such as through SVD).
438 :param values: a `dict` with a key for each field qty of shape `(..., qty.shape)` and a `coords` key of shape
439 `(qty.shape, dim)` that gives the coordinates of each point. Only a single `latent` key should
440 be given instead if `reconstruct=True`.
441 :param coords: the coordinates of each point in `values` if `values` did not contain a `coords` key;
442 defaults to the compression grid coordinates
443 :param reconstruct: whether to reconstruct values instead of compress
444 :returns: the compressed values with key `latent` and shape `(..., latent_size)`; if `reconstruct=True`,
445 then the reconstructed values with shape `(..., qty.shape)` for each `qty` key are returned.
446 The return `dict` also has a `coords` key with shape `(qty.shape, dim)`.
447 """
448 if not self.compression:
449 raise ValueError(f'Compression is not supported for variable "{self.name}". Please specify a compression'
450 f' method for this variable.')
451 if not self.compression.map_exists:
452 raise ValueError(f'Compression map not computed yet for "{self.name}".')
454 # Default field coordinates to the compression coordinates if they are not provided
455 field_coords = values.pop('coords', coords)
456 if field_coords is None:
457 field_coords = self.compression.coords
458 ret_dict = {'coords': field_coords}
460 # For reconstruction: decompress -> denormalize -> interpolate
461 if reconstruct:
462 try:
463 states = np.atleast_1d(values['latent']) # (..., rank)
464 except KeyError as e:
465 raise ValueError('Must pass values["latent"] in for reconstruction.') from e
466 states = self.compression.reconstruct(states) # (..., dof)
467 states = self.denormalize(states) # (..., dof)
468 states = self.compression.interpolate_from_grid(states, field_coords)
469 ret_dict.update(states)
471 # For compression: interpolate -> normalize -> compress
472 else:
473 states = self.compression.interpolate_to_grid(field_coords, values)
474 states = self.normalize(states) # (..., dof)
475 states = self.compression.compress(states) # (..., rank)
476 ret_dict['latent'] = states
478 return ret_dict
480 def reconstruct(self, values, coords=None):
481 """Alias for `compress(reconstruct=True)`. See `compress` for more details."""
482 return self.compress(values, coords=coords, reconstruct=True)
484 def serialize(self, save_path: str | Path = '.') -> dict:
485 """Convert a `Variable` to a `dict` with only standard Python types
486 (i.e. convert custom objects like `dist` and `norm` to strings and save `compression` to a `.pkl`).
488 :param save_path: the path to save the compression data to (defaults to current directory)
489 :returns: the serialized `dict` of the `Variable` object
490 """
491 d = {}
492 for key, value in self.__dict__.items():
493 if value is not None and not key.startswith('_'):
494 if key == 'domain':
495 d[key] = [str(v) for v in value] if isinstance(value, list) else str(value)
496 elif key == 'distribution':
497 d[key] = str(value)
498 elif key == 'norm':
499 d[key] = [str(transform) for transform in value]
500 elif key == 'compression':
501 fname = f'{self.name}_compression.pkl'
502 d[key] = value.serialize(save_path=Path(save_path) / fname)
503 else:
504 d[key] = value
505 return d
507 @classmethod
508 def deserialize(cls, data: dict, search_paths: list[str | Path] = None) -> Variable:
509 """Convert a `dict` to a `Variable` object. Let `pydantic` handle validation and conversion of fields.
511 :param data: the `dict` to convert to a `Variable`
512 :param search_paths: the paths to search for compression files (if necessary)
513 :returns: the `Variable` object
514 """
515 if isinstance(data, Variable):
516 return data
517 elif isinstance(data, str):
518 return cls(name=data)
519 else:
520 if (compression := data.get('compression', None)) is not None:
521 if isinstance(compression, str):
522 data['compression'] = search_for_file(compression, search_paths=search_paths)
523 return cls(**data)
525 @staticmethod
526 def _yaml_representer(dumper: yaml.Dumper, data: Variable) -> yaml.MappingNode:
527 """Convert a single `Variable` object (`data`) to a yaml MappingNode (i.e. a `dict`)."""
528 save_path, save_file = _get_yaml_path(dumper)
529 return dumper.represent_mapping(Variable.yaml_tag, data.serialize(save_path=save_path))
531 @staticmethod
532 def _yaml_constructor(loader: yaml.Loader, node):
533 """Convert the `!Variable` tag in yaml to a single `Variable` object (or a list of `Variables`)."""
534 save_path, save_file = _get_yaml_path(loader)
535 if isinstance(node, yaml.SequenceNode):
536 return [ele if isinstance(ele, Variable) else Variable.deserialize(ele, search_paths=[save_path]) for ele in
537 loader.construct_sequence(node, deep=True)]
538 elif isinstance(node, yaml.MappingNode):
539 return Variable.deserialize(loader.construct_mapping(node), search_paths=[save_path])
540 else:
541 raise NotImplementedError(f'The "{Variable.yaml_tag}" yaml tag can only be used on a yaml sequence or '
542 f'mapping, not a "{type(node)}".')
545class VariableList(OrderedDict, Serializable):
546 """Store `Variables` as `str(var) : Variable` in the order they were passed in. You can:
548 - Initialize/update from a single `Variable` or a list of `Variables`
549 - Get/set a `Variable` directly or by name via `my_vars[var]` or `my_vars[str(var)]` etc.
550 - Retrieve the original order of insertion by `list(my_vars.items())`
551 - Access/delete elements by order of insertion using integer/slice indexing (i.e. `my_vars[1:3]`)
552 - Save/load from yaml file using the `!VariableList` tag
553 """
554 yaml_tag = '!VariableList'
556 def __init__(self, data: list[Variable] | Variable | OrderedDict | dict = None, **kwargs):
557 """Initialize a collection of `Variable` objects."""
558 super().__init__()
559 self.update(data, **kwargs)
561 def __iter__(self):
562 yield from self.values()
564 def __eq__(self, other):
565 if isinstance(other, VariableList):
566 for v1, v2 in zip(self.values(), other.values()):
567 if v1 != v2:
568 return False
569 return True
570 else:
571 return False
573 def append(self, data: Variable):
574 self.update(data)
576 def extend(self, data: list[Variable]):
577 self.update(data)
579 def index(self, key):
580 for i, k in enumerate(self.keys()):
581 if k == key:
582 return i
583 raise ValueError(f"'{key}' is not in list")
585 def get_domains(self, norm: bool = True):
586 """Get normalized variable domains (expand latent coefficient domains for field quantities). Assume a
587 domain of `(0, 1)` for variables if their domain is not specified.
589 :param norm: whether to normalize the domains using `Variable.norm` (useful for getting bds for surrogate);
590 latent coefficient domains do not get normalized
591 :returns: a `dict` of variables to their normalized domains; field quantities return a domain for each
592 of their latent coefficients
593 """
594 domains = {}
595 for var in self:
596 var_domain = var.get_domain()
597 if isinstance(var_domain, list): # only field qtys return a list of domains, one for each latent coeff
598 for i, domain in enumerate(var_domain):
599 domains[f'{var.name}{LATENT_STR_ID}{i}'] = domain
600 elif var_domain is None:
601 domains[var.name] = (0, 1)
602 else:
603 domains[var.name] = var.normalize(var_domain) if norm else var_domain
604 return domains
606 def get_pdfs(self, norm: bool = True):
607 """Get callable pdfs for all variables (skipping field quantities for now)
609 :param norm: whether values passed to the pdf functions are normalized and should be denormed first
610 before pdf evaluation (useful for surrogate construction where samples are gathered in the
611 normalized space)
612 :returns: a `dict` of variables to callable pdf functions; field quantities are skipped.
613 """
614 def _get_pdf(var, norm):
615 return lambda z: var.pdf(var.denormalize(z) if norm else z)
617 pdf_fcns = {}
618 for var in self:
619 var_domain = var.get_domain()
620 if isinstance(var_domain, list): # only field qtys return a list of domains, one for each latent coeff
621 # for i, domain in enumerate(var_domain):
622 # pdf_fcns[f'{var.name}{LATENT_STR_ID}{i}'] = var.latent_pdfs[i] TODO: Implement latent pdfs
623 pass
624 else:
625 pdf_fcns[var.name] = _get_pdf(var, norm)
626 return pdf_fcns
628 def update(self, data: list[Variable | str] | str | Variable | OrderedDict | dict = None, **kwargs):
629 """Update from a list or dict of `Variable` objects, or from `key=value` pairs."""
630 if data:
631 if isinstance(data, OrderedDict | dict):
632 for key, value in data.items():
633 self.__setitem__(key, value)
634 else:
635 data = [data] if not isinstance(data, list | tuple) else data
636 for variable in data:
637 self.__setitem__(str(variable), variable)
638 if kwargs:
639 for key, value in kwargs.items():
640 self.__setitem__(key, value)
642 def get(self, key, default=None):
643 """Make sure this passes through `__getitem__()`"""
644 try:
645 return self.__getitem__(key)
646 except Exception:
647 return default
649 def __setitem__(self, key, value):
650 """Only allow `str(var): Variable` items. Or normal list indexing via `my_vars[0] = var`."""
651 if isinstance(key, int):
652 k = list(self.keys())[key]
653 self.__setitem__(k, value)
654 return
655 if isinstance(value, str):
656 value = Variable(name=value)
657 if not isinstance(key, str | Variable):
658 raise TypeError(f'VariableList key "{key}" is not a Variable or string.')
659 if not isinstance(value, Variable):
660 raise TypeError(f'VariableList value "{value}" is not a Variable.')
661 super().__setitem__(str(key), value)
663 def __getitem__(self, key):
664 """Allow accessing variable(s) directly via `my_vars[var]` or by index/slicing."""
665 if isinstance(key, list | tuple):
666 return [self.__getitem__(ele) for ele in key]
667 elif isinstance(key, int | slice):
668 return list(self.values())[key]
669 elif isinstance(key, str | Variable):
670 return super().__getitem__(str(key))
671 else:
672 raise TypeError(f'VariableList key "{key}" is not valid.')
674 def __delitem__(self, key):
675 """Allow deleting variable(s) directly or by index/slicing."""
676 if isinstance(key, list | tuple):
677 for ele in key:
678 self.__delitem__(ele)
679 elif isinstance(key, int | slice):
680 ele = list(self.keys())[key]
681 if isinstance(ele, list):
682 for item in ele:
683 super().__delitem__(item)
684 else:
685 super().__delitem__(ele)
686 elif isinstance(key, str | Variable):
687 super().__delitem__(str(key))
688 else:
689 raise TypeError(f'VariableList key "{key}" is not valid.')
691 def __str__(self):
692 return str(list(self.values()))
694 def __repr__(self):
695 return self.__str__()
697 def serialize(self, save_path='.') -> list[dict]:
698 """Convert to a list of `dict` objects for each `Variable` in the list.
700 :param save_path: the path to save the compression data to (defaults to current directory)
701 """
702 return [var.serialize(save_path=save_path) for var in self.values()]
704 @classmethod
705 def merge(cls, *variable_lists) -> VariableList:
706 """Merge multiple sets of variables into a single `VariableList` object.
708 !!! Note
709 Variables with the same name will be merged by keeping the one with the most information provided.
711 :param variable_lists: the variables/lists to merge
712 :returns: the merged `VariableList` object
713 """
714 merged_vars = cls()
716 def _get_best_variable(var1, var2):
717 var1_dict = {key: value for key, value in var1.__dict__.items() if value is not None}
718 var2_dict = {key: value for key, value in var2.__dict__.items() if value is not None}
719 return var1 if len(var1_dict) >= len(var2_dict) else var2
721 for var_list in variable_lists:
722 for var in cls(var_list):
723 if var.name in merged_vars:
724 merged_vars[var.name] = _get_best_variable(merged_vars[var.name], var)
725 else:
726 merged_vars[var.name] = var
728 return merged_vars
730 @classmethod
731 def deserialize(cls, data: dict | list[dict], search_paths=None) -> VariableList:
732 """Convert a `dict` or list of `dict` objects to a `VariableList` object. Let `pydantic` handle validation."""
733 if not isinstance(data, list):
734 data = [data]
735 return cls([Variable.deserialize(d, search_paths=search_paths) for d in data])
737 @staticmethod
738 def _yaml_representer(dumper: yaml.Dumper, data: VariableList) -> yaml.SequenceNode:
739 """Convert a single `VariableList` object (`data`) to a yaml SequenceNode (i.e. a list)."""
740 save_path, save_file = _get_yaml_path(dumper)
741 return dumper.represent_sequence(VariableList.yaml_tag, data.serialize(save_path=save_path))
743 @staticmethod
744 def _yaml_constructor(loader: yaml.Loader, node):
745 """Convert the `!VariableList` tag in yaml to a `VariableList` object."""
746 save_path, save_file = _get_yaml_path(loader)
747 if isinstance(node, yaml.SequenceNode):
748 return VariableList.deserialize(loader.construct_sequence(node, deep=True), search_paths=[save_path])
749 elif isinstance(node, yaml.MappingNode):
750 return VariableList.deserialize(loader.construct_mapping(node), search_paths=[save_path])
751 else:
752 raise NotImplementedError(f'The "{VariableList.yaml_tag}" yaml tag can only be used on a yaml sequence or '
753 f'mapping, not a "{type(node)}".')