Coverage for src/hallmd/models/plume.py: 93%

1"""Module for Hall thruster plume models. 

2 

3Includes: 

4 

5- `current_density()` - Semi-empirical ion current density model with $1/r^2$ Gaussian beam. 

6""" 

7 

8import numpy as np 

9from amisc.typing import Dataset 

10from amisc.utils import get_logger 

11from scipy.integrate import simpson 

12from scipy.special import erfi 

13 

14from hallmd.utils import TORR_2_PA 

15 

16__all__ = ['current_density'] 

17 

18LOGGER = get_logger(__name__) 

19 

20 

21def current_density(inputs: Dataset, sweep_radius: float | list = 1.0): 

22 """Compute the semi-empirical ion current density ($j_{ion}$) plume model over a 90 deg sweep, with 0 deg at 

23 thruster centerline. Also compute the plume divergence angle. Will return the ion current density at 91 points, 

24 from 0 to 90 deg in 1 deg increments. The angular locations are returned as `j_ion_coords` in radians. 

25 

26 :param inputs: input arrays - `P_b`, `c0`, `c1`, `c2`, `c3`, `c4`, `c5`, `sigma_cex`, `I_B0` for background 

27 pressure (Torr), plume fit coefficients, charge-exchange cross-section ($m^2$), 

28 and total initial ion beam current (A). If `T` is provided, then 

29 also compute corrected thrust using the divergence angle. 

30 :param sweep_radius: the location(s) at which to compute the ion current density 90 deg sweep, in units of radial 

31 distance (m) from the thruster exit plane. If multiple locations are provided, then the 

32 returned $j_{ion}$ array's last dimension will match the length of `sweep_radius`. Defaults to 

33 1 meter. 

34 :returns outputs: output arrays - `j_ion` for ion current density ($A/m^2$) at the `j_ion_coords` locations, 

35 and `div_angle` in radians for the divergence angle of the plume. Optionally, 

36 `T_c` for corrected thrust (N) if `T` is provided in the inputs. 

37 """ 

38 # Load plume inputs 

39 P_B = inputs['P_b'] * TORR_2_PA # Background pressure (Torr) 

40 c0 = inputs['c0'] # Fit coefficients (-) 

41 c1 = inputs['c1'] # (-) 

42 c2 = inputs['c2'] # (rad/Pa) 

43 c3 = inputs['c3'] # (rad) 

44 c4 = inputs['c4'] # (m^-3/Pa) 

45 c5 = inputs['c5'] # (m^-3) 

46 sigma_cex = inputs['sigma_cex'] # Charge-exchange cross-section (m^2) 

47 I_B0 = inputs['I_B0'] # Total initial ion beam current (A) 

48 thrust = inputs.get('T', None) # Thrust (N) 

49 sweep_radius = np.atleast_1d(sweep_radius) 

50 

51 # 90 deg angle sweep for ion current density 

52 alpha_rad = np.linspace(0, np.pi / 2, 91) 

53 

54 # Neutral density 

55 n = c4 * P_B + c5 # m^-3 

56 

57 # Divergence angles 

58 alpha1 = np.atleast_1d(c2 * P_B + c3) # Main beam divergence (rad) 

59 alpha1[alpha1 > np.pi / 2] = np.pi / 2 

60 alpha2 = alpha1 / c1 # Scattered beam divergence (rad) 

61 

62 with np.errstate(invalid='ignore', divide='ignore'): 

63 A1 = (1 - c0) / ( 

64 (np.pi ** (3 / 2)) 

65 / 2 

66 * alpha1 

67 * np.exp(-((alpha1 / 2) ** 2)) 

68 * ( 

69 2 * erfi(alpha1 / 2) 

70 + erfi((np.pi * 1j - (alpha1**2)) / (2 * alpha1)) 

71 - erfi((np.pi * 1j + (alpha1**2)) / (2 * alpha1)) 

72 ) 

73 ) 

74 A2 = c0 / ( 

75 (np.pi ** (3 / 2)) 

76 / 2 

77 * alpha2 

78 * np.exp(-((alpha2 / 2) ** 2)) 

79 * ( 

80 2 * erfi(alpha2 / 2) 

81 + erfi((np.pi * 1j - (alpha2**2)) / (2 * alpha2)) 

82 - erfi((np.pi * 1j + (alpha2**2)) / (2 * alpha2)) 

83 ) 

84 ) 

85 # Broadcast over angles and radii (..., a, r) 

86 A1 = np.expand_dims(A1, axis=(-1, -2)) # (..., 1, 1) 

87 A2 = np.expand_dims(A2, axis=(-1, -2)) 

88 alpha1 = np.expand_dims(alpha1, axis=(-1, -2)) 

89 alpha2 = np.expand_dims(alpha2, axis=(-1, -2)) 

90 I_B0 = np.expand_dims(I_B0, axis=(-1, -2)) 

91 n = np.expand_dims(n , axis=(-1, -2)) 

92 sigma_cex = np.expand_dims(sigma_cex, axis=(-1, -2)) 

93 

94 decay = np.exp(-sweep_radius * n * sigma_cex) # (..., 1, r) 

95 j_cex = I_B0 * (1 - decay) / (2 * np.pi * sweep_radius ** 2) 

96 

97 base_density = I_B0 * decay / sweep_radius ** 2 

98 j_beam = base_density * A1 * np.exp(-((alpha_rad[..., np.newaxis] / alpha1) ** 2)) 

99 j_scat = base_density * A2 * np.exp(-((alpha_rad[..., np.newaxis] / alpha2) ** 2)) 

100 

101 j_ion = j_beam + j_scat + j_cex # (..., 91, r) the current density 1d profile at r radial locations 

102 

103 # Set j~0 where alpha1 < 0 (invalid cases) 

104 invalid_idx = np.logical_or(np.any(alpha1 <= 0, axis=(-1, -2)), np.any(j_ion <= 0, axis=(-1, -2))) 

105 j_ion[invalid_idx, ...] = 1e-20 

106 j_cex[invalid_idx, ...] = 1e-20 

107 

108 if np.any(abs(j_ion.imag) > 0): 

109 LOGGER.warning('Predicted beam current has non-zero imaginary component.') 

110 j_ion = j_ion.real 

111 

112 # Calculate divergence angle from https://aip.scitation.org/doi/10.1063/5.0066849 

113 # Requires alpha = [0, ..., 90] deg, from thruster exit-plane to thruster centerline (need to flip) 

114 # do j_beam + j_scat instead of j_ion - j_cex to avoid catastrophic loss of precision when 

115 # j_beam and j_scat << j_cex 

116 j_non_cex = np.flip((j_beam + j_scat).real, axis=-2) 

117 den_integrand = j_non_cex * np.cos(alpha_rad[..., np.newaxis]) 

118 num_integrand = den_integrand * np.sin(alpha_rad[..., np.newaxis]) 

119 

120 with np.errstate(divide='ignore', invalid='ignore'): 

121 num = simpson(num_integrand, x=alpha_rad, axis=-2) 

122 den = simpson(den_integrand, x=alpha_rad, axis=-2) 

123 cos_div = np.atleast_1d(num / den) 

124 cos_div[cos_div == np.inf] = np.nan 

125 

126 div_angle = np.arccos(cos_div) # Divergence angle (rad) - (..., r) 

127 

128 # Squeeze last dim if only a single radius was passed 

129 if sweep_radius.shape[0] == 1: 

130 j_ion = np.squeeze(j_ion, axis=-1) 

131 div_angle = np.squeeze(div_angle, axis=-1) 

132 

133 ret = {'j_ion': j_ion, 'div_angle': div_angle} 

134 

135 if thrust is not None: 

136 thrust_corrected = np.expand_dims(thrust, axis=-1) * cos_div 

137 if sweep_radius.shape[0] == 1: 

138 thrust_corrected = np.squeeze(thrust_corrected, axis=-1) 

139 ret['T_c'] = thrust_corrected 

140 

141 # Interpolate to requested angles 

142 # if j_ion_coords is not None: 

143 # # Extend to range (-90, 90) deg 

144 # alpha_grid = np.concatenate((-np.flip(alpha_rad)[:-1], alpha_rad)) # (2M-1,) 

145 # jion_grid = np.concatenate((np.flip(j_ion, axis=-1)[..., :-1], j_ion), axis=-1) # (..., 2M-1) 

146 # 

147 # f = interp1d(alpha_grid, jion_grid, axis=-1) 

148 # j_ion = f(j_ion_coords) # (..., num_pts) 

149 

150 # Broadcast coords to same loop shape as j_ion (all use the same coords -- store in object array) 

151 last_axis = -1 if sweep_radius.shape[0] == 1 else -2 

152 j_ion_coords = np.empty(j_ion.shape[:last_axis], dtype=object) 

153 for index in np.ndindex(j_ion.shape[:last_axis]): 

154 j_ion_coords[index] = alpha_rad 

155 

156 ret['j_ion_coords'] = j_ion_coords 

157 

158 return ret