Coverage for python/gsfit/database_readers/st40_mdsplus/setup_dialoop.py: 0%

1import typing 

2from typing import TYPE_CHECKING 

3 

4import mdsthin 

5import numpy as np 

6import numpy.typing as npt 

7from gsfit_rs import Dialoop 

8from st40_database import GetData 

9 

10if TYPE_CHECKING: 

11 from . import DatabaseReader 

12 

13 

14def setup_dialoop( 

15 self: "DatabaseReader", 

16 pulseNo: int, 

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

18) -> Dialoop: 

19 """ 

20 This method initialises the Rust `Dialoop` class. 

21 

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

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

24 

25 **This method is specific to ST40's experimental MDSplus database.** 

26 

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

28 """ 

29 

30 # Initialise the Dialoop Rust class 

31 dialoop = Dialoop() 

32 

33 # We use a single diamagnetic flux loop 

34 sensor_name = "DIALOOP" 

35 

36 if sensor_name in settings["sensor_weights_dialoop.json"]: 

37 fit_settings_comment = settings["sensor_weights_dialoop.json"][sensor_name]["fit_settings"]["comment"] 

38 fit_settings_expected_value = settings["sensor_weights_dialoop.json"][sensor_name]["fit_settings"]["expected_value"] 

39 fit_settings_include = settings["sensor_weights_dialoop.json"][sensor_name]["fit_settings"]["include"] 

40 fit_settings_weight = settings["sensor_weights_dialoop.json"][sensor_name]["fit_settings"]["weight"] 

41 else: 

42 fit_settings_comment = "" 

43 fit_settings_expected_value = np.nan 

44 fit_settings_include = False 

45 fit_settings_weight = np.nan 

46 

47 # If the diamagnetic flux loop is not being fitted, skip reading the DIALOOP (and MAG) 

48 # trees entirely and return an empty `Dialoop` object. 

49 if not fit_settings_include: 

50 return dialoop 

51 

52 # Read in the diamagnetic flux from MDSplus 

53 dialoop_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_mdsplus"]["workflow"]["dialoop"]["run_name"] 

54 mag_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_mdsplus"]["workflow"]["mag"]["run_name"] 

55 dialoop_data = GetData(pulseNo, f"DIALOOP#{dialoop_run_name}", is_fail_quiet=False) 

56 

57 # Get geometry (path of integration) directly via mdsthin from the MAG tree 

58 # (the geometry is stored in MAG, not in the DIALOOP tree) 

59 conn = mdsthin.Connection("smaug") 

60 conn.openTree("MAG", pulseNo) 

61 path_r = conn.get(f"\\MAG::TOP.{mag_run_name}.DIALOOP.L000:R_PATH").data().astype(np.float64) 

62 path_z = conn.get(f"\\MAG::TOP.{mag_run_name}.DIALOOP.L000:Z_PATH").data().astype(np.float64) 

63 

64 # Measured signal: \DIALOOP::TOP.<run_name>.GLOBAL:PHI_DIA 

65 time = typing.cast(npt.NDArray[np.float64], dialoop_data.get("TIME")).astype(np.float64) 

66 measured = typing.cast(npt.NDArray[np.float64], dialoop_data.get("GLOBAL.PHI_DIA")).astype(np.float64) 

67 

68 # Add the sensor to the Rust class 

69 dialoop.add_sensor( 

70 name=sensor_name, 

71 r=path_r, 

72 z=path_z, 

73 fit_settings_comment=fit_settings_comment, 

74 fit_settings_expected_value=fit_settings_expected_value, 

75 fit_settings_include=fit_settings_include, 

76 fit_settings_weight=fit_settings_weight, 

77 time=time, 

78 measured=measured, 

79 ) 

80 

81 return dialoop