Coverage for python/gsfit/database_readers/st40_spider_mdsplus/setup_flux_loops.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 FluxLoops 

11from st40_database import GetData 

12 

13if TYPE_CHECKING: 

14 from . import DatabaseReader 

15 

16 

17def setup_flux_loops( 

18 self: "DatabaseReader", 

19 pulseNo: int, 

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

21) -> FluxLoops: 

22 """ 

23 This method initialises the Rust `FluxLoops` 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 FluxLoops Rust class 

34 flux_loops = FluxLoops() 

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 flux_loops 

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}.FLOOP.ALL:PSI").data().astype(np.float64) 

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

47 

48 # read in mag data from MDSplus 

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

50 

51 # FL-probes 

52 names_long = mag.get("FLOOP.ALL.NAMES") 

53 sensors_names = np.char.replace(names_long, "FLOOP_", "L") 

54 sensors_r = mag.get("FLOOP.ALL.R") 

55 sensors_z = mag.get("FLOOP.ALL.Z") 

56 

57 n_sensors = len(sensors_names) 

58 for i_sensor in range(0, n_sensors): 

59 sensor_name = sensors_names[i_sensor] 

60 if sensor_name in settings["sensor_weights_flux_loops.json"]: 

61 fit_settings_comment = settings["sensor_weights_flux_loops.json"][sensor_name]["fit_settings"]["comment"] 

62 fit_settings_expected_value = settings["sensor_weights_flux_loops.json"][sensor_name]["fit_settings"]["expected_value"] 

63 fit_settings_include = settings["sensor_weights_flux_loops.json"][sensor_name]["fit_settings"]["include"] 

64 fit_settings_weight = settings["sensor_weights_flux_loops.json"][sensor_name]["fit_settings"]["weight"] / (2.0 * np.pi) 

65 else: 

66 fit_settings_comment = "" 

67 fit_settings_expected_value = np.nan 

68 fit_settings_include = False 

69 fit_settings_weight = np.nan 

70 

71 # Measured values 

72 index = names == sensor_name.replace("L", "FLOOP_") 

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

74 

75 # Add the sensor to the Rust class 

76 flux_loops.add_sensor( 

77 name=sensors_names[i_sensor], 

78 geometry_r=sensors_r[i_sensor], 

79 geometry_z=sensors_z[i_sensor], 

80 fit_settings_comment=fit_settings_comment, 

81 fit_settings_expected_value=fit_settings_expected_value, 

82 fit_settings_include=fit_settings_include, 

83 fit_settings_weight=fit_settings_weight, 

84 time=time, 

85 measured=measured, 

86 ) 

87 

88 return flux_loops