Coverage for tcvx21/grillix_post/observables/velocity_m.py: 81%

1""" 

2Parallel velocity and ExB velocity, as vector arrays 

3""" 

4import xarray as xr 

5import numpy as np 

6from tcvx21.grillix_post.components.vector_m import toroidal_vector, vector_cross, vector_dot 

7from tcvx21.grillix_post.observables.perpendicular_gradient_m import perpendicular_gradient 

8 

9 

10def parallel_ion_velocity_vector(grid, equi, velocity: xr.DataArray): 

11 """ 

12 Converts the parallel ion velocity ion a vector array 

13 """ 

14 

15 return xr.DataArray(equi.parallel_unit_vector(grid.r_u, grid.z_u, grid=False) * velocity).assign_attrs( 

16 norm=velocity.norm, name="Parallel velocity vector") 

17 

18 

19def sound_speed(electron_temp: xr.DataArray, ion_temp: xr.DataArray, norm, adiabatic_index: float = 1.0): 

20 """ 

21 Ion sound speed for longitudinal waves. 

22 

23 Use adiabatic_index = 0.0 to exclude ion_temperature from the calculation 

24 Normally use adiabatic_index = 1.0, but it is possible that 3.0 could also be valid since some kinetic modelling 

25 suggests ions have an additional degree of freedom 

26 """ 

27 

28 return xr.DataArray(np.sqrt(electron_temp + adiabatic_index * norm.Z * ion_temp / norm.zeta.magnitude)).assign_attrs( 

29 norm=norm.c_s0, name="Ion sound speed") 

30 

31 

32def electric_field(grid, norm, potential: xr.DataArray): 

33 """ 

34 Returns the electric field calculated from the scalar potential, as a vector array 

35 

36 The potential gradient is normalised to the perpendicular length scale (i.e. units of phi0 / rho_s0) 

37 """ 

38 shaped_potential = grid.shape(potential) 

39 assert shaped_potential.norm 

40 

41 potential_gradient = perpendicular_gradient(grid, norm, shaped_potential) 

42 

43 return grid.flatten(xr.DataArray(potential_gradient * -1.0))\ 

44 .assign_attrs(norm=potential_gradient.norm, name="Electric field") 

45 

46 

47def exb_velocity(grid, equi, norm, potential: xr.DataArray): 

48 """ 

49 Returns the E cross B velocity as a vector array. 

50 Assumes that the magnetic field can be approximated as purely toroidal. 

51 

52 It is possible to calculate the exb_velocity with the full magnetic field, but the result is no longer within a 

53 poloidal plane. 

54 

55 Units 

56 [vE] = [B]/[B]**2 * [grad_perp] * [phi] 

57 = B0 / B0**2 * 1/rho_s0 * Te0/e 

58 = Te0/(e * rho_s0 * B0) 

59 We can use rho_s0 = sqrt(Te0 * Mi)/(e * B0), so 

60 [vE] = sqrt(Te0 / Mi) 

61 = c_s0 

62 

63 Need to multiply by a factor of delta to get to perpendicular velocity scale (which is what is used in GRILLIX). 

64 However, in order to easily compare to parallel velocities, we leave the drift velocity normalised to c_s0 

65 """ 

66 

67 E_field = electric_field(grid, norm, potential) 

68 

69 B_toroidal = norm.convert_norm(equi.magnetic_field_toroidal(grid.r_u, grid.z_u, grid=False)) 

70 

71 v_E = (vector_cross(E_field, toroidal_vector(B_toroidal)) / B_toroidal ** 2).assign_attrs( 

72 norm=E_field.norm / B_toroidal.norm, 

73 name='ExB velocity' 

74 ).persist() 

75 

76 return v_E 

77 

78 

79def effective_parallel_exb_velocity(grid, equi, norm, potential: xr.DataArray): 

80 """ 

81 The ExB velocity has components both across and along magnetic flux surfaces. 

82 

83 For the along-flux-surface component, we are interested in how this compares to the parallel transport. 

84 

85 We can consider this in one of two ways. We can compute the total_velocity, which is a 3D vector, and then 

86 use vector_dot to project this into some direction of interest. 

87 

88 Another approach is to consider what parallel velocity would lead to the same poloidal transport as the ExB 

89 velocity. This is the effective parallel ExB velocity (it is NOT the parallel component of the ExB, which is small) 

90 

91 We first compute the poloidal component of the ExB velocity, and then use the inverse magnetic field pitch to 

92 find the equivalent parallel vector (essentially, calculating the hypotenuse) 

93 

94 Because this term is primarily useful for comparing to diagnostics, we return as a flattened array 

95 """ 

96 

97 v_E = exb_velocity(grid, equi, norm, potential) 

98 

99 e_pol = equi.poloidal_unit_vector(grid.r_u, grid.z_u, grid=False) 

100 

101 b_x = equi.magnetic_field_r(grid.r_u, grid.z_u, grid=False) 

102 b_y = equi.magnetic_field_z(grid.r_u, grid.z_u, grid=False) 

103 b_t = equi.magnetic_field_toroidal(grid.r_u, grid.z_u, grid=False) 

104 

105 return vector_dot(v_E, e_pol) * np.sqrt((b_x**2+b_y**2+b_t**2)/(b_x**2+b_y**2)) 

106 

107 

108def total_velocity(grid, equi, norm, potential: xr.DataArray, velocity: xr.DataArray): 

109 """Returns the sum of the ExB and parallel velocity (it is assumed that the diamagnetic component is small""" 

110 

111 v_E = exb_velocity(grid, equi, norm, potential) 

112 u_par = parallel_ion_velocity_vector(grid, equi, velocity) 

113 

114 # The ExB velocity norm is equal to the sound speed. We can check this 

115 assert np.isclose(norm.c_s0, v_E.norm) 

116 assert np.isclose(norm.c_s0, u_par.norm) 

117 

118 return (v_E + u_par).assign_attrs(name='Total velocity', norm=norm.c_s0) 

119