Coverage for tcvx21/grillix_post/observables/parallel_gradient_m.py: 100%

1import xarray as xr 

2import numpy as np 

3from pathlib import Path 

4from tcvx21.grillix_post.components import FieldlineTracer 

5from tcvx21.grillix_post.lineouts import Lineout 

6xr.set_options(keep_attrs=True) 

7 

8 

9def initialise_lineout_for_parallel_gradient(lineout, grid, equi, norm, npol, stored_trace: Path=None): 

10 """ 

11 Traces to find the forward and reverse lineouts for a given lineout 

12 

13 Expensive! Needs to be done once per lineout that you want to take gradients with 

14 """ 

15 fieldline_tracer = FieldlineTracer(equi) 

16 

17 try: 

18 print(f"Attempting to read stored trace from {stored_trace}") 

19 ds = xr.open_dataset(stored_trace) 

20 assert np.allclose(ds['lineout_x'], lineout.r_points) 

21 assert np.allclose(ds['lineout_y'], lineout.z_points) 

22 except (FileNotFoundError, ValueError): 

23 forward_trace, reverse_trace = \ 

24 fieldline_tracer.find_neighbouring_points(lineout.r_points, lineout.z_points, 

25 n_toroidal_planes=int(npol)) 

26 

27 ds = xr.Dataset(data_vars=dict( 

28 forward_x=('points', forward_trace[:, 0]), 

29 forward_y=('points', forward_trace[:, 1]), 

30 forward_l=('points', forward_trace[:, 2]), 

31 reverse_x=('points', reverse_trace[:, 0]), 

32 reverse_y=('points', reverse_trace[:, 1]), 

33 reverse_l=('points', reverse_trace[:, 2]), 

34 lineout_x=('points', lineout.r_points), 

35 lineout_y=('points', lineout.z_points), 

36 )) 

37 

38 if stored_trace is not None: 

39 if stored_trace.exists(): stored_trace.unlink() 

40 ds.to_netcdf(stored_trace) 

41 

42 lineout.forward_lineout = Lineout(ds['forward_x'], ds['forward_y']) 

43 lineout.forward_lineout.setup_interpolation_matrix(grid, use_source_points=True) 

44 lineout.reverse_lineout = Lineout(ds['reverse_x'], ds['reverse_y']) 

45 lineout.reverse_lineout.setup_interpolation_matrix(grid, use_source_points=True) 

46 

47 lineout.forward_distance = xr.DataArray(ds['forward_l'], 

48 dims='interp_points').assign_attrs(norm=norm.R0) 

49 lineout.reverse_distance = xr.DataArray(ds['reverse_l'], 

50 dims='interp_points').assign_attrs(norm=norm.R0) 

51 

52 

53def compute_parallel_gradient(lineout, field): 

54 """ 

55 Computes the parallel gradient via centred differences 

56 

57 Note that you should multiply this by the penalisation direction function to get the direction 'towards the 

58 wall'. This isn't quite the same as projecting onto the wall normal, but for computing the parallel 

59 heat flux this is actually more helpful 

60 """ 

61 

62 assert hasattr(lineout, 'forward_lineout') and hasattr(lineout, 'reverse_lineout'), \ 

63 f"Have to call initialise_lineout_for_parallel_gradient on lineout before trying to compute_parallel_gradient" 

64 

65 parallel_gradients = [compute_gradient_on_plane(lineout, field, plane) 

66 for plane in range(field.sizes['phi'])] 

67 

68 return xr.concat(parallel_gradients, dim='phi') 

69 

70 

71def compute_gradient_on_plane(lineout, field, plane): 

72 """Computes the parallel gradient on a single plane""" 

73 

74 forward_value = lineout.forward_lineout.interpolate(field.isel(phi=np.mod(plane+1, field.sizes['phi']))) 

75 reverse_value = lineout.forward_lineout.interpolate(field.isel(phi=np.mod(plane-1, field.sizes['phi']))) 

76 

77 two_plane_distance = (lineout.forward_distance - lineout.reverse_distance) 

78 

79 centred_difference = forward_value - reverse_value 

80 return (centred_difference/two_plane_distance).\ 

81 assign_coords(phi=plane).\ 

82 assign_attrs(norm=field.norm/two_plane_distance.norm)