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

1import typing 

2from typing import TYPE_CHECKING 

3 

4import numpy as np 

5import numpy.typing as npt 

6from gsfit_rs import Coils 

7from st40_database import GetData 

8 

9if TYPE_CHECKING: 

10 from . import DatabaseReader 

11 

12 

13def setup_coils( 

14 self: "DatabaseReader", 

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

16 pulseNo: int, 

17) -> Coils: 

18 """ 

19 This method initialises the Rust `Coils` class. 

20 

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

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

23 

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

25 

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

27 """ 

28 

29 # Initialise the Coils Rust class 

30 coils = Coils() 

31 

32 elmag_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_mdsplus"]["workflow"]["elmag"]["run_name"] 

33 # elmag = GetData(pulseNo, f"ELMAG#{elmag_run_name}") 

34 elmag = GetData(11012050, f"ELMAG#{elmag_run_name}") 

35 

36 coils_r = typing.cast(npt.NDArray[np.float64], elmag.get("COILS.R")) 

37 coils_z = typing.cast(npt.NDArray[np.float64], elmag.get("COILS.Z")) 

38 coils_d_r = typing.cast(npt.NDArray[np.float64], elmag.get("COILS.DR")) 

39 coils_d_z = typing.cast(npt.NDArray[np.float64], elmag.get("COILS.DZ")) 

40 coil_names = typing.cast(list[str], elmag.get("COILS.COIL_NAMES")) 

41 fils2coils = typing.cast(npt.NDArray[np.bool_], elmag.get("COILS.FILS2COILS") == 1.0) 

42 

43 psu2coil_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_mdsplus"]["workflow"]["psu2coil"]["run_name"] 

44 psu2coil = GetData(pulseNo, f"PSU2COIL#{psu2coil_run_name}", is_fail_quiet=False, use_redis=False) 

45 time = typing.cast(npt.NDArray[np.float64], psu2coil.get("TIME")) 

46 pf_i = typing.cast(npt.NDArray[np.float64], psu2coil.get("PF.ALL.I")) 

47 

48 # Retrieve the coils which are connected to each PSU 

49 coils_connected_to_psus = typing.cast(list[list[str]], psu2coil.get("PF.ALL.COILS")) 

50 

51 n_time, n_psu = pf_i.shape 

52 

53 for i_psu in range(0, n_psu): 

54 current_this_psu = pf_i[:, i_psu] 

55 coils_connected_to_this_psu = coils_connected_to_psus[i_psu] 

56 

57 for coil_name in coils_connected_to_this_psu: 

58 # Skip empty coil names (MDSplus does not allow jagged array) 

59 if coil_name != "": 

60 i_pf = coil_names.index(coil_name) 

61 i_filaments = fils2coils[:, i_pf] 

62 coil_r = coils_r[i_filaments] 

63 coil_z = coils_z[i_filaments] 

64 coil_d_r = coils_d_r[i_filaments] 

65 coil_d_z = coils_d_z[i_filaments] 

66 

67 # Add coil to Rust class 

68 coils.add_pf_coil( 

69 coil_name, 

70 coil_r, 

71 coil_z, 

72 coil_d_r, 

73 coil_d_z, 

74 time=time, 

75 measured=current_this_psu, 

76 ) 

77 

78 # Add TF coil 

79 i_rod = typing.cast(npt.NDArray[np.float64], psu2coil.get("TF.I_ROD")) 

80 coils.add_tf_coil( 

81 time, 

82 i_rod, 

83 ) 

84 

85 return coils