Source code for htrdrPy.postprocess

import numpy as np

from htrdrPy.helperFunctions import *
from htrdrPy.data import *
from htrdrPy.geometry import *
from htrdrPy.script import *


[docs] class Postprocess: def __init__(self, script: Script=None, exposure=1, cmap="inferno", heatCapacity=None, rho=None, threshold=0, skip=False, localPath=False): ''' The Postprocess module aims at processing the outputs of htrdr (based on the script launched or not). # Inputs: - script (optional, Script): Script object which as been called on a Data object - exposure (optional, float [s]): exposure time. Needed if the post-process operation consists in generating an image. - heatCapacity (optional, float [J/kg/K]): heat capacity. Neede if the post-process operation consists in temperature calculations from the flux divergence. - mmw (optional, float [kg/mol]): mean molecular weight. Needed if the post-process operation consists in temperature calculations from the flux divergence. - threshold (optional, float []): threshold limit below which the data are consider zero (expressed relatively to the data maximale value: threshold * max(data)). Default 0 - skip (optional, bool): whether or not to skip the default post-processing operation. Default is False - GCMmesh (optional, bool): whether or not the flux diveregnce calculation was realized on the original GCM mesh. Default is True (the other behavior is not implemented yet) # Note: Im most cases, the module recognize the kind of script and automatically apply the required post-process function. This is not the case for script based on "startMultipleObsGeometry", as the post-processing rootine to be used is not trivial. A bunch of additional methods are therefore provided. ''' if script: self.script = script self.data = script.data self.resultsPath = f'{self.data.workingDirectory}/results_{self.data.name}/' else: self.script = None self.data = None self.resultsPath = f'{pathlib.Path().resolve()}/results_/' if localPath: self.resultsPath = f'{pathlib.Path().resolve()}/results_{self.data.name}/' try: os.mkdir(self.resultsPath) except FileExistsError: pass if self.script==None: return if self.script.predefinedScript == "image": self.processImages(exposure=exposure, cmap=cmap) elif self.script.predefinedScript == "imageRatio": self.__imageRatio(threshold=threshold) elif self.script.predefinedScript == "compositeRGB": self.__compositeRGB() elif self.script.predefinedScript == "reflectanceSpectrum": self.__spectrum() elif self.script.predefinedScript == "spectrum": self.__spectrum() elif self.script.predefinedScript == "twoStream" and self.script.method==1: self.__twoStreamV1(heatCapacity=heatCapacity) elif self.script.predefinedScript == "twoStream" and self.script.method==2: self.__twoStreamV2(heatCapacity=heatCapacity) elif self.script.predefinedScript == "twoStream" and self.script.method==3: self.__twoStreamV3(heatCapacity=heatCapacity) elif self.script.predefinedScript == "radiativeBudget": self.__radiativeBudget(heatCapacity=heatCapacity, rho=rho) elif self.script.predefinedScript == "startMultipleObsGeometry": print("No predefined routine for such script") return
[docs] def getImage(self, file): ''' Recover the image data from a htrdr output file. # Input: - file (str): path to the htrdr output file # Ouputs: - radiance (2-D array): radiances associated to each pixel - deviation (2-D array): corresponding standard deviation ''' with open(file, 'r') as f: definition = np.array(f.readline().split(), dtype=int) npx, npy = int(definition[0]), int(definition[1]) radiance = np.zeros((npy,npx, 3)) deviation = np.zeros_like(radiance) for (i,j) in np.ndindex(npy,npx): line = np.array(f.readline().split(), dtype=float) radiance[i,j,:] = line[[0,2,4]] deviation[i,j,:] = line[[1, 3, 5]] return radiance, deviation
[docs] def getImages(self): ''' Recover the images data from all htrdr output files generated by the Script. # Outputs: - dictionnary with radiances and standard deviations map for all output files. ''' results = {} for file in self.script.outputFiles: name = file.split('/')[-1].replace('output_','').replace('.txt','') rad, dev = self.getImage(file) results[name] = { "radiance": rad, "deviation": dev } return results
[docs] def extractMeanRadiances(self, indices=False): ''' Extract the mean radiances of all (or a subset) images generated by the Script. # Outputs: - dictionnary with mean radiances and standard deviations (c.f. extractMeanRadianceFromOutput()) for all required output files. ''' results = {} if indices: for index in indices: file = f"{self.outputPath}output_{index}.txt" radiance, sigma = self.extractMeanRadianceFromOutput(file) results[file.split('/')[-1].replace("output_","").replace(".txt","")] = { "radiance": radiance, "std dev": sigma } else: for file in self.script.outputFiles: print(file) radiance, sigma = self.extractMeanRadianceFromOutput(file) results[file.split('/')[-1].replace("output_","").replace(".txt","")] = { "radiance": radiance, "std dev": sigma } return results
[docs] def processSingleArrayObsSW(self): ''' Processes the previously generated files in the case of single array (one of the image definition is 1) # Output: Return a 2D array containing the list of radiances (in W/m2/sr) and the associated standard deviation ''' res = {} for obs,file in zip(self.script.obsList, self.script.outputFiles): name = obs.case with open(file, 'r') as f: definition = np.array(f.readline().split(), dtype=int) if definition[0] == 1 or definition[1] == 1: nval = definition[0] * definition[1] results = np.zeros((nval,2)) ind = 0 for line in f: if line[0] == "#": continue if len(line.split()) == 0: continue data = np.array(line.split(), dtype=float) results[ind,:] = data[[2,3]] ind +=1 else: raise IndexError("Error processing output files: wrong image definition") res[name] = results return res
[docs] def processImages(self, exposure=1, cmap="inferno"): self.imageNames = [] for file in self.script.outputFiles: imageName = f"{self.resultsPath}/{file.split('/')[-1].replace('.txt', '.ppm').replace('output', f'image_ex{exposure}')}" if self.script.wavelengths['type'] == "cie_xyz": command = f"htpp -f -i exposure={exposure} -o {imageName} {file}" else: command = f"htpp -f -i exposure={exposure} -m palette={cmap}:pixcpnt=2 -o {imageName} {file}" subprocess.run(command, shell=True, check=True) self.imageNames.append(imageName)
[docs] def extractMeanRadianceFromOutput(self, file, time=False): ''' Calculate the average radiance and standard deviation over all pixels # Input: - file (str): path to htrdr output file containing the pixels information # Output: Average radiance and standard deviation over the pixel grid (units are similar to htrdr outputs) ''' with open(file, 'r') as f: definition = np.array(f.readline().split(), dtype=int) nPixel = definition[0] * definition[1] radiance = 0 variance = 0 radiances = np.zeros(nPixel) tpp = 0 varTimes = 0 times = np.zeros(nPixel) index = 0 for line in f: if line[0] == "#": continue if len(line.split()) == 0: continue data = np.array(line.split(), dtype=float) radiance += data[2] variance += (data[3])**2 radiances[index] = data[2] tpp += data[6] varTimes += (data[7])**2 times[index] = data[6] index += 1 if any(np.isnan(radiances)): print("Nan in radiances") mean = radiance/nPixel variance /= nPixel varTimes /= nPixel meanTime = tpp/nPixel if nPixel == 1: nPixel += 1 deviation = np.sqrt(variance / (nPixel-1)) variance2 = ((radiances - mean)**2).sum() / nPixel deviation2 = np.sqrt(variance2 / (nPixel-1)) # print(deviation, deviation2) devTime = np.sqrt(varTimes / (nPixel-1)) variance2 = ((times - meanTime)**2).sum() / nPixel devTime2 = np.sqrt(variance2 / (nPixel-1)) if not time: return mean, max(deviation, deviation2) return mean, max(deviation, deviation2), meanTime, max(devTime, devTime2)
def __spectrum(self): self.radiances = np.zeros_like(self.script.wavelengths) self.stdDeviation = np.zeros_like(self.script.wavelengths) self.times = np.zeros_like(self.script.wavelengths) self.timeStdDev= np.zeros_like(self.script.wavelengths) for i,file in enumerate(self.script.outputFiles): self.radiances[i], self.stdDeviation[i], self.times[i], self.timeStdDev[i] = self.extractMeanRadianceFromOutput(file, time=True) radUnit = "W/m2/sr" if not hasattr(self.script, 'bandWidths'): self.radiances /= cst.nano self.stdDeviation /= cst.nano radUnit = "W/m2/sr/m" self.file = f"{self.resultsPath}{self.data.name}_reflectance_spectrum.json" data = { "nWavlength": len(self.script.wavelengths), "wavelength unit": "m", "radiance unit":radUnit, "wavelength": self.script.wavelengths.tolist(), "radiance": self.radiances.tolist(), "radiance std deviation": self.stdDeviation.tolist(), "time per path": self.times.tolist(), "time std deviation": self.timeStdDev.tolist(), } if self.script.predefinedScript == "reflectanceSpectrum": if 'spectrum' in self.script.geometry.source.keys(): wavelengths, spectrum = self.script.geometry.source['spectrum'] wavelengths *= cst.nano # m if hasattr(self.script, 'bandWidths'): bandLows = self.script.wavelengths - self.script.bandWidths / 2 bandUps = self.script.wavelengths + self.script.bandWidths / 2 flux = np.array([integral(wavelengths, spectrum, low, up)[0] for low,up in zip(bandLows, bandUps)]) else: flux = interpolate(wavelengths, spectrum, self.script.wavelengths) elif 'temperature' in self.script.source.keys(): flux = planck(self.script.geometry.source['temperature'], self.script.wavelengths, (self.script.geometry.source["radius"], self.script.geometry.source["distance"] ) ) self.reflectances = (self.radiances) / flux data["reflectance spectrum"] = self.reflectances.tolist(), data["reflectance std deviation"] = (self.stdDeviation/flux).tolist() with open(self.file, 'w') as f: f.write(json.dumps(data, indent=4)) return data def __imageRatio(self, threshold=0): if len(self.script.outputFiles) != 2: raise IndexError("Wrong number of files") numerator, numDev = self.getImage(self.script.outputFiles[0]) mn = threshold * numerator.max() print(mn) numerator = np.where(numerator>mn, numerator, 0) print(numerator) denominator, denDev = self.getImage(self.script.outputFiles[1]) mn = threshold * denominator.max() print(mn) denominator = np.where(denominator>mn, denominator, 0) print(denominator) ratio = numerator[:,:,1] / denominator[:,:,1] print(ratio) ratioMasked = np.where(np.isfinite(ratio), ratio, 1) self.fileName = f"{self.resultsPath}/{self.data.name}_{self.script.wavelengths[0]:.2e}_{self.script.wavelengths[1]:.2e}" np.savetxt(f"{self.fileName}.txt", ratioMasked, fmt="%5f") plt.imsave(f"{self.fileName}.png", ratioMasked, cmap="gray") return def __compositeRGB(self): if len(self.script.outputFiles) != 3: raise IndexError("Wrong number of files") red, numDev = self.getImage(self.script.outputFiles[0]) green, denDev = self.getImage(self.script.outputFiles[1]) blue, denDev = self.getImage(self.script.outputFiles[2]) image = np.array([red, green, blue]).T fileName = f"{self.resultsPath}/{self.data.name}_R{self.script.wavelengths[0]:.2e}_G{self.script.wavelengths[1]:.2e}_B{self.script.wavelengths[1]:.2e}" np.savetxt(f"{fileName}.txt", image, fmt="%5f") plt.imsave(f"{fileName}.png", image) return def __twoStreamV1(self, heatCapacity): results = self.extractMeanRadiances() fluxUp = np.zeros((self.data.atmosphereCellCoord.shape[0]+1, self.data.latitudes.shape[0], self.data.longitudes.shape[0])) fluxDown = np.zeros_like(fluxUp) for case, res in results.items(): k, i, j, cs = case.split('_')[:4] k, i, j = int(k), int(i), int(j) if cs == 'down1': fluxUp[k,i,j] += res["radiance"] #* self.script.angle * cst.degree elif cs == 'down2': fluxUp[k,i,j] -= res["radiance"] #* self.script.angle * cst.degree elif cs == 'top1': fluxDown[k,i,j] += res["radiance"] #* self.script.angle * cst.degree elif cs == 'top2': fluxDown[k,i,j] -= res["radiance"] #* self.script.angle * cst.degree fluxUp *= (2*np.pi * np.cos(self.script.angle * cst.degree)) fluxDown *= (2*np.pi * np.cos(self.script.angle * cst.degree)) z = self.data.atmosphereCellCoord[:,:,:,0] # array of altitudes dz = np.zeros_like(z) dz[1:-1,:,:] = (z[2:,:,:] - z[:-2,:,:]) * 0.5 dz[0,:,:] = z[0,:,:] - self.data.topoMap[::-1,:] dz[-1,:,:] = dz[-2,:,:] dfdz_up = (fluxUp[1:,:,:] - fluxUp[:-1,:,:]) / dz dfdz_down = (fluxDown[1:,:,:] - fluxDown[:-1,:,:]) / dz fluxDivergence = dfdz_down - dfdz_up p = self.data.pressure dp = np.zeros_like(p) dp[1:-1,:,:] = (p[2:,:,:] - p[:-2,:,:]) * 0.5 dp[0,:,:] = dp[1,:,:] dp[-1,:,:] = dp[-2,:,:] gravity = cst.G * self.data.mass / (z**2) rho = - (1/gravity) * (dp/dz) dTdt = fluxDivergence / (rho * heatCapacity) file = f"{self.resultsPath}twoStreamResults_{self.script.kind}.json" with open(file, 'w') as f: f.write(json.dumps({ 'cells coordonates': self.data.atmosphereCellCoord.tolist(), # shape = (nAlt, nLat, nLon, 3) 'flux up': fluxUp.tolist(), 'flux down': fluxDown.tolist(), 'flux divergence': fluxDivergence.tolist(), 'heating rate': dTdt.tolist() }, indent=4)) return def __twoStreamV2(self, heatCapacity): results = self.extractMeanRadiances() print(results) fluxUp = np.zeros((self.data.atmosphereCellCoord.shape[0]+1, self.data.latitudes.shape[0], self.data.longitudes.shape[0])) fluxDown = np.zeros_like(fluxUp) fluxDirect = np.zeros_like(fluxUp) fluxUpDev = np.zeros_like(fluxUp) fluxDownDev = np.zeros_like(fluxDown) fluxDirectDev = np.zeros_like(fluxDirect) for case, res in results.items(): k, i, j, cs = case.split('_')[-4:] k, i, j = int(k), int(i), int(j) if 'top' in cs: fluxDown[k,i,j] += res["radiance"] fluxDownDev[k,i,j] += res["std dev"]**2 elif 'down' in cs: fluxUp[k,i,j] += res["radiance"] fluxUpDev[k,i,j] += res["std dev"]**2 elif 'zenith' in cs: fluxDirect[k,i,j] += res["radiance"] fluxDirectDev[k,i,j]+= res["std dev"] fluxUp /= (self.script.nAngle) fluxDown /= (self.script.nAngle) fluxUpDev = np.sqrt(fluxUpDev) / self.script.nAngle fluxDownDev = np.sqrt(fluxDownDev) / self.script.nAngle # fluxUp *= ( 2*np.pi * np.cos(self.script.angle * cst.degree) ) # fluxDown *= ( 2*np.pi * np.cos(self.script.angle * cst.degree) ) # fluxUp *= ( np.cos(self.script.angle * cst.degree) ) # fluxDown *= ( np.cos(self.script.angle * cst.degree) ) # fluxUp /= np.cos(self.script.angle * cst.degree) # fluxDown /= np.cos(self.script.angle * cst.degree) # fluxUp *= np.pi # fluxDown *= np.pi # fluxDirect *= ( 2*np.pi * ( 1 - np.cos(self.script.FOVzenith*cst.degree/2) ) ) surface = (2 * self.script.source['distance'] * np.tan(self.script.FOVzenith*cst.degree/2))**2 fluxDirect *= ( surface / (self.script.source['distance']**2) ) fluxDirectDev *= ( surface / (self.script.source['distance']**2) ) cellCoord = self.data.atmosphereCellCoord z = cellCoord[:,:,:,0] - self.data.radius # array of altitudes dz = np.zeros_like(z) dz[1:-1,:,:] = (z[2:,:,:] - z[:-2,:,:]) * 0.5 dz[0,:,:] = (z[1,:,:] + z[0,:,:]) * 0.5 - self.data.topoMap[self.data.latitudes.shape[0]-i-1,j] dz[-1,:,:] = dz[-2,:,:] # # calculating the direct flux from the star transmitted through the column # starflux = np.pi * planck(self.script.source["temperature"], self.data.wavelengths*cst.nano, r_d=(self.script.source['radius'],self.script.source['distance'])) # # print(starflux[:]) # vecCol = np.zeros(cellCoord.shape[1:]) # # print(vecCol.shape) # for (i,j) in np.ndindex(vecCol.shape[0], vecCol.shape[1]): # vecCol[i,j,:] = sphere2cart(np.array([1, cellCoord[0,i,j,1], cellCoord[0,i,j,2]])) # vecSrc = sphere2cart(np.array([1, self.script.source["latitude"], self.script.source["longitude"]])) # cosIncAngle = np.dot(vecCol, vecSrc) # cosIncAngle = np.where(cosIncAngle >= 0, cosIncAngle, 0) # # print(cosIncAngle) # starflux = starflux[:,None,None,None] * cosIncAngle[None,:,:] # # print(starflux.flatten()) # ext = self.data.absorption[:,:,:,:,:]+ self.data.scattering[:,:,:,:,None] # tau = (ext * dz[None,:,:,:,None]) # shape = np.array(tau.shape) # shape[1] += 1 # we add one to the altitude to have the optical depth at the cells interface (the first being at zero: interface with space) # opticalDepth = np.zeros(shape) # opticalDepth[:,:-1,:,:,:] = np.cumsum(tau[:,::-1,:,:,:], axis=1)[:,::-1,:,:,:] # (nWvl, nLev, nLat, nLon, nWeight) # transmission = (np.exp(-opticalDepth/cosIncAngle[None,None,:,:, None]) * self.data.weights[:,None,None,None,:]).sum(axis=-1) # (nWvl, nAlt, nLat, nLon) # # print(ext.flatten()) # # print(dz.flatten()) # # print(tau.flatten()) # # print(opticalDepth.flatten()) # # print(transmission.flatten()) # if self.data.wavelengths.shape[0] == 1: directFlux = (transmission * starflux).reshape(transmission.shape[1:]) * self.data.wavelengths[0]*cst.nano # else: directFlux = np.trapz(transmission * starflux, self.data.wavelengths*cst.nano, axis=0) # print(directFlux.flatten()) # directFlux /= (4*np.pi) fluxDownTotal = fluxDown + fluxDirect # directFlux # dfdz_up = (fluxUp[1:,:,:] - fluxUp[:-1,:,:]) / dz # dfdz_down = (fluxDown[1:,:,:] - fluxDown[:-1,:,:]) / dz # fluxDivergence = dfdz_down - dfdz_up fluxDivergence = (fluxDownTotal[1:,:,:] + fluxUp[:-1,:,:] - fluxDownTotal[:-1,:,:] - fluxUp[1:,:,:]) / dz file = f"{self.resultsPath}twoStreamResults_{self.script.data.name}_{self.script.name}_{self.script.kind}.json" with open(file, 'w') as f: f.write(json.dumps({ 'cells coordinates': self.data.atmosphereCellCoord.tolist(), # shape = (nAlt, nLat, nLon, 3) 'flux up': fluxUp.tolist(), 'flux down': fluxDown.tolist(), 'direct flux': fluxDirect.tolist(), 'flux up dev': fluxUpDev.tolist(), 'flux down dev': fluxDownDev.tolist(), 'direct flux dev': fluxDirectDev.tolist(), 'flux divergence': fluxDivergence.tolist(), # 'heating rate': dTdt.tolist() }, indent=4)) return p = self.data.pressure dp = np.zeros_like(p) dp[1:-1,:,:] = (p[2:,:,:] - p[:-2,:,:]) * 0.5 dp[0,:,:] = dp[1,:,:] dp[-1,:,:] = dp[-2,:,:] gravity = cst.G * self.data.mass / (z**2) rho = - (1/gravity) * (dp/dz) dTdt = fluxDivergence / (rho * heatCapacity) file = f"{self.resultsPath}twoStreamResults_{self.script.kind}.json" with open(file, 'w') as f: f.write(json.dumps({ 'cells coordinates': self.data.atmosphereCellCoord.tolist(), # shape = (nAlt, nLat, nLon, 3) 'flux up': fluxUp.tolist(), 'flux down': fluxDown.tolist(), 'direct flux': directFlux.tolist(), 'flux divergence': fluxDivergence.tolist(), 'heating rate': dTdt.tolist() }, indent=4)) return def __twoStreamV3(self, heatCapacity): results = self.extractMeanRadiances() print(results) fluxUp = np.zeros((self.data.atmosphereCellCoord.shape[0]+1, self.data.latitudes.shape[0], self.data.longitudes.shape[0])) fluxDown = np.zeros_like(fluxUp) for case, res in results.items(): k, i, j, cs = case.split('_')[:4] k, i, j = int(k), int(i), int(j) if cs == 'top': fluxDown[k,i,j] = res["radiance"] * np.pi * (1 - np.cos(self.script.angle * cst.degree)) elif cs == 'down': fluxUp[k,i,j] = res["radiance"] * np.pi * (1 - np.cos(self.script.angle * cst.degree)) z = self.data.atmosphereCellCoord[:,:,:,0] # array of altitudes dz = np.zeros_like(z) dz[1:-1,:,:] = (z[2:,:,:] - z[:-2,:,:]) * 0.5 dz[0,:,:] = z[0,:,:] - self.data.topoMap[self.data.latitudes.shape[0]-i-1,j] dz[-1,:,:] = dz[-2,:,:] # dfdz_up = (fluxUp[1:,:,:] - fluxUp[:-1,:,:]) / dz # dfdz_down = (fluxDown[1:,:,:] - fluxDown[:-1,:,:]) / dz # fluxDivergence = dfdz_down - dfdz_up fluxDivergence = (fluxDown[1:,:,:] + fluxUp[:-1,:,:] - fluxUp[1:,:,:] - fluxDown[:-1,:,:]) / dz p = self.data.pressure dp = np.zeros_like(p) dp[1:-1,:,:] = (p[2:,:,:] - p[:-2,:,:]) * 0.5 dp[0,:,:] = dp[1,:,:] dp[-1,:,:] = dp[-2,:,:] gravity = cst.G * self.data.mass / (z**2) rho = - (1/gravity) * (dp/dz) dTdt = fluxDivergence / (rho * heatCapacity) file = f"{self.resultsPath}twoStreamResults_{self.script.kind}.json" with open(file, 'w') as f: f.write(json.dumps({ 'cells coordonates': self.data.atmosphereCellCoord.tolist(), # shape = (nAlt, nLat, nLon, 3) 'flux up': fluxUp.tolist(), 'flux down': fluxDown.tolist(), 'flux divergence': fluxDivergence.tolist(), 'heating rate': dTdt.tolist() }, indent=4)) return def __radiativeBudget(self, heatCapacity, rho): meansTotal = [] deviationTotal = [] sumXTotal = [] sumX2Total = [] meansDirect = [] deviationDirect = [] sumXDirect = [] sumX2Direct = [] meansDiffuse = [] deviationDiffuse = [] sumXDiffuse = [] sumX2Diffuse = [] numbers = [] with open(self.script.outputFiles[0], 'r') as f: for line in f: xT, sigT, sxT, sx2T, xDir, sigDir, sxDir, sx2Dir, xDiff, sigDiff, sxDiff, sx2Diff, N, dump, dump = line.split() meansTotal.append(float(xT)) deviationTotal.append(float(sigT)) sumXTotal.append(float(sxT)) sumX2Total.append(float(sx2T)) meansDirect.append(float(xDir)) deviationDirect.append(float(sigDir)) sumXDirect.append(float(sxDir)) sumX2Direct.append(float(sx2Dir)) meansDiffuse.append(float(xDiff)) deviationDiffuse.append(float(sigDiff)) sumXDiffuse.append(float(sxDiff)) sumX2Diffuse.append(float(sx2Diff)) numbers.append(int(N)) meansTotal = np.array(meansTotal) sumXTotal = np.array(sumXTotal) sumX2Total = np.array(sumX2Total) meansDirect = np.array(meansDirect) sumXDirect = np.array(sumXDirect) sumX2Direct = np.array(sumX2Direct) meansDiffuse = np.array(meansDiffuse) sumXDiffuse = np.array(sumXDiffuse) sumX2Diffuse = np.array(sumX2Diffuse) numbers = np.array(numbers) if self.script.geometry.GCMmesh: cellIDS = self.data.cellIds cellIndices = self.data.cellIndices latitudes = self.data.latitudes longitudes = self.data.longitudes altitudes = self.data.altitudes else: cellIDS = self.script.geometry.cellIds cellIndices = self.script.geometry.cellIndices latitudes = self.script.geometry.latitudes longitudes = self.script.geometry.longitudes altitudes = self.script.geometry.altitudes cellIndices = np.array(cellIndices).astype(int) shape = cellIndices.max(axis=0) + 1 meanCellTotal = np.zeros(shape) varianceCellTotal = np.zeros(shape) deviationCellTotal = np.zeros(shape) meanCellDirect = np.zeros(shape) varianceCellDirect = np.zeros(shape) deviationCellDirect = np.zeros(shape) meanCellDiffuse = np.zeros(shape) varianceCellDiffuse = np.zeros(shape) deviationCellDiffuse = np.zeros(shape) for cell,indices in zip(cellIDS, cellIndices): i,j,k = indices meanCellTotal[i,j,k], varianceCellTotal[i,j,k], deviationCellTotal[i,j,k] = combineEstimates(sumXTotal[cell], sumX2Total[cell], numbers[cell]) meanCellDirect[i,j,k], varianceCellDirect[i,j,k], deviationCellDirect[i,j,k] = combineEstimates(sumXDirect[cell], sumX2Direct[cell], numbers[cell]) meanCellDiffuse[i,j,k], varianceCellDiffuse[i,j,k], deviationCellDiffuse[i,j,k] = combineEstimates(sumXDiffuse[cell], sumX2Diffuse[cell], numbers[cell]) meanCellTotal[:,0,:] = meanCellTotal[:,0,0][:,None] varianceCellTotal[:,0,:] = varianceCellTotal[:,0,0][:,None] deviationCellTotal[:,0,:] = deviationCellTotal[:,0,0][:,None] meanCellTotal[:,-1,:] = meanCellTotal[:,-1,0][:,None] deviationCellTotal[:,-1,:] = deviationCellTotal[:,-1,0][:,None] varianceCellTotal[:,-1,:] = varianceCellTotal[:,-1,0][:,None] meanCellDirect[:,0,:] = meanCellDirect[:,0,0][:,None] varianceCellDirect[:,0,:] = varianceCellDirect[:,0,0][:,None] deviationCellDirect[:,0,:] = deviationCellDirect[:,0,0][:,None] meanCellDirect[:,-1,:] = meanCellDirect[:,-1,0][:,None] deviationCellDirect[:,-1,:] = deviationCellDirect[:,-1,0][:,None] varianceCellDirect[:,-1,:] = varianceCellDirect[:,-1,0][:,None] meanCellDiffuse[:,0,:] = meanCellDiffuse[:,0,0][:,None] varianceCellDiffuse[:,0,:] = varianceCellDiffuse[:,0,0][:,None] deviationCellDiffuse[:,0,:] = deviationCellDiffuse[:,0,0][:,None] meanCellDiffuse[:,-1,:] = meanCellDiffuse[:,-1,0][:,None] deviationCellDiffuse[:,-1,:] = deviationCellDiffuse[:,-1,0][:,None] varianceCellDiffuse[:,-1,:] = varianceCellDiffuse[:,-1,0][:,None] dic = { 'latitudes': latitudes.tolist(), 'longitudes': longitudes.tolist(), 'altitudes': altitudes.tolist(), 'mean': meanCellTotal.tolist(), 'variance': varianceCellTotal.tolist(), 'std deviation': deviationCellTotal.tolist(), 'mean direct': meanCellDirect.tolist(), 'variance direct': varianceCellDirect.tolist(), 'std deviation direct': deviationCellDirect.tolist(), 'mean diffuse': meanCellDiffuse.tolist(), 'variance diffuse': varianceCellDiffuse.tolist(), 'std deviation diffuse': deviationCellDiffuse.tolist() } if heatCapacity is not None and rho is not None: dT_dt_total = meanCellTotal / (rho * heatCapacity) dT_dt_direct = meanCellDirect / (rho * heatCapacity) dT_dt_diffuse = meanCellDiffuse / (rho * heatCapacity) dic['heating rate'] = dT_dt_total.tolist(), dic['heating rate deviation'] = (deviationCellTotal / (rho * heatCapacity)).tolist(), dic['heating rate direct'] = dT_dt_direct.tolist(), dic['heating rate direct deviation'] = (deviationCellDirect / (rho * heatCapacity)).tolist(), dic['heating rate diffuse'] = dT_dt_diffuse.tolist(), dic['heating rate diffuse deviation'] = (deviationCellDiffuse / (rho * heatCapacity)).tolist() with open(f"{self.resultsPath}radiativeBudget.json", 'w') as f: f.write(json.dumps(dic, indent=4)) else: with open(f"{self.resultsPath}radiativeBudget.json", 'w') as f: f.write(json.dumps(dic, indent=4)) return