Coverage for python/gsfit/database_readers/mock_st40_mdsplus/mock_get_data.py: 0%

1import json 

2import typing 

3from pathlib import Path 

4 

5import numpy as np 

6import numpy.typing as npt 

7 

8# Files are named `{FILE_PREFIX}_{tree_name}_{pulseNo}.npz`, with the tree name lowercased 

9FILE_PREFIX = "mdsplus_mock" 

10 

11# Archive member holding the capture provenance; not an MDSplus node 

12METADATA_KEY = "__metadata__" 

13 

14 

15class MockGetData: 

16 """ 

17 Drop-in replacement for `st40_database.GetData`, backed by mocked MDSplus tree files. 

18 

19 Each (tree, pulse) which the `st40_mdsplus` reader would have fetched over the network is 

20 mocked by a single `.npz`, mirroring how MDSplus itself stores one tree per shot: 

21 `mdsplus_mock_mag_12050.npz` 

22 

23 Each archive member is named for the bare node path (e.g. `BPPROBE.P101.B`), because the 

24 tree and pulse are already carried by the filename. Capture provenance is stored inside the 

25 same archive under `__metadata__`, so it cannot drift away from the data it describes; read 

26 it with `get_metadata()`, or from a shell with: 

27 python -c "import numpy as np; print(np.load('<file>.npz')['__metadata__'])" 

28 

29 Which tree, pulse and run each read maps to comes from the `workflow` settings, exactly as 

30 it does for the live `st40_mdsplus` reader; use `from_workflow(...)` rather than naming a 

31 tree directly. The requested run is checked against the run the mock was captured from, so 

32 asking for a run which was never captured is an error rather than silently wrong data. 

33 

34 The files are produced by `investigation/capture_snapshot.py`. 

35 """ 

36 

37 # Each tree file is read once and shared by all readers 

38 _cache: dict[tuple[str, int, str], typing.Any] = {} 

39 

40 def __init__(self, mock_dir: str, pulseNo: int, tree_run: str) -> None: 

41 """ 

42 :param mock_dir: Directory containing the `mdsplus_mock_*.npz` files 

43 :param pulseNo: Pulse number, used to select which mocked tree file to read 

44 :param tree_run: MDSplus tree name, optionally with a `#run_name` suffix 

45 """ 

46 

47 self.pulseNo = pulseNo 

48 self.tree_name = tree_run.split("#")[0] 

49 self.run_name = tree_run.split("#")[1] if "#" in tree_run else "" 

50 self._nodes = MockGetData._load(mock_dir, pulseNo, self.tree_name) 

51 

52 # Guard against reading a mock which was captured from a different run 

53 captured_run_name = self.get_metadata()["run_name"] 

54 if self.run_name != "" and captured_run_name != "" and self.run_name != captured_run_name: 

55 raise ValueError( 

56 f"mock_st40_mdsplus: {self.tree_name}#{self.run_name} was requested for pulseNo={pulseNo}, " 

57 f"but the mock was captured from {self.tree_name}#{captured_run_name}" 

58 ) 

59 

60 @classmethod 

61 def from_workflow(cls, settings: dict[str, typing.Any], pulseNo: int, workflow_name: str) -> "MockGetData": 

62 """ 

63 Open the mocked tree which `workflow_name` maps to. 

64 

65 :param settings: Dictionary containing the JSON settings read from the `settings` directory 

66 :param pulseNo: The shot's pulse number, used when the workflow entry does not pin its own 

67 :param workflow_name: Key into `["database_reader"]["mock_st40_mdsplus"]["workflow"]` 

68 """ 

69 

70 reader_settings = settings["GSFIT_code_settings.json"]["database_reader"]["mock_st40_mdsplus"] 

71 workflow = reader_settings["workflow"] 

72 

73 if workflow_name not in workflow: 

74 raise KeyError(f"mock_st40_mdsplus: no `workflow` entry named '{workflow_name}'; the settings define {sorted(workflow)}") 

75 workflow_entry = workflow[workflow_name] 

76 

77 # `pulseNo = null` means "use the shot's pulse"; the machine-description reads pin their own 

78 tree_pulse_no = pulseNo if workflow_entry["pulseNo"] is None else workflow_entry["pulseNo"] 

79 

80 return cls(str(reader_settings["mock_dir"]), tree_pulse_no, f"{workflow_entry['tree_name']}#{workflow_entry['run_name']}") 

81 

82 @classmethod 

83 def _load(cls, mock_dir: str, pulseNo: int, tree_name: str) -> typing.Any: 

84 cache_key = (mock_dir, pulseNo, tree_name) 

85 if cache_key not in cls._cache: 

86 npz_path = Path(mock_dir) / f"{FILE_PREFIX}_{tree_name.lower()}_{pulseNo}.npz" 

87 if not npz_path.exists(): 

88 raise FileNotFoundError(f"mock_st40_mdsplus: no mocked MDSplus tree for pulseNo={pulseNo}, tree={tree_name}; expected {npz_path}") 

89 cls._cache[cache_key] = np.load(npz_path) 

90 return cls._cache[cache_key] 

91 

92 def get(self, node: str) -> typing.Any: 

93 """ 

94 :param node: MDSplus node path, relative to the top of the tree, e.g. `"BPPROBE.ALL.NAMES"` 

95 """ 

96 

97 if node == METADATA_KEY or node not in self._nodes: 

98 raise KeyError(f"mock_st40_mdsplus: node not found in mocked tree {self.tree_name}#{self.pulseNo}: {node}") 

99 

100 value = typing.cast(npt.NDArray[typing.Any], self._nodes[node]) 

101 

102 # `st40_database.GetData` hands back Python lists of `str` for string nodes, never numpy 

103 # arrays, so match that here. `.tolist()` also un-nests, e.g. for `PF.ALL.COILS` 

104 if value.dtype.kind == "U": 

105 return value.tolist() 

106 

107 return value 

108 

109 def get_metadata(self) -> dict[str, typing.Any]: 

110 """Return the provenance recorded when this tree was captured.""" 

111 

112 return typing.cast(dict[str, typing.Any], json.loads(str(self._nodes[METADATA_KEY])))