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

1import os 

2import re 

3import numpy as np 

4 

5 

6def giveMeasureTypeUnits(measureName, measureType=None): 

7 """ Given a measure name, returns its units and its type. 

8 

9 :param measureName: Name of the measure 

10 :type measureName: str 

11 :param measureType: Type of the measure (optional) 

12 :type measureType: str 

13 :return: A 2-tuple containing the type of the measure and the units 

14 :rtype: str,str 

15 """ 

16 if measureType is None: 

17 if measureName == "ER" or measureName == "DIFF": 

18 measureType = "SV" 

19 else: 

20 measureType = measureName 

21 

22 if measureType == "MF": 

23 units = "nT" 

24 elif measureType == "SV": 

25 units = "nT/yr" 

26 elif measureType == "U" or measureType == "S": 

27 units = "km/yr" 

28 elif measureType == "DIVHU": 

29 units = "/yr" 

30 else: 

31 print("Measure name not recognized for %s in giveMeasureTypeUnits. Setting units to nT.", measureName) 

32 units = "nT" 

33 

34 return measureType, units 

35 

36 

37def load(dataDirectory, dataModel, keepRealisations=True): 

38 """ Loading function for default type. Also adds the data to the dataModel. 

39 

40 The files to load should have one of the following namepatterns : 

41 - 'MeasureType_iRealisation.dat' (e.g. MF_0001.dat [type should be 'MF' or 'SV' or 'U']) 

42 - 'MeasureName-MeasureType_iRealisation.dat' (e.g. ER-SV_0001.dat ; MagField-MF_0001.dat ; ... [type should be 'MF' or 'SV' or 'U']) 

43 

44 In addition : 

45 - Time should be in "times.dat" 

46 - Units must be : nT, km, yr. 

47 

48 :param dataDirectory: location of the files to load 

49 :type dataDirectory: str 

50 :param dataModel: data model to which the loaded data are added 

51 :type dataModel: Model 

52 :param keepRealisations: indicates if realisations must be kept or not (not used) 

53 :type keepRealisations: bool 

54 :return: 0 if no errors are raised 

55 :rtype: int 

56 """ 

57 

58 if os.path.isfile(os.path.join(dataDirectory, "times.dat")): 

59 times = np.loadtxt(os.path.join(dataDirectory, "times.dat")) 

60 else: 

61 raise ValueError("%s does not contain times.dat file, did you forget to specify dataFormat?" % dataDirectory) 

62 

63 measures = {} 

64 for file in os.listdir(dataDirectory): 

65 print("[%s]" % file) 

66 matchfile = re.match('^(.*)_(.*).dat$', file) 

67 

68 if matchfile is not None: 

69 if len(matchfile.group(1).split('-')) == 2: 

70 measureName, measureType = matchfile.group(1).split('-') 

71 else: 

72 measureName = matchfile.group(1).split('-') 

73 measureType = measureName 

74 irealisation = int(matchfile.group(2))-1 

75 print(" ", measureName, measureType, irealisation) 

76 

77 if measureName not in measures: 

78 measures[measureName] = {"type": measureType, "filelist": []} 

79 measures[measureName]["filelist"].append([irealisation, file]) 

80 

81 if len(measures) == 0: 

82 raise ValueError("Unable to read %s as default format: no file matching *_*.dat" % dataDirectory) 

83 

84 for measureName in measures: 

85 measureData = None 

86 measureType = measures[measureName]["type"] 

87 nrealisations = len(measures[measureName]["filelist"]) 

88 measureType, units = giveMeasureTypeUnits(measureName, measureType) 

89 for ireal, fileName in measures[measureName]["filelist"]: 

90 # Reading realisation file 

91 sourceData = np.loadtxt(os.path.join(dataDirectory, fileName)) 

92 

93 # Compute lmax 

94 if measureType == 'U' or measureType == 'S': 

95 lmax = (-2+np.sqrt(4+4*(sourceData.shape[1])/2))/2 

96 else: 

97 lmax = (-2+np.sqrt(4+4*sourceData.shape[1]))/2 

98 

99 if int(lmax) != lmax: 

100 if measureType == 'U' or measureType == 'S': 

101 raise ValueError("Data length %i in file %s does not correspond to any 2*lmax*(lmax+2)" % (sourceData.shape[1], fileName)) 

102 else: 

103 raise ValueError("Data length %i in file %s does not correspond to any lmax*(lmax+2)" % (sourceData.shape[1], fileName)) 

104 else: 

105 lmax = int(lmax) 

106 

107 if times.shape[0] != sourceData.shape[0]: 

108 raise ValueError('times.dat file give %i times, %s file contains %i times' % (times.shape[0], fileName, sourceData.shape[0])) 

109 

110 if measureData is not None: 

111 measureData[:, :, ireal] = sourceData 

112 else: 

113 if nrealisations > 1: 

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

115 measureData = np.zeros((sourceData.shape[0], sourceData.shape[1], nrealisations)) 

116 measureData[:, :, ireal] = sourceData 

117 else: 

118 measureData = sourceData 

119 

120 dataModel.addMeasure(measureName, measureType, lmax, units, measureData, times=times) 

121 

122 # If everything went fine 

123 return 0 

124 

125 

126def save(dataDirectory, dataModel, forceOverwrite=False): 

127 """ Saves the given model in .dat files. 

128 

129 :param dataDirectory: location of the saving directory 

130 :type dataDirectory: str 

131 :param dataModel: data model to save 

132 :type dataModel: Model 

133 :param forceOverwrite: force the overwrite of existing files that bear the same name 

134 :type forceOverwrite: bool 

135 :return: 0 if no errors are raised 

136 :rtype: int 

137 

138 """ 

139 dataDirectory = os.path.realpath(dataDirectory) 

140 

141 if not forceOverwrite: 

142 # Test the existence of all file names before starting to write 

143 if os.path.isfile(os.path.join(dataDirectory, "times.dat")): 

144 if not forceOverwrite: 

145 raise ValueError("%s already contains a times.dat file, please set forceOverwrite to True if you wish to overwrite it." % dataDirectory) 

146 for measureName in dataModel.measures: 

147 measure = dataModel.measures[measureName] 

148 measureType = measure.measureType 

149 if measure.has_realisations: 

150 for ireal in range(measure.data.shape[2]): 

151 fileName = measureName + "-" + measureType + "_" + str(ireal+1).zfill(6) + ".dat" 

152 if os.path.isfile(os.path.join(dataDirectory, fileName)): 

153 raise ValueError("%s already contains a file named %s, please set forceOverwrite to True if you wish to overwrite it." 

154 % (dataDirectory, fileName)) 

155 else: 

156 fileName = measureName + "-" + measureType + "_" + str(1).zfill(6) + ".dat" 

157 if os.path.isfile(os.path.join(dataDirectory, fileName)): 

158 raise ValueError("%s already contains a file named %s, please set forceOverwrite to True if you wish to overwrite it." 

159 % (dataDirectory, fileName)) 

160 

161 # Write files 

162 os.makedirs(dataDirectory, exist_ok=True) 

163 print("Saving times to", os.path.join(dataDirectory, "times.dat")) 

164 np.savetxt(os.path.join(dataDirectory, "times.dat"), dataModel.times) 

165 for measureName in dataModel.measures: 

166 measure = dataModel.measures[measureName] 

167 measureType = measure.measureType 

168 print("Saving", measureName, "to", 

169 os.path.join(dataDirectory, measureName + "-" + measureType + "_######.dat")) 

170 if measure.has_realisations: 

171 for ireal in range(measure.data.shape[2]): 

172 fileName = os.path.join(dataDirectory, measureName + "-" + measureType + "_" + str(ireal+1).zfill(6) + ".dat") 

173 np.savetxt(fileName, measure.data[:, :, ireal]) 

174 else: 

175 fileName = os.path.join(dataDirectory,measureName + "-" + measureType + "_" + str(1).zfill(6) + ".dat") 

176 np.savetxt(fileName, measure.data) 

177 

178 # If everything went fine 

179 return 0