merge_afrt.py

Go to the documentation of this file.

#!/usr/bin/env python3
'''
* NAME: merge_afrt.py
*
* DESCRIPTION: Executable program to merge afrt (radiative
* transfer) NetCDF4 files
 
* Created on October 09, 2018
 
* Author: Samuel Anderson
* heavily modified on August 05, 2019 by Sean Bailey
'''
 
import os
import sys
import re
import time
from datetime import datetime as dt
import argparse
import logging
LOG = logging.getLogger('afrt-merge')
from netCDF4 import Dataset 
import numpy as np
import collections
 
 
 
#sensorInfoFile = os.path.join(os.environ['OCDATAROOT'] , 'common','SensorInfo.json')
#with open(sensorInfoFile, 'r') as myfile:
#    sensorDefs=myfile.read()
#
 
 
sensors = collections.defaultdict(dict)
 
sensors['modisa']['name'] = "MODIS-Aqua"
sensors['modisa']['instrument'] = "MODIS"
sensors['modisa']['platform'] = "Aqua"
sensors['modist']['name'] = "MODIS-Terra"
sensors['modist']['instrument'] = "MODIS"
sensors['modist']['platform'] = "Terra"
sensors['seawifs']['name'] = "SeaWiFS"
sensors['seawifs']['instrument'] = "SeaWiFS"
sensors['seawifs']['platform'] = "Orb-View 2"
sensors['meris']['name'] = "MERIS"
sensors['meris']['instrument'] = "MERIS"
sensors['meris']['platform'] = "ENVISAT"
sensors['octs']['name'] = "OCTS"
sensors['octs']['instrument'] = "OCTS"
sensors['octs']['platform'] = "ADEOS-I"
sensors['czcs']['name'] = "CZCS"
sensors['czcs']['instrument'] = "CZCS"
sensors['czcs']['platform'] = "NIMBUS-7"
sensors['viirsn']['name'] = "VIIRS-SNPP"
sensors['viirsn']['instrument'] = "VIIRS"
sensors['viirsn']['platform'] = "SNPP"
sensors['viirsj1']['name'] = "VIIRS-JPSS1"
sensors['viirsj1']['instrument'] = "VIIRS"
sensors['viirsj1']['platform'] = "JPSS1"
sensors['aviris']['name'] = "AVIRIS"
sensors['aviris']['instrument'] = "AVIRIS"
sensors['aviris']['platform'] = "ER-2"
sensors['hico']['name'] = "HICO"
sensors['hico']['instrument'] = "HICO"
sensors['hico']['platform'] = "ISS"
sensors['goci']['name'] = "GOCI"
sensors['goci']['instrument'] = "GOCI"
sensors['goci']['platform'] = "COMSAT-2"
sensors['misr']['name'] = "MISR"
sensors['misr']['instrument'] = "MISR"
sensors['misr']['platform'] = "Terra"
 
def sortbymodel(nameRegEx,fname):
    parsed = re.match(nameRegEx,fname)
    if parsed:
        return int(parsed.group(2)) + 0.1*int(parsed.group(3))  
    else:
        return 0
 
def createMergedNetCDF(ofile,dimensions,variables,sensor):
 
    LOG.info("Writing netCDF file: %s" % ofile)
    # Create the netCDF file and add dimensions
    if os.path.exists(ofile):
        LOG.error("%s exists! Bailing out..." % ofile)
        return None
 
    ncfile = Dataset(ofile, 'w',format='NETCDF4')
    for dim in dimensions:
        ncfile.createDimension(dimensions[dim]['name'], dimensions[dim]['size'])
 
    for var in variables:
        LOG.debug("var: %s dimensions: %s" %(var,variables[var]['dimensions']))
        LOG.debug(variables[var]['chunksize'])
        vardims = []
        for dim in variables[var]['dimensions']:
            vardims.append(dimensions[dim]['name'])
            
        ncfile.createVariable(var.lower(), variables[var]['dtype'], 
                              tuple(vardims), zlib=True, 
                              chunksizes=variables[var]['chunksize'])
 
    # Set global attributes
    ncfile.title = "Ahmad-Fraser Radiatitve Transfer Results"
    ncfile.description = "Ahmad-Fraser Radiatitve Transfer results for " + sensors[sensor]['name']
    ncfile.date_created = (dt.fromtimestamp(time.time()).strftime('%Y-%m-%dT%H:%M:%SZ')).encode('ascii')
    ncfile.creator_name = "NASA/GSFC/OBPG"
    ncfile.creator_email = "data@oceancolor.gsfc.nasa.gov" 
    ncfile.creator_url = "https://oceancolor.gsfc.nasa.gov" 
    ncfile.product_name=os.path.basename(ofile).encode("ascii")
    ncfile.history = ' '.join(sys.argv).encode("ascii")
 
 
    return ncfile
 
 
 
def main():
   
    description = '''Merge AFRT NetCDF output files.'''
    __version__ = "1.0.0"
 
    parser =  argparse.ArgumentParser(description=description)
    parser.add_argument('--version', action='version', version='%(prog)s ' + __version__)
    parser.add_argument('-i','--input-dir',
                      dest="inputDir" ,
                      type=str,
                      required=True,
                      help="The input directory")
    parser.add_argument('-s','--sensor',
                        dest="sensor",
                        choices=['modisa',
                                 'modist',
                                 'viirsn',
                                 'viirsj1',
                                 'seawifs',
                                 'meris',
                                 'aviris',
                                 'octs',
                                 'czcs',
                                 'misr',
                                 'oci',
                                 'goci',
                                 'hico'],
                        default='modis-aqua', help='Sensor for which to generate the table')
    # Optional arguments
    parser.add_argument('-n','--nameregex',
                        dest="nameRegEx",
                        type=str,
                        default= "^.*_WL(\d+)_MD(\d+)_TA(\d+)_WD(\d+)-(\d+).nc$",
                        help="regular expression for parsing filename to get parameter dimensions")
 
    parser.add_argument('-o','--output-dir',
                      dest="outputDir" ,
                      type=str,
                      default=os.getcwd(),
                      help="The output directory")
 
    parser.add_argument('-v', '--verbose',
                      dest='verbosity',
                      action="count",
                      default=0,
                      help='each occurrence increases verbosity 1 level from ERROR: -v=WARNING -vv=INFO -vvv=DEBUG')
 
 
    # Parse the arguments from the command line
    args = parser.parse_args()
 
    # Set up the logging levels
    levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
    logging.basicConfig(level = levels[args.verbosity])
 
    # Check for mutually exclusive options
    LOG.info('Merging radiative transfer tables in  %s' % (args.inputDir) )
 
    if os.path.isdir(args.outputDir):
        o_path = args.outputDir
    elif os.path.isdir(os.path.dirname(args.outputDir)):
        o_path = args.outputDir
        os.mkdir(o_path)            
    else:
        print ("Output path invalid")
        return
    if not os.path.isdir(args.inputDir):
        os.mkdir(args.inputDir)            
#    if args.fname:
#        with open(args.fname) as fp: 
#            contents = fp.readlines()
#            for fn in contents:
#                fn = fn.lstrip()
#                fn = fn.strip("\n")
#                command = "wget -N --content-disposition https://oceandata.sci.gsfc.nasa.gov/cgi/gethiddenfile/" + fn 
#                print (command)               
#                subprocess.run(command, shell=True)
 
    listing = os.listdir(args.inputDir)
    listing.sort(key=lambda fname:sortbymodel(args.nameRegEx,fname))   
    wvls = set()
    models = set()
    winds = set()
    taus = set()  
    inputFiles = []
    for fname in listing:   
        fpath = args.inputDir + "/" + fname        
        if os.path.isdir(fpath):
            continue
        parsed = re.match(args.nameRegEx,fname)
        if parsed:
            inputFiles.append(fpath)
            wvls.add(int(parsed.group(1)))
            models.add(int(parsed.group(2)))
            taus.add(int(parsed.group(3)))
            winds.add(int(parsed.group(4)))
            winds.add(int(parsed.group(5)))
 
    variables = collections.defaultdict(dict) 
    dimensions = collections.defaultdict(dict)
 
    dimensions['dim_size']['name'] = 'nmodels'
    dimensions['dim_size']['size'] = max(models)
    dimensions['dim_wave']['name'] = 'nwavelengths'
    dimensions['dim_wave']['size'] = max(wvls)
    dimensions['dim_wind']['name'] = 'nwindspeeds'
    dimensions['dim_wind']['size'] = max(winds)
    dimensions['dim_tau']['name'] = 'ntaus'
    dimensions['dim_tau']['size'] = max(taus)
 
    rtFile = Dataset(inputFiles[0],'r')
    dimensions['dim_phi']['name'] = 'nphi'
    dimensions['dim_phi']['size'] = rtFile.dimensions['dim_phi'].size
    dimensions['dim_sza']['name'] = 'nsolz'
    dimensions['dim_sza']['size'] = rtFile.dimensions['dim_sza'].size
    dimensions['dim_the']['name'] = 'nsenz'
    dimensions['dim_the']['size'] = rtFile.dimensions['dim_the'].size
    dimensions['dim_stokes']['name'] = 'nstokes'
    dimensions['dim_stokes']['size'] = rtFile.dimensions['dim_stokes'].size
 
    for var in rtFile.variables:
        lcvar = var.lower()
        if lcvar == 'aot':
            lcvar = 'tau_aerosol'
        if var == 'THETA':
            lcvar = 'sensor_zenith'
        if var == 'SZA':
            lcvar = 'solar_zenith'
        if var == 'PHI':
            lcvar = 'relative_azimuth'
        # skip a few...tau_aerosol - it's the same as AOT
        # sbar, tau_h20 and tau_ozone are currently zero
        # tau_total...yeah, we can do simple addition...
 
        if var in ('TAU_AEROSOL','TAU_H2O','TAU_OZONE','TAU_TOTAL','SBAR'):
            continue
        if 'index' not in var:
            variables[lcvar]['chunksize'] = None
            if var == 'TAU_RAYLEIGH':
                variables[lcvar]['dimensions'] = rtFile.variables['WAVELENGTH'].dimensions
            else:
                variables[lcvar]['dimensions'] = rtFile.variables[var].dimensions
            if var in ('WIND_SPEED','SZA','THETA','PHI'):
                variables[lcvar]['dtype'] = 'u2'
            else:    
                variables[lcvar]['dtype'] = 'f4'
                if var in ('BOA_DOWN_DIRECT','SBAR','BOA_DOWN_DIFFUSE','TOA_UP_DIFFUSE','HEM_REFLECTANCE','TRANSMITTANCE'):
                    variables[lcvar]['chunksize'] = (1,1,1,1,dimensions['dim_sza']['size'])
                if var in ('TOA_RADIANCE','BOA_RADIANCE'):
                  variables[lcvar]['chunksize'] = (1,1,1,1,dimensions['dim_sza']['size'], dimensions['dim_the']['size'],1,1)
 
    rtFile.close()
 
    rstr = "_WL%(l0)d-%(l1)d_MD%(s0)d-%(s1)d_TA%(t0)d-%(t1)d_WD%(w0)d-%(w1)d" % \
    {"l0":min(wvls),"l1":max(wvls),"s0":min(models),"s1":max(models),"t0":min(taus),"t1":max(taus),"w0":min(winds),"w1":max(winds)}
 
    ofilepath = args.outputDir + "/" + args.sensor + rstr + '.nc'
    ncFile = createMergedNetCDF(ofilepath,dimensions,variables,args.sensor)
 
    if ncFile:
        fixed = True
        tauset = set()
        wvlset = set()
        rayset = set()
        fcnt = 0
        for fpath in inputFiles:
            fcnt += 1
            LOG.info("Processing [%d of %d] %s ..." % (fcnt,len(inputFiles),fpath))
            if fpath != inputFiles[0]:
                fixed = False
            
            rtFile = Dataset(fpath,'r')
            wvl_idx = np.unique(rtFile.variables['wavelength_index'][:]) -1
            model_idx = np.unique(rtFile.variables['model_index'][:]) - 1
            tau_idx = np.unique(rtFile.variables['tau_index'][:]) - 1
            wind_idx = np.unique(rtFile.variables['wind_index'][:]) - 1
            # TODO: should the order of the processing change in ODPS from the 
            # current 'windspeed as the last element', the dimension indexing 
            # will need to be modified to accomodate.  The current assumption:
            # 1 model, 1 wavelength, 1 tau, N windsppeds
            LOG.debug("wvl_idx: %d\tmodel_idx: %d\ttau_idx: %d" % (wvl_idx,model_idx,tau_idx))
            LOG.debug(wind_idx)
 
 
            for var in rtFile.variables:
                if var in ('TAU_AEROSOL','TAU_H2O','TAU_OZONE','TAU_TOTAL','SBAR'):
                    continue
                lcvar = var.lower()
                # These are fixed for all inputs, so only do the first
                if var == 'THETA' and fixed:
                    ncFile.variables['sensor_zenith'][:] = rtFile.variables[var][:]
                if var == 'SZA' and fixed:
                    ncFile.variables['solar_zenith'][:] = rtFile.variables[var][:]
                if var == 'PHI' and fixed:
                    ncFile.variables['relative_azimuth'][:] = rtFile.variables[var][:]
                # TODO: since the current ODPS processing runs all windspeeds in a single run,
                # wind_speed is also fixed.  Should the order change, this will need to be
                # changed as well...
                if var == 'WIND_SPEED' and fixed:
                    ncFile.variables['wind_speed'][:] = rtFile.variables[var][:]
 
                if var == 'WAVELENGTH':
                    wvldata = rtFile.variables[var][:]
                    if wvldata[0] not in wvlset:
                        ncFile.variables[lcvar][wvl_idx] = wvldata[0]
                        wvlset.add(wvldata[0])
 
                if var == 'AOT':
                    taudata = rtFile.variables[var][:]
                    if taudata[0] not in tauset:
                        ncFile.variables['tau_aerosol'][tau_idx] = taudata[0]
                        tauset.add(taudata[0])
 
                if var in ('TAU_RAYLEIGH'):
                    raydata = rtFile.variables[var][0,:,0,0]
                    if raydata[0] not in rayset:
                        ncFile.variables[lcvar][wvl_idx] = raydata[0]
                        rayset.add(raydata[0])
 
                if var in ('BOA_DOWN_DIRECT','SBAR','BOA_DOWN_DIFFUSE','TOA_UP_DIFFUSE','HEM_REFLECTANCE','TRANSMITTANCE'):
                    ncFile.variables[lcvar][model_idx,wvl_idx,tau_idx,:,:] = rtFile.variables[var][:]
 
                if var in ('TOA_RADIANCE','BOA_RADIANCE'):
                    ncFile.variables[lcvar][model_idx,wvl_idx,tau_idx,:,:,:,:] = rtFile.variables[var][:]
 
            rtFile.close()
 
        ncFile.close()
 
        LOG.info("Merge complete for " + ofilepath)
    LOG.info('Exiting...')
    return 0
 
if __name__=='__main__':
    sys.exit(main())