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

1import typing 

2from typing import TYPE_CHECKING 

3 

4import freegs 

5import numpy as np 

6import numpy.typing as npt 

7from gsfit_rs import BpProbes 

8 

9if TYPE_CHECKING: 

10 from . import DatabaseReader 

11 

12 

13def setup_bp_probes( 

14 self: "DatabaseReader", 

15 pulseNo: int, 

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

17 time: npt.NDArray[np.float64], 

18 freegs_eqs: list[freegs.equilibrium.Equilibrium], 

19) -> BpProbes: 

20 """ 

21 This method initialises the Rust `BpProbes` class. 

22 

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

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

25 :param time: Measured time vector 

26 :param freegs_eqs: List of FreeGS equilibrium objects, one for each time-slice 

27 

28 **This method is specific to FreeGS.** 

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 n_time = len(time) 

37 

38 # Let's see how many flux loops there are 

39 n_bp_probes = 0 

40 for sensor in freegs_eqs[0].tokamak.sensors: 

41 if isinstance(sensor, freegs.machine.PoloidalFieldSensor): 

42 n_bp_probes += 1 

43 

44 # Collect the measured values 

45 measured_values = np.full((n_time, n_bp_probes), np.nan) 

46 for i_time in range(n_time): 

47 i_sensor = 0 

48 for sensor in freegs_eqs[i_time].tokamak.sensors: 

49 if isinstance(sensor, freegs.machine.PoloidalFieldSensor): 

50 measured_values[i_time, i_sensor] = sensor.measurement 

51 i_sensor += 1 

52 

53 # Add the sensors to the BpProbes class 

54 i_sensor = 0 

55 for sensor in freegs_eqs[0].tokamak.sensors: 

56 if isinstance(sensor, freegs.machine.PoloidalFieldSensor): 

57 # Get the sensor name 

58 sensor_name = sensor.name 

59 

60 # Get the sensor position 

61 r = sensor.R 

62 z = sensor.Z 

63 angle = sensor.theta 

64 

65 # Get the sensor data 

66 measured = sensor.measurement 

67 

68 # Only add sesnors which are included in the settings file 

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

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

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

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

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

74 else: 

75 fit_settings_comment = "" 

76 fit_settings_expected_value = np.nan 

77 fit_settings_include = False 

78 fit_settings_weight = np.nan 

79 

80 # Add the sensor to the BpProbes class 

81 bp_probes.add_sensor( 

82 name=sensor_name, 

83 geometry_angle_pol=angle, 

84 geometry_r=r, 

85 geometry_z=z, 

86 fit_settings_comment=fit_settings_comment, 

87 fit_settings_expected_value=fit_settings_expected_value, 

88 fit_settings_include=fit_settings_include, 

89 fit_settings_weight=fit_settings_weight, 

90 time=time, 

91 measured=measured_values[:, i_sensor], 

92 ) 

93 

94 # Increment the sensor index 

95 i_sensor += 1 

96 

97 return bp_probes