Coverage for python/gsfit/database_readers/freegsnke/setup_coils.py: 100%

1import typing 

2from typing import TYPE_CHECKING 

3 

4import freegs4e 

5import numpy as np 

6import numpy.typing as npt 

7from freegsnke.equilibrium_update import Equilibrium as FreeGsnkeEquilibrium 

8from gsfit_rs import Coils 

9 

10if TYPE_CHECKING: 

11 from . import DatabaseReader 

12 

13 

14def setup_coils( 

15 self: "DatabaseReader", 

16 pulseNo: int, 

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

18 time: npt.NDArray[np.float64], 

19 freegsnke_eqs: list[FreeGsnkeEquilibrium], 

20) -> Coils: 

21 """ 

22 This method initialises the Rust `Coils` class. 

23 

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

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

26 :param time: Measured time vector 

27 :param freegsnke_eqs: List of FreeGSNKE equilibrium objects, one for each time-slice 

28 

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

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 coils class 

35 coils = Coils() 

36 

37 # We assume that the static data is not changing in time 

38 freegsnke_tokamak = freegsnke_eqs[0].tokamak 

39 

40 # Get lengths 

41 n_time = len(time) 

42 

43 # Count the number of PF coils 

44 n_pf = 0 

45 for current_carrying_object in freegsnke_tokamak.coils: 

46 # `type(current_carrying_object) = (str, freegs4e.machine.Circuit) = (circuit_name, circuit)`. 

47 # `current_carrying_object` can be either a PF coil or passive 

48 circuit_name = current_carrying_object[0] 

49 circuit = current_carrying_object[1] 

50 if isinstance(circuit, freegs4e.machine.Circuit): 

51 pf_coils = circuit.coils 

52 n_pf = n_pf + len(pf_coils) 

53 

54 # Loop over time and get the currents 

55 pf_coil_currents = np.full((n_time, n_pf), np.nan) 

56 for i_time in range(0, n_time): 

57 i_pf = 0 

58 for current_carrying_object in freegsnke_eqs[i_time].tokamak.coils: 

59 circuit_name = current_carrying_object[0] 

60 circuit = current_carrying_object[1] 

61 circuit_current = circuit.current 

62 if isinstance(circuit, freegs4e.machine.Circuit): 

63 pf_coils = circuit.coils 

64 for pf_coil in pf_coils: 

65 # `type(pf_coils) = list[(str, freegs4e.multi_coil.MultiCoil)]` 

66 coil_name = pf_coil[0] 

67 coil_object = pf_coil[1] 

68 

69 # Store the current 

70 # Note, there is a recent bug in FreeGSNKE/FreeGS4E where the `coil_object.current` is not updated, and currents are all zero. 

71 # pf_coil_currents[i_time, i_pf] = coil_object.current 

72 pf_coil_currents[i_time, i_pf] = circuit_current 

73 

74 # set index for the next PF coil 

75 i_pf = i_pf + 1 

76 

77 # Add the PF coils 

78 i_pf = 0 

79 for current_carrying_object in freegsnke_tokamak.coils: 

80 circuit_name = current_carrying_object[0] 

81 circuit = current_carrying_object[1] 

82 if isinstance(circuit, freegs4e.machine.Circuit): 

83 pf_coils = circuit.coils 

84 for pf_coil in pf_coils: 

85 coil_name = pf_coil[0] 

86 coil_object = pf_coil[1] 

87 

88 coil_r = np.array(coil_object.Rfil).astype(np.float64) 

89 coil_z = np.array(coil_object.Zfil).astype(np.float64) 

90 coil_d_r = 0.0 * coil_r 

91 coil_d_z = 0.0 * coil_z 

92 

93 current = coil_object.current 

94 

95 coils.add_pf_coil( 

96 name=coil_name, 

97 r=coil_r, 

98 z=coil_z, 

99 d_r=coil_d_r, 

100 d_z=coil_d_z, 

101 time=time, 

102 measured=pf_coil_currents[:, i_pf], 

103 ) 

104 

105 # set index for the next PF coil 

106 i_pf = i_pf + 1 

107 

108 # Add TF coil 

109 coils.add_tf_coil( 

110 time=np.array([0.0, 1.0]), 

111 measured=np.array([1.0e6, 1.0e6]), 

112 ) 

113 

114 return coils