Coverage for tcvx21/grillix_post/components/sources_m.py: 92%

1""" 

2Routines to evaluate the density and temperature source strengths 

3 

4Since the sources were used with fixed positions and widths, these parameters are given 

5default parameters matching the values used. 

6""" 

7import xarray as xr 

8import numpy as np 

9 

10 

11def smoothstep(test_point: float, step_centre: float, step_width: float, order=3): 

12 """ 

13 Returns a hermite smoothstep of a specific order 

14 

15 The step will be exactly 0 for test_point values less than step_centre - step_width / 2 

16 The step will be exactly 1 for test_point values greater than step_centre + step_width / 2 

17 

18 For a step of order A, the (A-1)^th derivative will also be zero at the endpoints 

19 """ 

20 

21 xn = (np.atleast_1d(test_point) - step_centre) / step_width + 0.5 

22 

23 # Ignore the runtime warning that comes from comparing NaN 

24 # NaN will return False in any comparison 

25 with np.errstate(invalid='ignore'): 

26 xn[xn <= 0] = 0 

27 xn[xn >= 1] = 1 

28 

29 if order == 1: 

30 step_values = xn 

31 elif order == 2: 

32 step_values = -2.0 * xn ** 3 + 3.0 * xn ** 2 

33 elif order == 3: 

34 step_values = 6.0 * xn ** 5 - 15.0 * xn ** 4 + 10.0 * xn ** 3 

35 else: 

36 return NotImplemented 

37 

38 if isinstance(test_point, xr.DataArray): 

39 return xr.DataArray(step_values, coords=test_point.coords, dims=test_point.dims, attrs=test_point.attrs) 

40 else: 

41 return step_values 

42 

43 

44def density_source_function(rho, district, norm, source_strength, source_centre=0.915141, source_width=0.083797): 

45 """Annular source for the density, mimicking neutral particle ionisation""" 

46 annular_source = xr.DataArray(source_strength * np.exp(-((rho ** 2 - source_centre ** 2) / source_width) ** 2)) 

47 

48 return xr.DataArray(annular_source.where(district == "CLOSED", 0.0)).assign_attrs(norm=norm.n0 / norm.tau_0) 

49 

50 

51def core_temperature_source_function(rho, district, norm, source_strength, source_centre=0.3, source_width=0.3): 

52 core_source = source_strength * xr.DataArray(1 - smoothstep(rho, source_centre, source_width)) 

53 return xr.DataArray(core_source.where(district == "CLOSED", 0.0)).assign_attrs(norm=norm.Te0 / norm.tau_0) 

54 

55 

56def temperature_source_function(rho, district, norm, density: xr.DataArray, etemp: xr.DataArray, itemp: xr.DataArray, 

57 density_source: xr.DataArray, source_strength, source_centre=0.3, source_width=0.3): 

58 """ 

59 Smooth-step core power injection, mimicking Ohmic power deposition 

60 The temperature source takes into account the power from the density source, to give a constant-power source 

61 """ 

62 core_source = core_temperature_source_function(rho, district, norm, source_strength, source_centre, source_width) 

63 

64 total_source = xr.DataArray(core_source / density - density_source * (etemp + itemp) / density).assign_attrs( 

65 norm=norm.Te0 / norm.tau_0) 

66 core_source = xr.DataArray(core_source / density).assign_attrs( 

67 norm=norm.Te0 / norm.tau_0) 

68 annular_sink = xr.DataArray(-density_source * (etemp + itemp) / density).assign_attrs( 

69 norm=norm.Te0 / norm.tau_0) 

70 

71 return total_source, core_source, annular_sink 

72 

73 

74def _power_source_from_temperature_source(density_source: xr.DataArray, etemp_source: xr.DataArray, 

75 density: xr.DataArray, etemp: xr.DataArray, itemp: xr.DataArray, 

76 norm, kind: str = 'total'): 

77 """ 

78 Returns the power source corresponding to a given density source and temperature source 

79 

80 Can return the sum of the power contributions from the etemp and density sources, or just the etemp_source and 

81 density_source contributions, using the 'kind' parameter 

82 """ 

83 if kind == 'total': 

84 power_source = xr.DataArray(1.5 * (etemp_source * density + density_source * (etemp + itemp))) 

85 elif kind == 'etemp_source': 

86 power_source = xr.DataArray(1.5 * (etemp_source * density)) 

87 elif kind == 'density_source': 

88 power_source = xr.DataArray(1.5 * (density_source * (etemp + itemp))) 

89 else: 

90 raise NotImplementedError(f"Power source kind {kind} not recognised") 

91 

92 return power_source.assign_attrs(norm=norm.n0 * norm.Te0 / norm.tau_0) 

93 

94def _integrated_power_source(grid, norm, density_source, etemp_source, mean_density, mean_etemp, mean_itemp, 

95 kind: str = 'total'): 

96 """ 

97 Computes the integrated power source 

98 

99 Can return the sum of the power contributions from the etemp and density sources, or just the etemp_source and 

100 density_source contributions, using the 'kind' parameter 

101 """ 

102 from tcvx21.grillix_post.observables.integrations_m import axisymmetric_cylindrical_integration 

103 

104 power_source = _power_source_from_temperature_source(density_source, etemp_source, density=mean_density, 

105 etemp=mean_etemp, itemp=mean_itemp, norm=norm, kind=kind) 

106 

107 integrated_power_ = axisymmetric_cylindrical_integration(grid, norm, power_source) 

108 

109 return integrated_power_.to('kW') 

110 

111 

112def integrated_sources(grid, equi, norm, params, snaps): 

113 """ 

114 Compute the integrated particle and power sources 

115 

116 Integrated power = core_power + density_source_power - edge_sink_power 

117 = core_power (since the edge sink and density source powers cancel out) 

118 

119 We check that this is the case, to within 1kW 

120 """ 

121 from tcvx21.grillix_post.observables.integrations_m import axisymmetric_cylindrical_integration 

122 

123 rho = equi.normalised_flux_surface_label(grid.r_s, grid.z_s) 

124 

125 density_source_rate = params['params_srcsnk']['csrcn'] 

126 etemp_source_rate = params['params_srcsnk']['csrcte'] 

127 

128 density_source = density_source_function(rho, grid.districts, norm, density_source_rate) 

129 integrated_density_source = axisymmetric_cylindrical_integration(grid, norm, density_source) 

130 print(f"The integrated particle source rate is {integrated_density_source:4.3e}") 

131 

132 mean_density = grid.shape(snaps.density.mean(dim=('phi', 'tau'))) 

133 mean_etemp = grid.shape(snaps.electron_temp.mean(dim=('phi', 'tau'))) 

134 mean_itemp = grid.shape(snaps.ion_temp.mean(dim=('phi', 'tau'))) 

135 

136 etemp_source, etemp_core_source, etemp_edge_sink = temperature_source_function( 

137 rho=rho, district=grid.districts, norm=norm, 

138 density=mean_density, etemp=mean_etemp, itemp=mean_itemp, 

139 density_source=density_source, source_strength=etemp_source_rate) 

140 

141 # Compute the power injected into the core as the (positive) etemp source 

142 core_power = _integrated_power_source(grid, norm, density_source, etemp_core_source, 

143 mean_density, mean_etemp, mean_itemp, kind='etemp_source') 

144 

145 total_power = _integrated_power_source(grid, norm, density_source, etemp_source, 

146 mean_density, mean_etemp, mean_itemp, kind='total') 

147 

148 assert np.isclose(core_power.magnitude, total_power.magnitude, atol=1) 

149 

150 # Density source power 

151 density_source_power = _integrated_power_source(grid, norm, density_source, etemp_source, 

152 mean_density, mean_etemp, mean_itemp, kind='density_source') 

153 

154 edge_sink_power = _integrated_power_source(grid, norm, density_source, etemp_edge_sink, 

155 mean_density, mean_etemp, mean_itemp, kind='etemp_source') 

156 

157 assert np.isclose((density_source_power + edge_sink_power).magnitude, 0, atol=1) 

158 

159 print(f"The integrated power source is {total_power:4.3f}\n" 

160 f"The density source required {density_source_power:4.3f}, which was compensated by a temperature sink") 

161 

162 return {'density_source': integrated_density_source, 

163 'power_source': total_power}