Coverage for webgeodyn/inout/pygeodyn_asc.py: 0%

1import os 

2import re 

3import numpy as np 

4from .default import giveMeasureTypeUnits 

5 

6 

7def getDataAndLmax(data, measureType): 

8 """ Returns the data and the lmax for the data and type of measure given. 

9 

10 :param data: Array containing the measure data 

11 :type data: ndarray 

12 :param measureType: Type of the measure 

13 :type measureType: str 

14 :return: Data and the lmax computed from the data shape 

15 :rtype: ndarray, int 

16 """ 

17 sizeData = np.where(data[0] != 0)[0].shape[0] - 1 

18 if measureType == "U": 

19 lmax = (-2+np.sqrt(4+4*sizeData/2))/2 

20 nk = int(lmax*(lmax+2)) 

21 middle = int((data.shape[1] - 1)/2) 

22 data = data[:, np.r_[1:nk+1, middle+1:middle+nk+1]] 

23 else: 

24 lmax = (-2+np.sqrt(4+4*sizeData))/2 

25 nk = int(lmax*(lmax+2)) 

26 data = data[:, 1:nk+1] 

27 # Check if computed lmax is int 

28 if not (int(lmax) == lmax): 

29 raise ValueError('[{}] File size ({}) is not multiple of lmax*(lmax + 2), lmax={}'.format(measureType, sizeData, lmax)) 

30 

31 return data, lmax 

32 

33 

34def load(dataDirectory, dataModel, keepRealisations): 

35 """ Loading function for pygeodyn files of ASCII format. Also adds the data to the dataModel. 

36 

37 :param dataDirectory: Location of the pygeodyn files 

38 :type dataDirectory: os.path 

39 :param dataModel: Model in which to add the loaded measures 

40 :type dataModel: Model 

41 :param keepRealisations: If True, all realisations are kept in the data. Else, the data is averaged over the realisations 

42 :type keepRealisations: bool 

43 :return: 0 if everything went well, -1 otherwise 

44 :rtype: int 

45 """ 

46 print('Inside load function pygeodyn_asc') 

47 points_to_cut = 1 

48 times = None 

49 

50 valid_measures_names = ['DIFF', 'U', 'MF', 'SV', 'ER'] 

51 state_type = 'Analysed' #'Computed' 

52 

53 filelist = {} 

54 if keepRealisations: 

55 # Find all valid measures directories 

56 dirs_to_explore = [os.path.join(dataDirectory, x) for x in os.listdir(dataDirectory) if x in valid_measures_names] 

57 for realisation_dir in dirs_to_explore: 

58 for file in os.listdir(realisation_dir): 

59 matchfile = re.match('^'+state_type+'_Realisation(\d+)_(\w+)', file) 

60 if matchfile is not None: 

61 irealisation = int(matchfile.group(1)) 

62 measureName = matchfile.group(2) 

63 if measureName not in filelist: 

64 filelist[measureName] = [] 

65 filelist[measureName].append([irealisation, os.path.join(realisation_dir, file)]) 

66 

67 # Check if realisation data was found 

68 if len(filelist) == 0: 

69 raise IOError('No realisation data was found in {}. Check the path or try to load the model with keepRealisation to False.'.format(dataDirectory)) 

70 

71 # Treat the filelist previously build by extracting the data and adding it to the model 

72 for measureName in filelist: 

73 measureData = None 

74 measureType, units = giveMeasureTypeUnits(measureName) 

75 

76 i = 0 

77 for ireal, fileName in filelist[measureName]: 

78 print(i, f" / {len(filelist[measureName])} Reading real {ireal}", measureName, os.path.basename(fileName)) 

79 

80 # Reading realisation file 

81 sourceData = np.loadtxt(fileName)[points_to_cut:] 

82 # Extract times (once for all files) 

83 if times is None: 

84 times = sourceData[:, 0] 

85 # Prevent from loading times that are not the same for all files 

86 if not np.allclose(times, sourceData[:, 0]): 

87 print('[%s] Times are not the same in all files !', measureName) 

88 return -1 

89 

90 # Extract sourceData and the maximum sph. harm. coef lmax 

91 sourceData, lmax = getDataAndLmax(sourceData, measureType) 

92 if measureData is None: 

93 # Initialise data to be of size [ntimes, ncoef, nrealisations] 

94 measureData = np.zeros((sourceData.shape[0], sourceData.shape[1], len(filelist[measureName]))) 

95 measureData[:, :, ireal] = sourceData 

96 i +=1 

97 

98 # Add the data loaded to the model as measures 

99 dataModel.addMeasure(measureName, measureType, lmax, units, 

100 measureData, times=times) 

101 

102 # Else read Mean+RMS for all values and add them to the model 

103 else: 

104 mean_dir = os.path.join(dataDirectory, 'Mean') 

105 rms_dir = os.path.join(dataDirectory, 'RMS') 

106 for file in os.listdir(mean_dir): 

107 matchfile = re.match('^'+state_type+'_mean_(\w+)', file) 

108 if matchfile is not None: 

109 measureName = matchfile.group(1) 

110 rms_file = os.path.join(rms_dir, state_type+'_rms_'+measureName+'.dat') 

111 mean_file = os.path.join(mean_dir, file) 

112 measureType, units = giveMeasureTypeUnits(measureName) 

113 

114 # Reading measure file 

115 mean_data = np.loadtxt(mean_file)[points_to_cut:] 

116 rms_data = np.loadtxt(rms_file)[points_to_cut:] 

117 

118 # Extract times 

119 times = mean_data[:, 0] 

120 

121 mean_data, lmax = getDataAndLmax(mean_data, measureType) 

122 rms_data, lmax = getDataAndLmax(rms_data, measureType) 

123 

124 dataModel.addMeasure(measureName, measureType, lmax, units, 

125 mean_data, rms_data, times) 

126 

127 # Returns 0 if everything went well 

128 return 0