Coverage for src/amisc/typing.py: 100%

1"""Module with type hints for the AMISC package. 

2 

3Includes: 

4 

5- `MultiIndex` — tuples of integers or similar string representations 

6- `Dataset` — a type hint for the input/output `dicts` handled by `Component.model` 

7- `TrainIteration` — the results of a single training iteration 

8- `CompressionData` — a dictionary spec for passing data to/from `Variable.compress()` 

9- `LATENT_STR_ID` — a string identifier for latent coefficients of field quantities 

10""" 

11import ast as _ast 

12from pathlib import Path as _Path 

13from typing import Optional as _Optional 

14 

15import numpy as _np 

16from numpy.typing import ArrayLike as _ArrayLike 

17from typing_extensions import TypedDict as _TypedDict 

18 

19__all__ = ["MultiIndex", "Dataset", "TrainIteration", "CompressionData", "LATENT_STR_ID"] 

20 

21LATENT_STR_ID = "_LATENT" # String identifier for latent variables 

22 

23 

24class MultiIndex(tuple): 

25 """A multi-index is a tuple of integers, can be converted from a string.""" 

26 def __new__(cls, __tuple=()): 

27 if isinstance(__tuple, str): 

28 return super().__new__(cls, map(int, _ast.literal_eval(__tuple))) 

29 else: 

30 return super().__new__(cls, map(int, __tuple)) 

31 

32 

33class Dataset(_TypedDict, total=False): 

34 """Type hint for the input/output `dicts` of a call to `Component.model`. The keys are the variable names and the 

35 values are the corresponding `np.ndarrays`. There are also a few special keys that can be returned by the model 

36 that are described below. 

37 

38 The model can return additional items that are not part of `Component.outputs`. These items are returned as object 

39 arrays in the output. 

40 

41 This data structure is very similar to the `Dataset` class in the `xarray` package. Later versions might 

42 consider migrating to `xarray` for more advanced data manipulation. 

43 

44 :ivar model_cost: the computational cost (seconds of CPU time) of a single model evaluation 

45 :ivar output_path: the path to the output file or directory written by the model 

46 :ivar errors: a `dict` with the indices where the model evaluation failed with context about the errors 

47 """ 

48 model_cost: float | list | _ArrayLike 

49 output_path: str | _Path 

50 errors: dict 

51 

52 

53class TrainIteration(_TypedDict): 

54 """Gives the results of a single training iteration. 

55 

56 :ivar component: the name of the component selected for refinement at this iteration 

57 :ivar alpha: the selected candidate model fidelity multi-index 

58 :ivar beta: the selected candidate surrogate fidelity multi-index 

59 :ivar num_evals: the number of model evaluations performed during this iteration 

60 :ivar added_cost: the total added computational cost of the new model evaluations 

61 :ivar added_error: the error/difference between the refined surrogate and the previous surrogate 

62 :ivar test_error: the error of the refined surrogate on the test set (optional) 

63 """ 

64 component: str 

65 alpha: MultiIndex 

66 beta: MultiIndex 

67 num_evals: int 

68 added_cost: float 

69 added_error: float 

70 test_error: _Optional[dict[str, float]] 

71 

72 

73class CompressionData(_TypedDict, total=False): 

74 """Configuration `dict` for passing compression data to/from `Variable.compress()`. 

75 

76 !!! Info "Field quantity shapes" 

77 Field quantity data can take on any arbitrary shape, which we indicate with `qty.shape`. For example, a 3d 

78 structured grid might have `qty.shape = (10, 15, 10)`. Unstructured data might just have `qty.shape = (N,)` 

79 for $N$ points in an unstructured grid. Regardless, `Variable.compress()` will flatten this and compress 

80 to a single latent vector of size `latent_size`. That is, `qty.shape` → `latent_size`. 

81 

82 !!! Note "Compression coordinates" 

83 Field quantity data must be specified along with its coordinate locations. If the coordinate locations are 

84 different from what was used when building the compression map (i.e. the SVD data matrix), then they will be 

85 interpolated to/from the SVD coordinates. 

86 

87 :ivar coords: `(qty.shape, dim)` the coordinate locations of the qty data; coordinates exist in `dim` space (e.g. 

88 `dim=2` for 2d Cartesian coordinates). Defaults to the coordinates used when building the construction 

89 map (i.e. the coordinates of the data in the SVD data matrix) 

90 :ivar latent: `(..., latent_size)` array of latent space coefficients for a field quantity; this is what is 

91 _returned_ by `Variable.compress()` and what is _expected_ as input by `Variable.reconstruct()`. 

92 :ivar qty: `(..., qty.shape)` array of uncompressed field quantity data for this qty within 

93 the `fields` list. Each qty in this list will be its own `key:value` pair in the 

94 `CompressionData` structure 

95 """ 

96 coords: _np.ndarray 

97 latent: _np.ndarray 

98 qty: _np.ndarray