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

1""" 

2Defines routines for calculating integral quantities 

3 

4Automatically handles units 

5""" 

6import numpy as np 

7import xarray as xr 

8from scipy.interpolate import RectBivariateSpline 

9from tcvx21.units_m import Quantity 

10 

11 

12def poloidal_integration(grid, norm, poloidal_function: xr.DataArray): 

13 """ 

14 Returns the 2D poloidal integral of a function defined 

15 

16 Off-grid values are assigned a value of 0.0, and a filled grid spline interpolator is used to evaluate the integral 

17 """ 

18 assert hasattr(poloidal_function, 'norm') 

19 assert isinstance(poloidal_function.norm, Quantity) 

20 

21 grid_limits = { 

22 'xa': grid.r_s.min(), 

23 'xb': grid.r_s.max(), 

24 'ya': grid.z_s.min(), 

25 'yb': grid.z_s.max() 

26 } 

27 

28 poloidal_function_values = np.nan_to_num(poloidal_function.values, nan=0.0) 

29 

30 function_integral = RectBivariateSpline(grid.r_s, grid.z_s, 

31 poloidal_function_values.T 

32 ).integral(**grid_limits) 

33 

34 norm = poloidal_function.norm * norm.R0 **2 

35 

36 return function_integral * norm 

37 

38 

39def axisymmetric_cylindrical_integration(grid, norm, poloidal_function: xr.DataArray): 

40 """ 

41 Returns the 3D cylindrically-integrated value of a function defined over the poloidal grid 

42 

43 N.b. automatically takes care of the "R" factor introduced by the \theta integral 

44 """ 

45 

46 r_s_grid, _ = np.meshgrid(grid.r_s, grid.z_s) 

47 poloidal_integration_function = xr.DataArray(poloidal_function * r_s_grid).assign_attrs( 

48 norm = poloidal_function.norm) 

49 function_integral = 2.0 * np.pi * norm.R0 * poloidal_integration(grid, norm, poloidal_integration_function) 

50 

51 return function_integral