Coverage for src/hallmd/utils.py: 89%

1"""Module to provide utilities for the `hallmd` package. 

2 

3Includes: 

4 

5- `load_device()` - Load a device configuration from the `hallmd.devices` directory. 

6""" 

7import json 

8import os 

9from importlib import resources 

10from pathlib import Path 

11 

12import yaml 

13 

14__all__ = ['load_device'] 

15 

16AVOGADRO_CONSTANT = 6.02214076e23 # Avogadro constant (mol^-1) 

17FUNDAMENTAL_CHARGE = 1.602176634e-19 # Fundamental charge (C) 

18BOLTZMANN_CONSTANT = 1.380649e-23 # Boltzmann constant (J/K) 

19TORR_2_PA = 133.322 # Conversion factor from Torr to Pa 

20MOLECULAR_WEIGHTS = { 

21 'Xenon': 131.293, 

22 'Argon': 39.948, 

23 'Krypton': 83.798, 

24 'Bismuth': 208.98, 

25 'Hydrogen': 1.008, 

26 'Mercury': 200.59 

27} 

28 

29 

30def _path_in_dict(value, data: dict) -> list: 

31 """Recursively check if a value is in a dictionary and return the list of access keys to the value.""" 

32 if isinstance(data, dict): 

33 for key, v in data.items(): 

34 path = _path_in_dict(value, v) 

35 if path: 

36 return [key] + path 

37 elif data == value: 

38 return [value] # found the value 

39 return [] # value not found 

40 

41 

42def load_device(device_name: str, device_file: str = 'device.yml', device_dir: str | Path = None) -> dict: 

43 """Load a device configuration from the `device_dir` directory. The `device_file` must be located at 

44 `device_dir/device_name/device_file`. All other files in the directory, if referenced in `device_file`, will 

45 be converted to an absolute path. 

46 

47 !!! Example "Loading a device configuration" 

48 Currently, the only provided device configuration is for the SPT-100 thruster. 

49 ```python 

50 from hallmd.utils import load_device 

51 

52 device = load_device('SPT-100') 

53 ``` 

54 

55 You may put custom configurations in the `hallmd.devices` directory or specify a custom directory with a custom 

56 configuration file: 

57 ```yaml 

58 name: MyDevice 

59 geometry: 

60 channel_length: 1 

61 inner_radius: 2 

62 outer_radius: 3 

63 magnetic_field: 

64 file: bfield.csv 

65 shielded: false 

66 ``` 

67 

68 :param device_name: name of the device configuration to load 

69 :param device_file: name of the device configuration file (default: 'device.yml'). Only supported file types are 

70 `.yml` and `.json`. 

71 :param device_dir: directory containing the devices. If None, the `hallmd.devices` directory is used. 

72 :return: dictionary containing the device configuration 

73 """ 

74 device_dir = resources.files('hallmd.devices') if device_dir is None else Path(device_dir) 

75 if not (device_dir / device_name).exists(): 

76 raise FileNotFoundError(f'Device directory "{device_name}" not found in the device folder.') 

77 if not (device_dir / device_name / device_file).exists(): 

78 raise FileNotFoundError(f'Device configuration file "{device_file}" not found in the "{device_name}" ' 

79 f'directory. Please rename or specify the configuration file as "{device_file}".') 

80 

81 config_file = device_dir / device_name / device_file 

82 with open(config_file, 'r', encoding='utf-8') as fd: 

83 if config_file.suffix == '.yml': 

84 config = yaml.safe_load(fd) 

85 elif config_file.suffix == '.json': 

86 config = json.load(fd) 

87 else: 

88 raise ValueError(f'Unsupported file type "{config_file.suffix}". Only .yml and .json files are supported.') 

89 

90 # Convert all relative paths to absolute paths 

91 for root, _, files in os.walk(device_dir / device_name): 

92 for file in files: 

93 if file != device_file: 

94 # Check if the posix file path from root is in the config (i.e. "./file.csv") 

95 root_path = Path(root) / file # Path like "hallmd/devices/SPT-100/path/to/file.csv" 

96 rel_path = root_path.relative_to(device_dir / device_name) # Just the path/to/file.csv part (relative) 

97 dict_path = _path_in_dict(rel_path.as_posix(), config) 

98 if len(dict_path) == 0: 

99 # Check if the plain filename is in the config (i.e. file.csv); will only pick first match 

100 dict_path = _path_in_dict(file, config) 

101 

102 if dict_path: 

103 d = config # pointer to the nested location in config 

104 for key in dict_path[:-2]: 

105 d = config[key] 

106 d[dict_path[-2]] = root_path.resolve().as_posix() 

107 

108 return config