Coverage for python/gsfit/database_readers/st40_spider_mdsplus/setup_plasma.py: 0%

1import typing 

2from typing import TYPE_CHECKING 

3 

4import gsfit_rs 

5import numpy as np 

6import numpy.typing as npt 

7from gsfit_rs import Plasma 

8from st40_database import GetData 

9 

10if TYPE_CHECKING: 

11 from . import DatabaseReader 

12 

13 

14def setup_plasma( 

15 self: "DatabaseReader", 

16 pulseNo: int, 

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

18) -> Plasma: 

19 """ 

20 This method initialises the Rust `Plasma` 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 # Initial plasma conditions 

31 initial_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"] 

32 initial_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["r_cur"] 

33 initial_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["z_cur"] 

34 

35 # Set the source functions types 

36 p_prime_source_function: gsfit_rs.EfitPolynomial 

37 ff_prime_source_function: gsfit_rs.EfitPolynomial 

38 

39 if settings["source_function_p_prime.json"]["method"] == "efit_polynomial": 

40 n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"] 

41 regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"]) 

42 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0). 

43 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error 

44 if regularisations.shape == (1, 0): 

45 regularisations = np.zeros((0, n_dof), dtype=np.float64) 

46 p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

47 else: 

48 raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}") 

49 

50 if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial": 

51 n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"] 

52 regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"]) 

53 # If `regularisations` is [[]] in the json file, will be interpreted by numpy as having size (1, 0). 

54 # Which would be interpreted as (n_regularisations, n_dof). So it would cause an error 

55 if regularisations.shape == (1, 0): 

56 regularisations = np.zeros((0, n_dof), dtype=np.float64) 

57 ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

58 

59 # Grid size and shape 

60 n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"] 

61 n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"] 

62 r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"] 

63 r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"] 

64 z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"] 

65 z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"] 

66 

67 # Normalised poloidal flux grid 

68 n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"] 

69 psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64) 

70 

71 # Limiter 

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

73 elmag = GetData(13321, f"ELMAG#{elmag_run_name}", is_fail_quiet=False) 

74 limit_pts_r = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.R")) 

75 limit_pts_z = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.Z")) 

76 

77 # Vacuum vessel where the plasma is allowed to be 

78 vessel_r = limit_pts_r 

79 vessel_z = limit_pts_z 

80 

81 # Add lower MC tiles 

82 limit_pts_r = np.append(limit_pts_r, 0.7103) 

83 limit_pts_z = np.append(limit_pts_z, -0.3131) 

84 # Add upper MC tiles 

85 limit_pts_r = np.append(limit_pts_r, 0.7103) 

86 limit_pts_z = np.append(limit_pts_z, 0.3031) 

87 

88 # Initialise the Plasma Rust class 

89 plasma = Plasma( 

90 n_r, 

91 n_z, 

92 r_min, 

93 r_max, 

94 z_min, 

95 z_max, 

96 psi_n, # BUXTON: perhaps better to send in `n_psi_n` 

97 limit_pts_r, 

98 limit_pts_z, 

99 vessel_r, 

100 vessel_z, 

101 p_prime_source_function, 

102 ff_prime_source_function, 

103 initial_ip, 

104 initial_cur_r, 

105 initial_cur_z, 

106 ) 

107 

108 return plasma 

109 

110 

111# import typing 

112# from typing import TYPE_CHECKING 

113 

114# import gsfit_rs 

115# import numpy as np 

116# import numpy.typing as npt 

117# from gsfit_rs import Plasma 

118# from st40_database import GetData # type: ignore[import-not-found] 

119 

120# if TYPE_CHECKING: 

121# from . import DatabaseReader 

122 

123 

124# def setup_plasma( 

125# self: "DatabaseReader", 

126# pulseNo: int, 

127# settings: dict[str, typing.Any], 

128# ) -> Plasma: 

129# """ 

130# This method initialises the Rust `Plasma` class. 

131 

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

133# :param settings: Dictionary containing the JSON settings read from the `settings` directory 

134 

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

136 

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

138# """ 

139 

140# # Initial plasma conditions 

141# initial_ip = settings["GSFIT_code_settings.json"]["initial_guess"]["ip"] 

142# initial_cur_r = settings["GSFIT_code_settings.json"]["initial_guess"]["r_cur"] 

143# initial_cur_z = settings["GSFIT_code_settings.json"]["initial_guess"]["z_cur"] 

144 

145# # Set the source functions types 

146# p_prime_source_function: gsfit_rs.EfitPolynomial 

147# ff_prime_source_function: gsfit_rs.EfitPolynomial 

148 

149# if settings["source_function_p_prime.json"]["method"] == "efit_polynomial": 

150# n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"] 

151# regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"]) 

152# p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

153# elif settings["source_function_p_prime.json"]["method"] == "efit_polynomial": 

154# n_dof = settings["source_function_p_prime.json"]["efit_polynomial"]["n_dof"] 

155# regularisations = np.array(settings["source_function_p_prime.json"]["efit_polynomial"]["regularizations"]) 

156# p_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

157# else: 

158# raise ValueError(f"Unknown method for p_prime source function: {settings['source_function_p_prime.json']['method']}") 

159 

160# if settings["source_function_ff_prime.json"]["method"] == "efit_polynomial": 

161# n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"] 

162# regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"]) 

163# ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

164# elif settings["source_function_ff_prime.json"]["method"] == "efit_polynomial": 

165# n_dof = settings["source_function_ff_prime.json"]["efit_polynomial"]["n_dof"] 

166# regularisations = np.array(settings["source_function_ff_prime.json"]["efit_polynomial"]["regularizations"]) 

167# ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

168 

169# # Grid size and shape 

170# n_r = settings["GSFIT_code_settings.json"]["grid"]["n_r"] 

171# n_z = settings["GSFIT_code_settings.json"]["grid"]["n_z"] 

172# r_min = settings["GSFIT_code_settings.json"]["grid"]["r_min"] 

173# r_max = settings["GSFIT_code_settings.json"]["grid"]["r_max"] 

174# z_min = settings["GSFIT_code_settings.json"]["grid"]["z_min"] 

175# z_max = settings["GSFIT_code_settings.json"]["grid"]["z_max"] 

176 

177# # Normalised poloidal flux grid 

178# n_psi_n = settings["GSFIT_code_settings.json"]["n_psi_n"] 

179# psi_n = np.linspace(0.0, 1.0, n_psi_n).astype(np.float64) 

180 

181# # Limiter 

182# elmag_run_name = settings["GSFIT_code_settings.json"]["database_reader"]["st40_astra_mdsplus"]["elmag_run_name"] 

183# elmag = GetData(12050, f"ELMAG#{elmag_run_name}", is_fail_quiet=False) 

184# limit_pts_r = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.R")) 

185# limit_pts_z = typing.cast(npt.NDArray[np.float64], elmag.get("LIMITER.LIMIT_PTS.Z")) 

186 

187# # Vacuum vessel where the plasma is allowed to be 

188# vessel_r = limit_pts_r 

189# vessel_z = limit_pts_z 

190 

191# # Add lower MC tiles 

192# limit_pts_r = np.append(limit_pts_r, 0.7103) 

193# limit_pts_z = np.append(limit_pts_z, -0.3131) 

194# # Add upper MC tiles 

195# limit_pts_r = np.append(limit_pts_r, 0.7103) 

196# limit_pts_z = np.append(limit_pts_z, 0.3031) 

197 

198# # Initialise the Plasma Rust class 

199# plasma = Plasma( 

200# n_r, 

201# n_z, 

202# r_min, 

203# r_max, 

204# z_min, 

205# z_max, 

206# psi_n, # BUXTON: perhaps better to send in `n_psi_n` 

207# limit_pts_r, 

208# limit_pts_z, 

209# vessel_r, 

210# vessel_z, 

211# p_prime_source_function, 

212# ff_prime_source_function, 

213# initial_ip, 

214# initial_cur_r, 

215# initial_cur_z, 

216# ) 

217 

218# return plasma