Coverage for tcvx21/grillix_post/components/snaps_m.py: 88%

1""" 

2Interface to snapshot files 

3""" 

4from pathlib import Path 

5import numpy as np 

6import xarray as xr 

7import time 

8 

9variable_dict = { 

10 "logne": "density", 

11 "logte": "electron_temp", 

12 "logti": "ion_temp", 

13 "potxx": "potential", 

14 #"vortx": "vorticity", 

15 "jparx": "current", 

16 "uparx": "velocity", 

17 #"aparx": "apar", 

18 "tau": "tau" 

19} 

20 

21drop_vars = [ 

22 "logne_perpghost", "logte_perpghost", "logti_perpghost", "potxx_perpghost" 

23 "vortx", "vortx_perpghost", "jparx_perpghost", "uparx_perpghost" 

24 "aparx", "aparx_perpghost", "tau" 

25] 

26only_tau = [ 

27 "logne", "logne_perpghost", "logte", "logte_perpghost", "logti", "logti_perpghost", 

28 "potxx", "potxx_perpghost", "vortx", "vortx_perpghost", "jparx", "jparx_perpghost", 

29 "uparx", "uparx_perpghost", "aparx", "aparx_perpghost" 

30] 

31 

32 

33def get_snap_length_from_file(file_path: Path): 

34 return xr.open_dataset(file_path / 'snaps00000.nc').sizes['dim_tau'] 

35 

36 

37def read_snaps_from_file(file_path: Path, norm, time_slice: slice = slice(-1, None), 

38 all_planes: bool = False, verbose: bool = False) -> xr.Dataset: 

39 """ 

40 Reads a single snapshot file (corresponding to a single toroidal plane) 

41 

42 A subset of the time-interval can be loaded to reduce the memory requirements 

43 """ 

44 

45 if all_planes: 

46 if verbose: 

47 print("Reading all planes") 

48 

49 snaps = [] 

50 for plane in range(len([snap for snap in file_path.glob('snaps*.nc')])): 

51 

52 if verbose: 

53 print(f"Reading plane {plane}", end='. ') 

54 start = time.time() 

55 

56 snap = xr.open_dataset(file_path / f'snaps{plane:05d}.nc', chunks={'dim_tau': 50}, 

57 engine='netcdf4', drop_variables=drop_vars) 

58 snaps.append(snap.isel(dim_tau=time_slice).persist().expand_dims({'phi': [plane]})) 

59 snap.close() 

60 

61 if verbose: 

62 print(f"Finished in {time.time() - start:4.3f}s") 

63 

64 if verbose: 

65 print(f"Concatenating planes", end='. ') 

66 start = time.time() 

67 dataset = xr.concat(snaps, dim='phi') 

68 if verbose: 

69 print(f"Finished in {time.time() - start:4.3f}s") 

70 

71 else: 

72 if verbose: 

73 print("Reading one plane") 

74 dataset = xr.open_dataset(str(file_path / 'snaps00000.nc'), 

75 chunks={'dim_tau': 50}, 

76 drop_variables=drop_vars).expand_dims({'phi': [0]}) 

77 dataset = dataset.isel(dim_tau=time_slice) 

78 dataset.close() 

79 

80 tau = xr.open_dataset(file_path / 'snaps00000.nc', 

81 drop_variables=only_tau)['tau'].isel(dim_tau=time_slice) 

82 tau.close() 

83 

84 dataset['tau'] = tau 

85 dataset = dataset.rename({'dim_tau': 'tau'}) 

86 

87 normalisation = { 

88 'tau': norm.tau_0, 

89 'density': norm.n0, 

90 'electron_temp': norm.Te0, 

91 'ion_temp': norm.Ti0, 

92 'velocity': norm.c_s0, 

93 'current': (norm.c_s0 * norm.elementary_charge * norm.n0).to('kiloampere*meter**-2'), 

94 'potential': (norm.Te0 / norm.elementary_charge).to("kilovolt"), 

95 'vorticity': (norm.Mi * norm.n0 * norm.Te0 

96 / (norm.elementary_charge * norm.rho_s0 ** 2 * norm.B0 ** 2) 

97 ).to('coulomb/meter**3'), 

98 'apar': norm.beta_0 * norm.B0 * norm.rho_s0, 

99 } 

100 

101 snaps = xr.Dataset() 

102 

103 for var_access, variable in variable_dict.items(): 

104 attrs = {'name': variable, 'norm': normalisation[variable]} 

105 

106 if variable == "tau": 

107 snaps[variable] = dataset[variable].assign_attrs(**attrs) 

108 else: 

109 var = dataset[var_access].rename({'dim_vgrid': 'points'}) 

110 

111 if dataset[var_access].logquantity: 

112 snaps[variable] = np.exp(var).assign_attrs(**attrs) 

113 else: 

114 snaps[variable] = var.assign_attrs(**attrs) 

115 

116 return snaps