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 

8from typing import cast 

9 

10import numpy as np 

11from pem_core import get_logger 

12from pem_core.constants import TORR_2_PA 

13from pem_core.types import ArrayLike, Dataset 

14from scipy.integrate import simpson 

15from scipy.special import erfi 

16 

17__all__ = ['current_density'] 

18 

19LOGGER = get_logger(__name__) 

20 

21def current_density(inputs: Dataset | dict, sweep_radius: float | ArrayLike = 1.0) -> Dataset: 

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 input_dict = cast(dict, inputs) 

40 P_B = input_dict['P_b'] * TORR_2_PA # Background pressure (Torr) 

41 c0 = input_dict['c0'] # Fit coefficients (-) 

42 c1 = input_dict['c1'] # (-) 

43 c2 = input_dict['c2'] # (rad/Pa) 

44 c3 = input_dict['c3'] # (rad) 

45 c4 = input_dict['c4'] # (m^-3/Pa) 

46 c5 = input_dict['c5'] # (m^-3) 

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

48 I_B0 = input_dict['I_B0'] # Total initial ion beam current (A) 

49 thrust = input_dict.get('T', None) # Thrust (N) 

50 radii = np.atleast_1d(sweep_radius) 

51 

52 # 90 deg angle sweep for ion current density 

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

54 

55 # Neutral density 

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

57 

58 # Divergence angles 

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

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

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

62 

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

64 A1 = (1 - c0) / ( 

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

66 / 2 

67 * alpha1 

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

69 * ( 

70 2 * erfi(alpha1 / 2) 

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

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

73 ) 

74 ) 

75 A2 = c0 / ( 

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

77 / 2 

78 * alpha2 

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

80 * ( 

81 2 * erfi(alpha2 / 2) 

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

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

84 ) 

85 ) 

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

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

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

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

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

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

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

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

94 

95 decay = np.exp(-radii * n * sigma_cex) # (..., 1, r) 

96 j_cex = I_B0 * (1 - decay) / (2 * np.pi * radii**2) 

97 

98 base_density = I_B0 * decay / radii**2 

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

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

101 

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

103 

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

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

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

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

108 

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

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

111 j_ion = j_ion.real 

112 

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

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

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

116 # j_beam and j_scat << j_cex 

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

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

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

120 

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

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

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

124 cos_div = np.atleast_1d(num / den) 

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

126 

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

128 

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

130 if radii.shape[0] == 1: 

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

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

133 

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

135 

136 if thrust is not None: 

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

138 if radii.shape[0] == 1: 

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

140 ret['T_c'] = thrust_corrected 

141 

142 # Interpolate to requested angles 

143 # if j_ion_coords is not None: 

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

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

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

147 # 

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

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

150 

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

152 last_axis = -1 if radii.shape[0] == 1 else -2 

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

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

155 j_ion_coords[index] = alpha_rad 

156 

157 ret['j_ion_coords'] = j_ion_coords 

158 

159 return cast(Dataset, ret)