Coverage for python/diagnostic_and_simulation_base/diagnostic_and_simulation_base.py: 59%

1import json 

2import logging 

3import sys 

4import textwrap 

5import typing 

6from importlib import resources as resources_py # BUXTON: this is deprecated and needs removing 

7from pathlib import Path 

8from time import time as time_py 

9 

10import f90nml 

11 

12from . import version_storage 

13from .nested_dictionary import NestedDict 

14from .utility import logging as util_logging 

15from .utility.make_settings_json import make_settings_json 

16 

17 

18class DiagnosticAndSimulationBase: 

19 """Class for creating Diagnostic Analysis codes, such as: `Gas`, `Efit`, `Ppts`, `Zeff_pi`""" 

20 

21 log_records: list[logging.LogRecord] = [] 

22 

23 @property 

24 def log_string(self) -> str: 

25 return util_logging.format_logs(self.log_records) 

26 

27 def __init__( 

28 self, 

29 pulseNo: int, 

30 run_name: str, 

31 run_description: str = "Standard run with default settings", 

32 settings_path: str = "default", 

33 write_to_mds: bool = True, 

34 pulseNo_write: int | None = None, 

35 analysis_name: str | None = None, 

36 link_run_to_best: bool = False, 

37 ) -> None: 

38 """Class constructor 

39 

40 :param pulseNo: pulse number 

41 :param run_name: run_name to save to MDSplus 

42 :param run_description: help string for MDSplus Tree 

43 :param settings_path: location where code inputs are stored 

44 :param write_to_mds: flag to turn on / off writing to MDSplus 

45 :param pulseNo_write: pulse number in which data is to be written, if different from the current pulse 

46 :param analysis_name: for example "EFITP" is analysis using the "EFIT code" 

47 :param link_run_to_best: BEST will be linked to the run 

48 """ 

49 self.analysis_time_start = time_py() 

50 

51 # code_name: is the package name from "pip list" in upper case (note, this includes "_") 

52 code_name = self.__class__.__module__.split(".")[0].upper() # BUXTON: I don't like this!! 

53 

54 # analysis_name 

55 if analysis_name is None: 

56 analysis_name = code_name 

57 

58 # Create a logger with the module name 

59 # Get the logging_level 

60 logging_level = util_logging.logger.getEffectiveLevel() 

61 

62 # Create logger 

63 self.logger = logging.getLogger(analysis_name) 

64 # Don't propagate messages to the root logger; `mdsthin` uses the root logger and this avoids double logging 

65 self.logger.propagate = False 

66 # Set the logging level 

67 self.logger.setLevel(level=logging_level) 

68 

69 # Check if the logger already has handlers to avoid adding duplicates 

70 if not self.logger.handlers: 

71 # Create and add handler to emit the logs to "stadnard out", 

72 # i.e. printing to the terminal (standard behaviour) 

73 handler_to_standard_out = logging.StreamHandler(sys.stdout) 

74 handler_to_standard_out.setFormatter(util_logging.CustomFormatter()) 

75 self.logger.addHandler(handler_to_standard_out) 

76 

77 # Create and add handler to emit the log "records" into `self.log_records` 

78 handler_to_list = util_logging.EmitLogToListHandler(self.log_records) 

79 self.logger.addHandler(handler_to_list) 

80 

81 # Store inputs 

82 self.pulseNo = pulseNo 

83 self.code_name = code_name 

84 self.analysis_name = analysis_name 

85 self.run_name = run_name.upper() 

86 self.run_description = run_description 

87 self.write_to_mds = write_to_mds 

88 self.link_run_to_best = link_run_to_best 

89 

90 # Determine which pulseNo to write to and store in class object 

91 if pulseNo_write is None: 

92 pulseNo_write = pulseNo 

93 self.pulseNo_write = pulseNo_write 

94 

95 # Test if "settings_path" is a directory. 

96 # If directory doesn't exist, then treat as "relative path" 

97 if not (Path(settings_path).is_dir()): 

98 python_module = self.code_name.lower() 

99 settings_path = f"{resources_py.files(python_module)}/settings/{settings_path}" 

100 

101 # Test if settings directory exists 

102 if not Path(settings_path).is_dir(): 

103 raise FileNotFoundError("settings directory not found") 

104 

105 # Store the resolved settings path 

106 self.settings_path = settings_path 

107 

108 # Create empty settings dictionary 

109 self.settings = {} # type: dict[str, typing.Any] 

110 

111 self._load_settings_from_files() 

112 

113 # Create results dictionary, with the stuff we already know 

114 # Important: results is a 1:1 mapping to MDSplus data-strucutre 

115 # TO-DO: Change this to a pre-populated dictionary from the *.csv file 

116 self.results = NestedDict() 

117 self.results["CODE_VERSION"]["COMPUTER"] = version_storage.__computer__ 

118 self.results["CODE_VERSION"]["DATETIME"] = version_storage.__datetime__ 

119 self.results["CODE_VERSION"]["GIT_ID"] = version_storage.__git_short_hash__ 

120 self.results["CODE_VERSION"]["LIBRARY"] = version_storage.__python_library__ 

121 self.results["CODE_VERSION"]["PYTHON"] = version_storage.__python__ 

122 self.results["CODE_VERSION"]["USER"] = version_storage.__user__ 

123 self.results["CODE_VERSION"]["VERSION"] = version_storage.__version__ 

124 

125 def __repr__(self) -> str: 

126 """Print to screen""" 

127 string_output = "" 

128 string_output += "╔═════════════════════════════════════════════════════════════════════════════╗\n" 

129 string_output += f"║ {f' <{self.__class__.__name__}>':<75} ║\n" 

130 string_output += f"║ {f' {version_storage.__version__}':<75} ║\n" 

131 string_output += f"║ {' ':<75} ║\n" 

132 string_output += f"║ {' pulseNo = ' + f'{self.pulseNo:_}':<75} ║\n" 

133 string_output += f"║ {' pulseNo_write = ' + f'{self.pulseNo_write:_}':<75} ║\n" 

134 string_output += f"║ {' run_name = ' + str(self.run_name):<75} ║\n" 

135 string_output += f"║ {' run_description = ' + str(self.run_description):<75} ║\n" 

136 string_output += f"║ {' settings_path = ...':<75} ║\n" 

137 return string_output 

138 

139 def _load_settings_from_files(self) -> None: 

140 """Look in the "settings_path" directory and load all settings files. 

141 Will recursively load *.json and *.nml files 

142 

143 TODO: add *.csv reader #SUNDAR - is ths even needed? 

144 """ 

145 

146 settings_path = self.settings_path 

147 

148 # Load *.json settings, including sub-directories 

149 for file in Path(settings_path).glob("**/*.json"): 

150 with open(file, "r") as file_id: 

151 step_name = f'Loading settings from: "{file.name}"' 

152 try: 

153 relative_path = str(file.relative_to(settings_path)) 

154 self.settings[relative_path] = json.load(file_id) 

155 self.logger.info(msg=step_name) 

156 except Exception as exception_obj: 

157 self.logger.exception(msg=step_name) 

158 raise exception_obj 

159 

160 # Load *.nml (namelist) settings, including sub-directories 

161 for file in Path(settings_path).glob("**/*.nml"): 

162 with open(file, "r") as file_id: 

163 step_name = f"Loading settings from {file.name}" 

164 try: 

165 relative_path = str(file.relative_to(settings_path)) 

166 self.settings[relative_path] = f90nml.read(file_id) 

167 self.logger.info(msg=step_name) 

168 except Exception as exception_obj: 

169 self.logger.exception(msg=step_name) 

170 raise exception_obj 

171 

172 def _write_to_mds(self) -> None: 

173 """Writes data to MDSplus""" 

174 

175 # Lazy loading of `standard_utility` because it's specific to Tokamak Energy. 

176 import standard_utility as util 

177 

178 # Add settings files to results. 

179 # We do this right at the end, as they can be programatially changed, e.g. for scans 

180 self.results["INPUT"]["SETTINGS"] = make_settings_json(data=self.settings, json_indent=2) 

181 

182 # Create MDSplus nodes 

183 mdsplus_settings_file = f"{self.analysis_name}_mdsplus_settings.json" 

184 if mdsplus_settings_file in self.settings: 

185 pulseNo_cal = self.settings[mdsplus_settings_file].get("calibration", {}).get("pulse", None) 

186 else: 

187 pulseNo_cal = None 

188 

189 try: 

190 workflow = list(self.results["INPUT"]["WORKFLOW"].keys()) 

191 except KeyError: 

192 workflow = None 

193 

194 util.create_script_nodes( 

195 script_name=self.analysis_name, 

196 pulseNo_write=self.pulseNo_write, 

197 pulseNo_cal=pulseNo_cal, 

198 run_name=self.run_name, 

199 run_info=self.run_description, 

200 workflows=workflow, 

201 link_best=self.link_run_to_best, 

202 ) 

203 

204 # Write to MDSplus 

205 util.write_script_data( 

206 script_name=self.analysis_name, 

207 pulseNo_write=self.pulseNo_write, 

208 data_to_write=self.results.to_dictionary(), 

209 pulseNo_cal=pulseNo_cal, 

210 run_name=self.run_name, 

211 run_description=self.run_description, 

212 force_write=True, 

213 ) 

214 

215 def run_with_log( 

216 self, 

217 operator: typing.Callable[..., typing.Any], 

218 message: str | None, 

219 args: dict[typing.Any, typing.Any] | None = None, 

220 ) -> typing.Any: 

221 if message is None: 

222 step_name = "" 

223 else: 

224 step_name = message 

225 

226 try: 

227 if args is None: 

228 ret_data = operator() 

229 else: 

230 ret_data = operator(**args) 

231 self.logger.info(msg=step_name) 

232 return ret_data 

233 except Exception as exception_obj: 

234 self.logger.exception(msg=step_name) 

235 raise exception_obj