Coverage for python/gsfit/database_readers/st40_spider_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 

9import numpy.typing as npt 

10from gsfit_rs import BpProbes 

11from st40_database import GetData 

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 SPIDER 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 spider_run_name from settings 

37 spider_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_spider_mdsplus"]["workflow"]["spider"]["run_name"] 

38 

39 # Connect to MDSplus 

40 conn = mdsthin.Connection("smaug") 

41 conn.openTree("SPIDER", pulseNo) 

42 

43 # SPIDER bp_probes 

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

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

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

47 

48 # Read mag data from MDSplus 

49 # FIXME: using a fixed shot is not a good idea! 

50 mag = GetData(12050, "MAG#BEST", is_fail_quiet=False) 

51 

52 # Bp-probes 

53 names_long = mag.get("BPPROBE.ALL.NAMES") 

54 sensors_names = np.char.replace(names_long, "BPPROBE_", "P") 

55 sensors_r = mag.get("BPPROBE.ALL.R") 

56 sensors_z = mag.get("BPPROBE.ALL.Z") 

57 sensors_angle_pol = mag.get("BPPROBE.ALL.THETA") 

58 # names_long = conn.get(f"\\SPIDER::TOP.{spider_run_name}.BPPROBE.ALL:NAME").data().astype(str) 

59 # sensors_names = np.char.replace(names_long, "BPPROBE_", "P") 

60 # sensors_r = conn.get(f"\\SPIDER::TOP.{spider_run_name}.BPPROBE.ALL:R").data().astype(np.float64) 

61 # sensors_z = conn.get(f"\\SPIDER::TOP.{spider_run_name}.BPPROBE.ALL:Z").data().astype(np.float64) 

62 # sensors_angle_pol = conn.get(f"\\SPIDER::TOP.{spider_run_name}.BPPROBE.ALL:ANGLE").data().astype(np.float64) 

63 

64 n_sensors = len(sensors_names) 

65 for i_sensor in range(0, n_sensors): 

66 sensor_name = sensors_names[i_sensor] 

67 

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

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

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

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

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

73 else: 

74 fit_settings_comment = "" 

75 fit_settings_expected_value = np.nan 

76 fit_settings_include = False 

77 fit_settings_weight = np.nan 

78 

79 # Measured values 

80 index = names == sensor_name.replace("P", "BPPROBE_") 

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

82 

83 # Add the sensor to the Rust class 

84 bp_probes.add_sensor( 

85 name=sensor_name, 

86 geometry_angle_pol=sensors_angle_pol[i_sensor], 

87 geometry_r=sensors_r[i_sensor], 

88 geometry_z=sensors_z[i_sensor], 

89 fit_settings_comment=fit_settings_comment, 

90 fit_settings_expected_value=fit_settings_expected_value, 

91 fit_settings_include=fit_settings_include, 

92 fit_settings_weight=fit_settings_weight, 

93 time=time, 

94 measured=measured, 

95 ) 

96 

97 return bp_probes