Coverage for python/gsfit/database_readers/st40_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 | gsfit_rs.TensionedCubicBSpline 

37 ff_prime_source_function: gsfit_rs.EfitPolynomial | gsfit_rs.TensionedCubicBSpline 

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 elif settings["source_function_p_prime.json"]["method"] == "tensioned_cubic_b_spline": 

48 regularisations = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["regularizations"]) 

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

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

51 interior_knots = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"]) 

52 n_dof = len(interior_knots) + 4 

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

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

55 interval_tensions = np.array(settings["source_function_p_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"]) 

56 p_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions) 

57 else: 

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

59 

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

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

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

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

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

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

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

67 ff_prime_source_function = gsfit_rs.EfitPolynomial(n_dof, regularisations) 

68 elif settings["source_function_ff_prime.json"]["method"] == "tensioned_cubic_b_spline": 

69 regularisations = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["regularizations"]) 

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

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

72 interior_knots = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interior_knots"]) 

73 n_dof = len(interior_knots) + 4 

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

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

76 interval_tensions = np.array(settings["source_function_ff_prime.json"]["tensioned_cubic_b_spline"]["interval_tensions"]) 

77 ff_prime_source_function = gsfit_rs.TensionedCubicBSpline(regularisations, interior_knots, interval_tensions) 

78 else: 

79 raise ValueError(f"Unknown method for ff_prime source function: {settings['source_function_ff_prime.json']['method']}") 

80 

81 # Grid size and shape 

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

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

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

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

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

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

88 

89 # Normalised poloidal flux grid 

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

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

92 

93 # Limiter 

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

95 elmag = GetData(pulseNo, f"ELMAG#{elmag_run_name}", is_fail_quiet=False) 

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

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

98 

99 # Vacuum vessel where the plasma is allowed to be 

100 vessel_r = limit_pts_r 

101 vessel_z = limit_pts_z 

102 

103 # Add lower MC tiles 

104 limit_pts_r = np.append(limit_pts_r, 0.7103) 

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

106 # Add upper MC tiles 

107 limit_pts_r = np.append(limit_pts_r, 0.7103) 

108 limit_pts_z = np.append(limit_pts_z, 0.3031) 

109 

110 # Initialise the Plasma Rust class 

111 plasma = Plasma( 

112 n_r, 

113 n_z, 

114 r_min, 

115 r_max, 

116 z_min, 

117 z_max, 

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

119 limit_pts_r, 

120 limit_pts_z, 

121 vessel_r, 

122 vessel_z, 

123 p_prime_source_function, 

124 ff_prime_source_function, 

125 initial_ip, 

126 initial_cur_r, 

127 initial_cur_z, 

128 ) 

129 

130 return plasma