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

1import numpy as np 

2import xarray as xr 

3from tcvx21.grillix_post.components import poloidal_vector 

4 

5 

6def perpendicular_gradient(grid, norm, shaped_array: xr.DataArray, normalised_to_Lperp: bool = True): 

7 """ 

8 Takes the perpendicular gradient of a shaped array (i.e. already has had vector_to_matrix applied) 

9 

10 Returns as a vector array. 

11 

12 If normalised_to_Lperp, returns an array which is the gradient in units of 1/rho_s0 (perpendicular length scale) 

13 Otherwise, returns an array which is the gradient in units of 1/R0 (parallel length scale) 

14 """ 

15 

16 dims = shaped_array.dims 

17 assert 'R' in dims and 'Z' in dims, f'Should pass shaped array (i.e. apply vector_to_matrix) to perpendicular_gradient' 

18 

19 delta = (norm.R0 / norm.rho_s0).to('').magnitude 

20 if normalised_to_Lperp: 

21 # Take the gradient w.r.t. perpendicular length scale 

22 gradient = np.gradient(shaped_array, grid.spacing * delta, axis=(dims.index('R'), dims.index('Z'))) 

23 gradient_scale_length = norm.rho_s0 

24 else: 

25 # Take the gradient w.r.t. parallel length scale 

26 gradient = np.gradient(shaped_array, grid.spacing, axis=(dims.index('R'), dims.index('Z'))) 

27 gradient_scale_length = norm.R0 

28 

29 xr_data = {'dims': shaped_array.dims, 'coords': shaped_array.coords, 'attrs': shaped_array.attrs} 

30 vector_grad = poloidal_vector(input_r=gradient[0], input_z=gradient[1], **xr_data) 

31 

32 if 'norm' in vector_grad.attrs: 

33 assert not isinstance( 

34 vector_grad.norm, str), f"Cannot operate on units in string format (in perpendicular_gradient). Norm was\ 

35 {vector_grad.norm}. Convert to a Quantity via the Normalisation before calling this function." 

36 vector_grad.attrs['norm'] *= 1.0 / gradient_scale_length 

37 else: 

38 vector_grad.attrs['norm'] = 1.0 / gradient_scale_length 

39 

40 return vector_grad