diff --git a/srcPython/postAether.py b/srcPython/postAether.py index 613dd843..07933189 100755 --- a/srcPython/postAether.py +++ b/srcPython/postAether.py @@ -9,6 +9,7 @@ import numpy as np import matplotlib.cm as cm from netCDF4 import Dataset +from h5py import File import argparse import os import json @@ -67,6 +68,8 @@ def epoch_to_datetime(epoch_time): return dtime +# ---------------------------------------------------------------------- +# Define a class for helping handle data class DataArray(np.ndarray): def __new__(cls, input_array, attrs={}): @@ -435,6 +438,9 @@ def read_aether_file(filename, file_vars=None, epoch_name='time'): def parse_args(): parser = argparse.ArgumentParser(description = 'Post process Aether files') + parser.add_argument('-hdf5', \ + help='output HDF5 files', \ + action="store_true") parser.add_argument('-rm', \ help='removes processed files', \ action="store_true") @@ -748,6 +754,52 @@ def write_netcdf(allBlockData, fileName, isVerbose = True): ncfile.close() +#---------------------------------------------------------------------------- +# Write a hdf5 file from the data +#---------------------------------------------------------------------------- + +def write_hdf5(allBlockData, fileName, isVerbose = True): + + if (isVerbose): + print(' Outputting file : ', fileName) + hdf5file = File(fileName, 'w') + + nBlocks = len(allBlockData) + nLons, nLats, nZ = get_sizes(allBlockData) + + lon_dim = hdf5file.create_dataset('nlons', data = [nLons]) + lat_dim = hdf5file.create_dataset('nlats', data = [nLats]) + z_dim = hdf5file.create_dataset('nalts', data =[nZ]) + block_dim = hdf5file.create_dataset('nblocks', data = [nBlocks]) + + oneBlock = allBlockData[0] + + time_out = [datetime_to_epoch(oneBlock["time"])] + hdf5file.create_dataset('time', data = time_out, dtype = np.float64) + + allHDF5Datasets = [] + # create all of the variables + varList = [] + for iV, v in enumerate(oneBlock['vars']): + if (v == 'alt'): + v = 'z' + varList.append(v) + if ('long_name' in oneBlock): + longName = oneBlock['long_name'][iV] + else: + longName = v + unitName = oneBlock['units'][iV] + allHDF5Datasets.append(hdf5file.create_dataset(v, dtype = np.float32, \ + shape = (nBlocks, nLons, nLats, nZ))) + allHDF5Datasets[-1].units = unitName + allHDF5Datasets[-1].long_name = longName + + for iB, oneBlock in enumerate(allBlockData): + tmp = np.asarray(oneBlock[iV]) + allHDF5Datasets[-1][iB,:,:,:] = tmp + + hdf5file.close() + #---------------------------------------------------------------------------- # copy block data in one file #---------------------------------------------------------------------------- @@ -868,11 +920,17 @@ def write_and_plot_data(dataToWrite, fileAddon, iVar, iAlt, + output_netcdf, isVerbose = True): - netcdfFile = fileStart + fileAddon + '.nc' - print(' --> Outputting nc file : ', netcdfFile) - write_netcdf(dataToWrite, netcdfFile, isVerbose = isVerbose) + if output_netcdf: + netcdfFile = fileStart + fileAddon + '.nc' + print(' --> Outputting nc file : ', netcdfFile) + write_netcdf(dataToWrite, netcdfFile, isVerbose = isVerbose) + else: + hdf5File = fileStart + fileAddon + '.hdf5' + print(' --> Outputting hdf5 file : ', hdf5File) + write_hdf5(dataToWrite, hdf5File, isVerbose = isVerbose) plotFile = fileStart + fileAddon + '.png' var = dataToWrite[0]['vars'][iVar] @@ -895,6 +953,8 @@ def write_and_plot_data(dataToWrite, iVar = 3 iAlt = 10 + output_netcdf = False if args.hdf5 else True + for iFile, fileInfo in enumerate(filesInfo): coreFile = fileInfo['coreFile'] isNetCDF = fileInfo['isNetCDF'] @@ -903,7 +963,7 @@ def write_and_plot_data(dataToWrite, isVerbose = isVerbose) write_and_plot_data(allBlockData, coreFile, '', iVar, iAlt, - isVerbose = isVerbose) + output_netcdf, isVerbose = isVerbose) if (fileInfo['isEnsemble']): factor = 1.0 / float(fileInfo['ensembleMembers']) @@ -919,13 +979,13 @@ def write_and_plot_data(dataToWrite, if (fileInfo['ensembleNumber'] == fileInfo['ensembleMembers']): write_and_plot_data(ensembleData, fileInfo['ensembleFile'], - '_mean', iVar, iAlt) + '_mean', iVar, iAlt, output_netcdf) stdData = calc_std_of_ensembles(filesInfo, ensembleIndexList, ensembleData) write_and_plot_data(stdData, fileInfo['ensembleFile'], - '_std', iVar, iAlt) + '_std', iVar, iAlt, output_netcdf) if (args.rm): print(' --> Removing files ...') diff --git a/tests/hdf5/run_test.sh b/tests/hdf5/run_test.sh new file mode 100755 index 00000000..59f0901d --- /dev/null +++ b/tests/hdf5/run_test.sh @@ -0,0 +1,17 @@ +#!/bin/sh + +# simple test to make sure that the hdf5 output and plotting works ok! + +rm -rf run.hdf5 + +cp -r ../../share/run ./run.hdf5 +cd run.hdf5 +mpirun -np 4 ./aether + +cd UA/output + +../../../../../srcPython/postAether.py -rm -hdf5 +~/bin/run_plot_block_model_results.py -var=Temperature -alt=30 3DALL_20110320_001000.hdf5 + +ls +