expand_fluxes.py

#!/usr/bin/env python
# expand_fluxes.py
import sys
sys.path.append('../../')
import os
import getopt
from pylab import array, arange, transpose, date2num
from datetime import datetime
from da.tools.general import CreateDirs
from da.tools.io4 import CT_Read, CT_CDF, std_savedict, GetVariable

"""
Author: Wouter Peters (Wouter.Peters@noaa.gov)

Revision History:
File created on 21 Ocotber 2008.

"""

def proceed_dialog(txt, yes=['y','yes'], all=['a','all', 'yes-to-all']):
    """ function to ask whether to proceed or not """
    response=raw_input(txt)
    if response.lower() in yes:
       return 1
    if response.lower() in all:
       return 2
    return 0

def SaveWeeklyAvg1x1Data(DaCycle):
    """ Function SaveFlux1x1Data saves 3-hourly surface flux data to NetCDF files
        
        *** Inputs ***
        rundat : a RunInfo object

        *** Outputs ***
        daily NetCDF file containing 1-hourly global surface fluxes at 1x1 degree

        *** Example ***
        ./expand_savestate project=enkf_release sd=20000101 ed=20010101 """

    import da.tools.io4 as io
    import logging
#
    dirname     = 'data_flux1x1_weekly'
#
    dirname     = CreateDirs(os.path.join(DaCycle['dir.analysis'],dirname))
#
# Some help variables
#
    dectime0    = date2num(datetime(2000,1,1))
    dt          = DaCycle['cyclelength']
    startdate   = DaCycle['time.start'] 
    enddate     = DaCycle['time.end'] 

    msg = "DA Cycle start date is %s"   % startdate.strftime('%Y-%m-%d %H:%M')      ; logging.info(msg)
    msg = "DA Cycle end   date is %s"   % enddate.strftime('%Y-%m-%d %H:%M')        ; logging.info(msg)

#
# Create or open NetCDF output file
#
    saveas      = os.path.join(dirname,'flux_1x1.nc')
    ncf         = io.CT_CDF(saveas,'write')

#
# Create dimensions and lat/lon grid
#
    dimgrid     = ncf.AddLatLonDim()
    dimdate     = ncf.AddDateDim()
#
# set title and tell GMT that we are using "pixel registration"
#
    setattr(ncf,'Title','CarbonTracker fluxes')
    setattr(ncf,'node_offset',1)
#
# skip dataset if already in file
#
    ncfdate = date2num(startdate)-dectime0+dt.days/2.0
    skip    = ncf.has_date(ncfdate)
    if skip:
        msg = 'Skipping writing of data for date %s : already present in file %s' % (startdate.strftime('%Y-%m-%d'),saveas) ; logging.warning(msg)
    else:
        
#
# if not, process this cycle. Start by getting input data from CTDAS
#
        filename = os.path.join(DaCycle['dir.output'],'flux1x1_%s_%s.nc'%(startdate.strftime('%Y%m%d%H'),enddate.strftime('%Y%m%d%H'),) )

        file                = io.CT_Read(filename,'read')
        bio                 = array(file.GetVariable(DaCycle.DaSystem['background.co2.bio.flux']))
        ocean               = array(file.GetVariable(DaCycle.DaSystem['background.co2.ocean.flux']))
        fire                = array(file.GetVariable(DaCycle.DaSystem['background.co2.fires.flux']))
        fossil              = array(file.GetVariable(DaCycle.DaSystem['background.co2.fossil.flux']))
        mapped_parameters   = array(file.GetVariable(DaCycle.DaSystem['final.param.mean.1x1']))
        file.close()

        next=ncf.inq_unlimlen()[0]

# Start adding datasets from here on, both prior and posterior datasets for bio and ocn

        for prior in [False, True]:
#
# if prior, do not multiply fluxes with parameters, otherwise do
#
            if prior:
                qual_short='prior'
                biomapped=bio
                oceanmapped=ocean
            else:
                qual_short='opt'
                biomapped=bio*mapped_parameters 
                oceanmapped=ocean*mapped_parameters 
#
#  For each dataset, get the standard definitions from the module mysettings, add values, dimensions, and unlimited count, then write
#
            savedict=ncf.StandardVar(varname='bio_flux_'+qual_short)
            savedict['values']=biomapped.tolist()
            savedict['dims']=dimdate+dimgrid
            savedict['count']=next
            savedict['nsets']=1
            ncf.AddData(savedict)
#
            savedict=ncf.StandardVar(varname='ocn_flux_'+qual_short)
            savedict['name']='ocn_flux_'+qual_short
            savedict['values']=oceanmapped.tolist()
            savedict['dims']=dimdate+dimgrid
            savedict['count']=next
            ncf.AddData(savedict)
#
        savedict=ncf.StandardVar(varname='fire_flux_imp')
        savedict['values']=fire.tolist()
        savedict['dims']=dimdate+dimgrid
        savedict['count']=next
        ncf.AddData(savedict)
#
        savedict=ncf.StandardVar(varname='fossil_flux_imp')
        savedict['values']=fossil.tolist()
        savedict['dims']=dimdate+dimgrid
        savedict['count']=next
        ncf.AddData(savedict)
#
        savedict=ncf.StandardVar(varname='date')
        savedict['values']=date2num(startdate)-dectime0+dt.days/2.0
        savedict['dims']=dimdate
        savedict['count']=next
        ncf.AddData(savedict)

        sys.stdout.write('.')
        sys.stdout.flush()
#
#   Done, close the new NetCDF file
#
    ncf.close()
#
#   Return the full name of the NetCDF file so it can be processed by the next routine
#
    msg="Gridded weekly average fluxes now written" ; logging.info(msg)

    return saveas

def SaveWeeklyAvgTCData(DaCycle):
    """ Function SaveEcoData saves weekly surface flux data to NetCDF files
        
        *** Inputs ***
        DaCycle : a DaCycle object
    """
    from da.analysis.tools_regions import globarea, StateToGrid
    from da.analysis.tools_transcom import StateToTranscom, StateCovToTranscom, transcommask
    import da.tools.io4 as io
    from numpy import dot
    import logging
#
    dirname     = 'data_tc_weekly'
#
    dirname     = CreateDirs(os.path.join(DaCycle['dir.analysis'],dirname))
#
# Some help variables
#
    dectime0    = date2num(datetime(2000,1,1))
    dt          = DaCycle['cyclelength']
    startdate   = DaCycle['time.start'] 
    enddate     = DaCycle['time.end'] 
    ncfdate     = date2num(startdate)-dectime0+dt.days/2.0

    msg = "DA Cycle start date is %s"   % startdate.strftime('%Y-%m-%d %H:%M')      ; logging.info(msg)
    msg = "DA Cycle end   date is %s"   % enddate.strftime('%Y-%m-%d %H:%M')        ; logging.info(msg)

    # Write/Create NetCDF output file
    #
    saveas          = os.path.join(dirname,'tcfluxes.nc')
    ncf             = io.CT_CDF(saveas,'write')
    dimdate         = ncf.AddDateDim()
    dimidateformat  = ncf.AddDateDimFormat()
    dimregs         = ncf.AddRegionDim(type='tc')
#
# set title and tell GMT that we are using "pixel registration"
#
    setattr(ncf,'Title','CarbonTracker TransCom fluxes')
    setattr(ncf,'node_offset',1)
    #

    skip             = ncf.has_date(ncfdate)
    if skip:
        msg = 'Skipping writing of data for date %s : already present in file %s' % (startdate.strftime('%Y-%m-%d'),saveas) ; logging.warning(msg)
    else:

        # Get input data

        area=globarea()

        infile=os.path.join(DaCycle['dir.analysis'],'data_flux1x1_weekly','flux_1x1.nc')
        if not os.path.exists(infile):
            msg="Needed input file (%s) does not exist yet, please create weekly eco flux files first, returning..."%infile ; logging.error(msg)
            return None

        ncf_in      = io.CT_Read(infile,'read')

        # Transform data one by one

        # Get the date variable, and find index corresponding to the DaCycle date

        try:
            dates       = ncf_in.variables['date'][:]
        except KeyError:
            msg         = "The variable date cannot be found in the requested input file (%s) " % infile ; logging.error(msg)
            msg         = "Please make sure you create gridded fluxes before making TC fluxes " ; logging.error(msg)
            raise KeyError

        try:
            index       = dates.tolist().index(ncfdate)
        except ValueError:
            msg         = "The requested cycle date is not yet available as gridded flux in file %s "% infile ; logging.error(msg)
            msg         = "Please make sure you create gridded fluxes before making TC fluxes " ; logging.error(msg)
            raise ValueError

        # First add the date for this cycle to the file, this grows the unlimited dimension

        savedict            = ncf.StandardVar(varname='date')
        savedict['values']  = ncfdate
        savedict['dims']    = dimdate
        savedict['count']   = index
        dummy               = ncf.AddData(savedict)

        # Now convert other variables that were inside the flux_1x1 file

        vardict     = ncf_in.variables
        for vname, vprop in vardict.iteritems():

            data    = ncf_in.GetVariable(vname)[index]

            if vname=='latitude': continue
            elif vname=='longitude': continue

            if vname not in ['date','idate']:

                if 'cov' in vname: 
                    alldata=[] 
                    dd=StateCovToGrid(data*area,transcommask,reverse=True)
                    data=array(dd) 
                    dims=dimdate+dimregs+dimregs
                else:
                    dd=StateToGrid(data*area,transcommask,reverse=True)
                    data=array(dd) 
                    dims=dimdate+dimregs

                savedict            = ncf.StandardVar(varname=vname)
                savedict['values']  = data.tolist()
                savedict['dims']    = dims
                savedict['units']   = 'mol/region/s'
                savedict['count']   = index
                ncf.AddData(savedict)
             
        ncf_in.close()
    ncf.close()

    msg="TransCom weekly average fluxes now written" ; logging.info(msg)

    return saveas

def SaveTCDataExt(rundat):
    """ Function SaveTCDataExt saves surface flux data to NetCDF files for extended TransCom regions
        
        *** Inputs ***
        rundat : a RunInfo object

        *** Outputs ***
        NetCDF file containing n-hourly global surface fluxes per TransCom region

        *** Example ***
        ./expand_savestate project=enkf_release sd=20000101 ed=20010101 """
    from da.analysis.tools_transcom import StateCovToGrid, transcommask, ExtendedTCRegions
    import da.tools.io4 as io

    infile=os.path.join(rundat.outputdir,'data_tc_weekly','tcfluxes.nc')
    if not os.path.exists(infile):
        print "Needed input file (%s) does not exist yet, please create weekly tc flux files first, returning..."%infile
        return None

    # Create NetCDF output file
    #
    saveas          = os.path.join(rundat.outputdir,'data_tc_weekly','tc_extfluxes.nc')
    ncf             = io.CT_CDF(saveas,'create')
    dimdate         = ncf.AddDateDim()
    dimidateformat  = ncf.AddDateDimFormat()
    dimregs         = ncf.AddRegionDim(type='tc_ext')


    ncf_in      = io.CT_Read(infile,'read')
    vardict     = ncf_in.variables
    for vname, vprop in vardict.iteritems():

        data    = ncf_in.GetVariable(vname)[:]

        if vname=='latitude': continue
        elif vname=='longitude': continue
        elif vname=='date': dims=dimdate
        elif vname=='idate': dims=dimdate+dimidateformat

        if vname not in ['date','idate']:

                if 'cov' in vname: 
                    alldata=[] 
                    for dd in data:
                        dd=StateCovToGrid(dd*area,transcommask,reverse=True)
                        alldata.append(dd)
                    data=array(alldata) 
                    dims=dimdate+dimregs+dimregs
                else:
                    data=ExtendedTCRegions(data)
                    dims=dimdate+dimregs

        print vname,data.shape
        savedict            = ncf.StandardVar(varname=vname)
        savedict['values']  = data.tolist()
        savedict['dims']    = dims
        savedict['units']   = 'mol/region/s'
        savedict['count']   = 0
        ncf.AddData(savedict,nsets=data.shape[0])
         
    ncf.close()

    return saveas

def SaveEcoDataExt(rundat):
    """ Function SaveEcoDataExt saves surface flux data to NetCDF files for extended ecoregions
        
        *** Inputs ***
        rundat : a RunInfo object

        *** Outputs ***
        NetCDF file containing n-hourly global surface fluxes per TransCom region

        *** Example ***
        ./expand_savestate project=enkf_release sd=20000101 ed=20010101 """
    from ecotools import xte
    from da.analysis.tools_transcom import LookUpName
    import pycdf as CDF

    infile=os.path.join(rundat.outputdir,'data_eco_weekly','ecofluxes.nc')
    if not os.path.exists(infile):
        print "Needed input file (%s) does not exist yet, please create weekly ecoflux files first, returning..."%infile
        return None

    # Create NetCDF output file
    #
    saveas=os.path.join(rundat.outputdir,'data_eco_weekly','eco_extfluxes.nc')
    ncf=CT_CDF(saveas,'create')
    dimdate=ncf.AddDateDim()
    dimidateformat=ncf.AddDateDimFormat()
    dimregs         = ncf.AddRegionDim(type='tc_ext')


    ncf_in=CDF.CDF(infile)
    vardict=ncf_in.variables()
    for vname, vprop in vardict.iteritems():
        dims=()
        data=array(ncf_in.var(vname)[:])
        atts=ncf_in.var(vname).attributes()
        if vname not in ['date','idate']:
                if 'cov' in vname: 
                    data=xte(data[:,0:rundat.n_land,0:rundat.n_land],cov=True)
                    dims=dimdate+dimregs+dimregs
                else: 
                    data=xte(data[:,0:rundat.n_land])
                    dims=dimdate+dimregs
        elif vname=='date': dims=dimdate
        elif vname=='idate': dims=dimdate+dimidateformat
        else:
             print 'Dataset with unknown dimensions encountered in file: %s'%vname

        savedict=ncf.StandardVar(varname=vname)
        savedict['units']=atts['units']
        savedict['values']=data.tolist()
        savedict['dims']=dims
        ncf.AddData(savedict,nsets=data.shape[0])
         
    ncf.close()

    return saveas

def SaveTimeAvgData(rundat,infile,avg='monthly'):
    """ Function saves time mean surface flux data to NetCDF files
        
        *** Inputs ***
        rundat : a RunInfo object

        *** Outputs ***
        daily NetCDF file containing 1-hourly global surface fluxes at 1x1 degree

        *** Example ***
        ./expand_savestate project=enkf_release sd=20000101 ed=20010101 """

    import da.analysis.tools_time as timetools
    import da.tools.io4 as io

    dirname,filename    = os.path.split(infile)
    dirname             = CreateDirs(os.path.join(rundat.outputdir,dirname.replace('weekly',avg) ) )

    dectime0            = date2num(datetime(2000,1,1))

# Create NetCDF output file
#
    saveas              = os.path.join(dirname,filename)
    ncf                 = io.CT_CDF(saveas,'create')
    dimdate             = ncf.AddDateDim()
#
# Open input file specified from the command line
#
    if not os.path.exists(infile):
        print "Needed input file (%s) not found. Please create this first:"%infile
        print "returning..."
        return None 
    else:
        pass

    file        = io.CT_Read(infile,'read')
    datasets    = file.variables.keys()
    date        = file.GetVariable('date')
    globatts    = file.ncattrs()

    time        = [datetime(2000,1,1)+timedelta(days=d) for d in date]

#
# Add global attributes, sorted 
#

    keys        = globatts
    keys.sort()
    for k in keys:
        setattr(ncf,k,k)

# loop over datasets in infile, skip idate and date as we will make new time axis for the averaged data

    for sds in datasets:
        if sds in ['idate','date'] : continue
        print '['

# get original data

        data        = array(file.GetVariable(sds))
        varatts     = file.variables[sds].ncattrs()
        vardims     = file.variables[sds].dimensions
#
# Depending on dims of input dataset, create dims for output dataset. Note that we add the new dimdate now.
#

        dims=()
        for d in vardims:
            if 'latitude' in d: 
                dimgrid     = ncf.AddLatLonDim()
                dims        += (dimgrid[0],)
            if 'longitude' in d: 
                dimgrid     = ncf.AddLatLonDim()
                dims        += (dimgrid[1],)
            for type in ['eco','eco_ext','tc','tc_ext','olson']:
                if 'regions_%s' % type == d: 
                    dimregs     = ncf.AddRegionDim(type)
                    dims        += dimregs
#
# If the variable name from the infile is not defined in the standard, simply copy attributes from the old to the new sds
# Do not copy over the grid information
#

        if ncf.StandardVar(sds)['name'] == 'unknown':

                savedict                    = ncf.StandardVar(sds)
                savedict['name']            = sds
                savedict['values']          = data.tolist()
                savedict['dims']            = dims
                savedict['units']           = file.variables[sds].units
                savedict['long_name']       = file.variables[sds].long_name
                savedict['comment']         = file.variables[sds].comment
                savedict['standard_name']   = file.variables[sds].standard_name
                savedict['count']           =0
                ncf.AddData(savedict)
                sys.stdout.write('.')
                sys.stdout.flush()

        else:

            if sds in ['latitude','longitude']: continue

            if avg == 'monthly':
                time_avg, data_avg      = timetools.monthly_avg(time,data)
            elif avg == 'seasonal':
                time_avg, data_avg      = timetools.season_avg(time,data)
            elif avg == 'yearly':
                time_avg, data_avg      = timetools.yearly_avg(time,data)
            elif avg == 'longterm':
                time_avg, data_avg      = timetools.longterm_avg(time,data)
                time_avg                = [time_avg]
                data_avg                = [data_avg]
            else:
                raise ValueError,'Averaging (%s) does not exist' % avg

            count=-1
            for dd,data in zip(time_avg,data_avg):
                count=count+1

                if not ncf.has_date(date2num(dd)-dectime0): 

                    savedict=ncf.StandardVar('date')
                    savedict['values']=date2num(dd)-dectime0
                    savedict['dims']=dimdate
                    savedict['count']=count
                    ncf.AddData(savedict)

                savedict=ncf.StandardVar(sds)
                savedict['values']=data.tolist()
                savedict['units']=file.variables[sds].units
                if 'cov' in sds: 
                    savedict['dims']=dimdate+dims+dims
                else:
                    savedict['dims']=dimdate+dims
                savedict['count']=count
                ncf.AddData(savedict)

                sys.stdout.write('.')
                sys.stdout.flush()
        print ']'


# end NetCDF file access
    file.close()
    ncf.close()

if __name__ == "__main__":

    import logging
    from da.tools.initexit import CycleControl
    from da.ct.dasystem import CtDaSystem 

    sys.path.append('../../')

    logging.root.setLevel(logging.INFO)

    DaCycle = CycleControl(args={'rc':'../../da.rc'})
    DaCycle.Initialize()
    DaCycle.ParseTimes()

    DaSystem    = CtDaSystem('../rc/carbontracker.rc')
    DaSystem.Initialize()

    DaCycle.DaSystem    = DaSystem

    no, yes, all = range(3)

    proceed = no

    # first, parse results from the inverse part of the run 
    if proceed != all:
        proceed = proceed_dialog('Create average 1x1 flux data files? [y|yes, n|no, a|all|yes-to-all ]')
    if proceed != no:
        while DaCycle['time.start'] < DaCycle['time.finish']:

            savedas=SaveWeeklyAvg1x1Data(DaCycle)
            savedas=SaveWeeklyAvgTCData(DaCycle)

            DaCycle.AdvanceCycleTimes()

        sys.exit(2)

        a=SaveTimeAvgData(rundat,savedas,avg='monthly')
        a=SaveTimeAvgData(rundat,savedas,avg='yearly')
        a=SaveTimeAvgData(rundat,savedas,avg='longterm')

    if proceed != all:
        proceed = proceed_dialog('Create average tc flux data files? [y|yes, n|no, a|all|yes-to-all] ')
    if proceed != no:
        a=SaveTimeAvgData(rundat,savedas,avg='monthly')
        a=SaveTimeAvgData(rundat,savedas,avg='seasonal')
        a=SaveTimeAvgData(rundat,savedas,avg='yearly')
        a=SaveTimeAvgData(rundat,savedas,avg='longterm')
    if proceed != all:
        proceed = proceed_dialog('Create average extended tc flux data files? [y|yes, n|no, a|all|yes-to-all] ')
    if proceed != no:
        savedas=SaveTCDataExt(rundat)
        a=SaveTimeAvgData(rundat,savedas,avg='monthly')
        a=SaveTimeAvgData(rundat,savedas,avg='seasonal')
        a=SaveTimeAvgData(rundat,savedas,avg='yearly')
        a=SaveTimeAvgData(rundat,savedas,avg='longterm')
    sys.exit(0)