Coverage for python/gsfit/database_readers/freegsnke/setup_bp_probes.py: 86%

1import typing 

2from typing import TYPE_CHECKING 

3 

4import numpy as np 

5import numpy.typing as npt 

6from freegsnke.equilibrium_update import Equilibrium as FreeGsnkeEquilibrium 

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 freegsnke_eqs: list[FreeGsnkeEquilibrium], 

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 freegsnke_eqs: List of FreeGSNKE equilibrium objects, one for each time-slice 

27 

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

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 # We assume that the static data is not changing in time 

37 freegsnke_tokamak = freegsnke_eqs[0].tokamak 

38 

39 # Get lengths 

40 n_time = len(time) 

41 n_sensors = len(freegsnke_tokamak.probes.pickups) 

42 

43 # Loop over time and calculate the pickup values 

44 pickups_vals = np.full((n_time, n_sensors), np.nan) 

45 for i_time in range(n_time): 

46 pickups_vals[i_time, :] = freegsnke_tokamak.probes.calculate_pickup_value(freegsnke_eqs[i_time]) 

47 

48 # Loop over sensors and add them to the BpProbes class 

49 for i_sensor, sensor in enumerate(freegsnke_tokamak.probes.pickups): 

50 # Only use sensors in the poloidal plane 

51 if sensor["orientation"] == "PARALLEL": 

52 # Get the sensor name 

53 sensor_name = sensor["name"] 

54 

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

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

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

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

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

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

61 else: 

62 fit_settings_comment = "" 

63 fit_settings_expected_value = np.nan 

64 fit_settings_include = False 

65 fit_settings_weight = np.nan 

66 

67 # Get the sensor position 

68 r = sensor["position"][0] 

69 z = sensor["position"][2] 

70 

71 orientation_vector = sensor["orientation_vector"] 

72 orientation_r = orientation_vector[0] 

73 orientation_z = orientation_vector[2] 

74 if orientation_r == 0.0 and orientation_z == 0.0: 

75 raise ValueError(f"Sensor {sensor_name} has zero orientation vector. Perhaps it's pointing in the toroidal direction?") 

76 angle = np.arctan2(orientation_z, orientation_r) 

77 

78 # Get the sensor data, all times 

79 measured = pickups_vals[:, i_sensor] 

80 

81 # Add the sensor to the BpProbes class 

82 bp_probes.add_sensor( 

83 name=sensor_name, 

84 geometry_angle_pol=angle, 

85 geometry_r=r, 

86 geometry_z=z, 

87 fit_settings_comment=fit_settings_comment, 

88 fit_settings_expected_value=fit_settings_expected_value, 

89 fit_settings_include=fit_settings_include, 

90 fit_settings_weight=fit_settings_weight, 

91 time=time, 

92 measured=measured, 

93 ) 

94 

95 return bp_probes