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

11 

12from .astra_flux_loop_reader import astra_flux_loop_reader 

13 

14if TYPE_CHECKING: 

15 from . import DatabaseReader 

16 

17 

18def setup_flux_loops( 

19 self: "DatabaseReader", 

20 pulseNo: int, 

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

22) -> FluxLoops: 

23 """ 

24 This method initialises the Rust `FluxLoops` class. 

25 

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

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

28 

29 **This method is specific to ST40's ASTRA stored on MDSplus.** 

30 

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

32 """ 

33 

34 # Initialise the FluxLoops Rust class 

35 flux_loops = FluxLoops() 

36 

37 # Extract the astra_run_name from settings 

38 astra_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_astra_mdsplus"]["workflow"]["astra"]["run_name"] 

39 

40 # Connect to MDSplus 

41 conn = mdsthin.Connection("smaug") 

42 conn.openTree("ASTRA", pulseNo) 

43 

44 # ASTRA flux_loops 

45 time = conn.get(f"\\ASTRA::TOP.{astra_run_name}:TIME").data().astype(np.float64) 

46 measurements = conn.get(f"\\ASTRA::TOP.{astra_run_name}.FLOOP.ALL:PSI").data().astype(np.float64) 

47 names = conn.get(f"\\ASTRA::TOP.{astra_run_name}.FLOOP.ALL:NAME").data().astype(str) 

48 

49 # Read flux loop geometry from fl_loop.dat 

50 flux_loops_data = astra_flux_loop_reader() 

51 

52 for flux_loop_name, data in flux_loops_data.items(): 

53 sensor_name = flux_loop_name.replace("FLOOP_", "L") 

54 

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

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

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

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

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

60 else: 

61 fit_settings_comment = "" 

62 fit_settings_expected_value = np.nan 

63 fit_settings_include = False 

64 fit_settings_weight = np.nan 

65 

66 # Measured values 

67 index = names == flux_loop_name 

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

69 

70 # Add the sensor to the Rust class 

71 flux_loops.add_sensor( 

72 name=sensor_name, 

73 geometry_r=data["r"], 

74 geometry_z=data["z"], 

75 fit_settings_comment=fit_settings_comment, 

76 fit_settings_expected_value=fit_settings_expected_value, 

77 fit_settings_include=fit_settings_include, 

78 fit_settings_weight=fit_settings_weight, 

79 time=time, 

80 measured=measured, 

81 ) 

82 

83 return flux_loops