Coverage for include / utils.py: 61%

1import torch 

2import ruamel.yaml as yaml 

3yamlc = yaml.YAML(typ="safe") 

4 

5from typing import Union 

6 

7""" General utility functions """ 

8 

9def load_config(path) -> tuple[dict, dict]: 

10 """ Loads the configuration and returns it together with the model-specific configuration 

11 

12 :param path: 

13 :return: 

14 """ 

15 

16 with open(path) as cfg_file: 

17 cfg = yamlc.load(cfg_file) 

18 model_name = cfg.get("root_model_name", "SIR") 

19 model_cfg = cfg[model_name] 

20 

21 return cfg, model_cfg 

22 

23def set_default_device(cfg) -> torch.device: 

24 """ Sets the default pytorch training device and number of threads. 

25 By default, the cpu is used, since pytorch may not have been installed with Cuda (or Metal) enabled. 

26 

27 :param cfg: config file 

28 :return: device 

29 """ 

30 device = cfg.get("device", 'cpu') 

31 torch.set_default_device(device) 

32 num_threads = cfg.get("num_threads", None) 

33 if num_threads is not None: 

34 torch.set_num_threads(num_threads) 

35 return torch.device(device) 

36 

37def random_tensor( 

38 cfg: Union[dict, list], *, size: tuple = None, device: str = "cpu", **__ 

39) -> torch.Tensor: 

40 """Generates a multi-dimensional random tensor. Each entry can be initialised separately, or a common 

41 initialisation configuration is used for each entry. For instance, the configuration 

42 

43 .. code-block:: 

44 

45 cfg: 

46 distribution: uniform 

47 parameters: 

48 lower: 0 

49 upper: 1 

50 

51 together with `size: (2, 2)` will initialise a 2x2 matrix with entries drawn from a uniform distribution on 

52 [0, 1]. The configuration 

53 

54 .. code-block:: 

55 

56 cfg: 

57 - distribution: uniform 

58 parameters: 

59 lower: 0 

60 upper: 1 

61 - distribution: normal 

62 parameters: 

63 mean: 0.5 

64 std: 0.1 

65 

66 will initialise a (2, 1) tensor with entries drawn from different distributions. 

67 

68 :param cfg: the configuration entry containing the initialisation data 

69 :param size (optional): the size of the tensor, in case the configuration is not a list 

70 :param device: the device onto which to load the data 

71 :param __: additional kwargs (ignored) 

72 :return: the tensor of random variables 

73 """ 

74 

75 def _random_tensor_1d( 

76 *, distribution: str, parameters: dict, s: tuple = (1,), **__ 

77 ) -> torch.Tensor: 

78 """Generates a random tensor according to a distribution. 

79 

80 :param distribution: the type of distribution. Can be 'uniform' or 'normal'. 

81 :param parameters: the parameters relevant to the respective distribution 

82 :param s: the size of the random tensor 

83 """ 

84 

85 # Uniform distribution in an interval 

86 if distribution == "uniform": 

87 l, u = parameters.get("lower"), parameters.get("upper") 

88 if l > u: 

89 raise ValueError( 

90 f"Upper bound must be greater or equal to lower bound; got {l} and {u}!" 

91 ) 

92 

93 return torch.tensor((u - l), dtype=torch.float) * torch.rand( 

94 s, dtype=torch.float, device=device 

95 ) + torch.tensor(l, dtype=torch.float) 

96 

97 # Normal distribution 

98 elif distribution == "normal": 

99 return torch.normal( 

100 parameters.get("mean"), 

101 parameters.get("std"), 

102 size=s, 

103 device=device, 

104 dtype=torch.float, 

105 ) 

106 

107 else: 

108 raise ValueError(f"Unrecognised distribution type {distribution}!") 

109 

110 if isinstance(cfg, list): 

111 return torch.tensor([_random_tensor_1d(**entry) for entry in cfg]).to(device) 

112 else: 

113 return _random_tensor_1d(**cfg, s=size).to(device)