Coverage for python/gsfit/database_readers/st40_astra_mdsplus/setup_bp_probes.py: 0%

1# mypy: ignore-errors 

2# TODO: need to fix mypy errors 

3 

4import typing 

5from typing import TYPE_CHECKING 

6 

7import mdsthin 

8import numpy as np 

9from gsfit_rs import BpProbes 

10 

11from .astra_bp_probe_reader import astra_bp_probe_reader 

12 

13if TYPE_CHECKING: 

14 from . import DatabaseReader 

15 

16 

17def setup_bp_probes( 

18 self: "DatabaseReader", 

19 pulseNo: int, 

20 settings: dict[str, typing.Any], 

21) -> BpProbes: 

22 """ 

23 This method initialises the Rust `BpProbes` class. 

24 

25 :param pulseNo: Pulse number, used to read from the database 

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

27 

28 **This method is specific to ST40's ASTRA stored on MDSplus.** 

29 

30 See `python/gsfit/database_readers/interface.py` for more details on how a new database_reader should be implemented. 

31 """ 

32 

33 # Initialise the BpProbes Rust class 

34 bp_probes = BpProbes() 

35 

36 # Extract the astra_run_name from settings 

37 astra_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_astra_mdsplus"]["workflow"]["astra"]["run_name"] 

38 

39 # Connect to MDSplus 

40 conn = mdsthin.Connection("smaug") 

41 conn.openTree("ASTRA", pulseNo) 

42 

43 # ASTRA bp_probes 

44 time = conn.get(f"\\ASTRA::TOP.{astra_run_name}:TIME").data().astype(np.float64) 

45 measurements = conn.get(f"\\ASTRA::TOP.{astra_run_name}.BPPROBE.ALL:B").data().astype(np.float64) 

46 names = conn.get(f"\\ASTRA::TOP.{astra_run_name}.BPPROBE.ALL:NAME").data().astype(str) 

47 

48 # Read Bp probe geometry from pf_probe.dat 

49 bp_probes_data = astra_bp_probe_reader() 

50 

51 for bp_probe_name, data in bp_probes_data.items(): 

52 sensor_name = bp_probe_name.replace("BPPROBE_", "P") 

53 

54 if sensor_name in settings["sensor_weights_bp_probe.json"]: 

55 fit_settings_comment = settings["sensor_weights_bp_probe.json"][sensor_name]["fit_settings"]["comment"] 

56 fit_settings_expected_value = settings["sensor_weights_bp_probe.json"][sensor_name]["fit_settings"]["expected_value"] 

57 fit_settings_include = settings["sensor_weights_bp_probe.json"][sensor_name]["fit_settings"]["include"] 

58 fit_settings_weight = settings["sensor_weights_bp_probe.json"][sensor_name]["fit_settings"]["weight"] 

59 else: 

60 fit_settings_comment = "" 

61 fit_settings_expected_value = np.nan 

62 fit_settings_include = False 

63 fit_settings_weight = np.nan 

64 

65 # Measured values 

66 index = names == bp_probe_name 

67 measured = measurements[:, index].squeeze() 

68 

69 # Add the sensor to the Rust class 

70 bp_probes.add_sensor( 

71 name=sensor_name, 

72 geometry_angle_pol=data["angle_pol"], 

73 geometry_r=data["r"], 

74 geometry_z=data["z"], 

75 fit_settings_comment=fit_settings_comment, 

76 fit_settings_expected_value=fit_settings_expected_value, 

77 fit_settings_include=fit_settings_include, 

78 fit_settings_weight=fit_settings_weight, 

79 time=time, 

80 measured=measured, 

81 ) 

82 

83 return bp_probes