Merge pull request #133 from AetherModel/hdf5_post

Add Hdf5 in post processing to allow people to use hdf5 files instead of netcdf files.

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 ...')

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
+