Coverage for tcvx21/grillix_post/components/fieldline_tracer_m.py: 100%

1""" 

2A class to handle tracing along field-lines 

3""" 

4from scipy.integrate import solve_ivp 

5import numpy as np 

6 

7 

8class FieldlineTracer: 

9 """ 

10 Routines to parallel-trace along a fieldline, based on the DOP853 integrator 

11 """ 

12 

13 def __init__(self, equi, rtol: float = 1E-3, atol: float = 1E-6): 

14 """ 

15 Simple initialiser 

16 equi should be an object which defines the magnetic field 

17 rtol and atol are used to set integration tolerances, for the relative and absolute tolerance 

18 """ 

19 self.equi = equi 

20 self.rtol, self.atol = rtol, atol 

21 

22 def integrate_ivp(self, func, initial, t_max, events=None): 

23 """Interface to DOP853 integrator""" 

24 initial = np.array(initial) # Throw a warning about ragged arrays early 

25 

26 solution = solve_ivp( 

27 fun=func, 

28 t_span=(0.0, t_max), 

29 y0=initial, 

30 rtol=self.rtol, 

31 atol=self.atol, 

32 method="RK45", 

33 events=events 

34 ) 

35 

36 return solution 

37 

38 def toroidal_integration(self, r_initial, z_initial, max_toroidal_trace, events=None): 

39 """ 

40 Performs an integration along a magnetic field line, with phi (toroidal angle) as the independent variable 

41 

42 Returns the (x,y) positions of the trace, as well as the fieldline length (in radians) 

43 

44 To get the x, y, L_parallel values in physical units, multiply by R0 

45 """ 

46 r_initial, z_initial = np.squeeze(r_initial), np.squeeze(z_initial) 

47 

48 solution = self.integrate_ivp( 

49 func=self.toroidal_integration_equation, 

50 initial=[r_initial, z_initial, 0.0], 

51 t_max=max_toroidal_trace, 

52 events=events 

53 ) 

54 

55 return solution 

56 

57 def toroidal_integration_equation(self, _, state): 

58 """ 

59 Integrate around the torus (independent variable = toroidal angle) 

60 

61 The fieldline length is returned as the third element of the state vector 

62 """ 

63 

64 r_norm = state[0] 

65 z_norm = state[1] 

66 

67 b_r = self.equi.magnetic_field_r(r_norm, z_norm, grid=False) 

68 b_z = self.equi.magnetic_field_z(r_norm, z_norm, grid=False) 

69 b_t = self.equi.magnetic_field_toroidal(r_norm, z_norm, grid=False) 

70 jacobian = r_norm 

71 

72 d_state = np.zeros_like(state) 

73 d_state[0] = b_r / b_t * jacobian 

74 d_state[1] = b_z / b_t * jacobian 

75 d_state[2] = np.sqrt(b_r * b_r / (b_t * b_t) + 

76 b_z * b_z / (b_t * b_t) + 1.0) * jacobian 

77 

78 return d_state 

79 

80 def find_neighbouring_points(self, r_initial, z_initial, n_toroidal_planes: int = 16): 

81 """ 

82 Finds the neighbours in both directions of an array of sample points 

83 Although this loop should be trivially parallel, dask parallelism doesn't seem to work here 

84 """ 

85 r_initial = np.atleast_1d(r_initial) 

86 z_initial = np.atleast_1d(z_initial) 

87 

88 assert r_initial.shape == z_initial.shape 

89 assert r_initial.ndim == 1 and z_initial.ndim == 1, f"Should provide r and z as 1D arrays" 

90 

91 forward_trace = np.zeros((r_initial.size, 3)) 

92 reverse_trace = np.zeros((r_initial.size, 3)) 

93 

94 print("Tracing", end=' ') 

95 for i in range(r_initial.size): 

96 print(f"{i}/{r_initial.size}", end=', ') 

97 

98 r_in, z_in = r_initial[i], z_initial[i] 

99 forward_trace[i, :] = self.toroidal_integration(r_in, z_in, +2.0 * np.pi / n_toroidal_planes).y[:, -1] 

100 reverse_trace[i, :] = self.toroidal_integration(r_in, z_in, -2.0 * np.pi / n_toroidal_planes).y[:, -1] 

101 print("Done") 

102 

103 return forward_trace, reverse_trace