source: TI02-CSML/trunk/csml/csmllibs/netCDFWriter.py @ 2633

import cdms,MV
import Numeric
import csml
import sys

class NCwriter(object):
    #This provides a simplified wrapper to CDMS to write a CF compliant NetCDF
    def __init__(self, fileName):
        #initiates a new NetCDF file
        self.nc=cdms.open(fileName,'w')
    
    def setGlobalAttributes(self, **kwargs):
        #sets global attributes
        pass
    
    def addAxis(self,axisName, data,isLon=None,isLat=None,isTime=None,**kwargs):
        #first ensure tuple does not contains a list
        if type(data[0]) is list:
            newdata=tuple(data[0])
            data=newdata
            
        # Now create the axis   
        dataarray=MV.array(data)
        ax=cdms.createAxis(dataarray)
        ax.id = axisName
        if not hasattr(self,'axes'):
            self.axes=[]
        if isLon is not None:
            ax.designateLongitude()
        elif isLat is not None:
            ax.designateLatitude()
        elif isTime is not None:
            ax.designateTime()
        for key in kwargs:
            setattr(ax, key,kwargs[key])
        self.axes.append(ax)
        
    
    def getAxis(self, axID):
        for axis in self.axes:
            if axis.id == axID:
                return axis
    
    
    
    def addTimeVariable(self,timeval):  
        dataarray=MV.array(timeval)
        #dataarray.id='time'
        #dataarray.name ='time'
        ax=self.getAxis('time')
        dataarray.setAxis(0,ax)            
        self.nc.write(dataarray)
        
        
    def addVariable(self,data, variableName,axesList, fillvalue,  **kwargs):        
        #creates a new variable containing data with named attributes from **kwargs       
             
        dataarray=MV.array(data)      
        dataarray.id=variableName
        dataarray.name=variableName
        for key in kwargs:
            setattr(dataarray, key,kwargs[key])
        axisCount=0
      
        #depending on whether time is modelled in the underlying data as a dimension or not the shapes may not match
        if len(dataarray.shape) < len(axesList):  
            #take time out and create a separate variable
            newaxesList=[]
            for a in axesList:
                if a != 'time':
                    newaxesList.append(a)
            axesList=newaxesList
        if hasattr(self, 'axes'):    
            for axis in axesList:
                for ax in self.axes:
                    if ax.id == axis:
                        dataarray.setAxis(axisCount,ax)                       
                        axisCount = axisCount +1      
        setattr(dataarray, 'missing_value' ,fillvalue)
        self.nc.write(dataarray)

    def genWriteVar(self,varid, ordinates, times, caltype, axisorder, unitlist, fulldata, fillvalue, **kwargs):
        #**kwargs may contain additional axes not contained in a GridCoordinatesTable - e.g latitude, longitude may be stored in other attribute not in the rectified grid.     
        axesdone=[]
        floatTimes=[]
        #determine base units for times:
        tOne=csml.csmllibs.csmltime.getCDtime(times[0])
        tbase=csml.csmllibs.csmltime.getBaseUnits(tOne)
        for time in times:
            time=csml.csmllibs.csmltime.getCDtime(time).torel(tbase)
            floatTimes.append(time.value)
        self.addAxis('time',floatTimes,isTime=1,units=tbase,calendar=caltype)
  
        axesdone.append('time')
        if ordinates is not None:
            for ord in enumerate(ordinates):
                vals=[]
                lon,lat=None,None
                if ord[1].coordAxisLabel.CONTENT=='time':
                    continue
                else:
                    for val in ord[1].coordAxisValues.coordinateList.CONTENT.split():
                        if val != ' ':
                            vals.append(eval(val))  
                if ord[1].coordAxisLabel.CONTENT=='longitude':
                    lon=1
                    name='longitude'
                elif ord[1].coordAxisLabel.CONTENT=='latitude':
                    lat=1
                    name='latitude'
                else:
                    name=ord[1].coordAxisLabel.CONTENT
                for ax in enumerate(axisorder):
                    if ax[1]==name:
                        position=ax[0]
                axesdone.append(name)
                self.addAxis(name,vals,isLon=lon,isLat=lat,units=unitlist[position])#to do, units attribute for CF compliance
        
        if kwargs is not None:   
            for kw in kwargs:
                lon,lat=None,None
                if kw=='longitude':
                    lon=1
                elif kw == 'latitude':
                    lat=1
                name = kw
                for ax in enumerate(axisorder):
                    if ax[1]==name:
                        position=ax[0]
                vals=(kwargs[kw],)
                self.addAxis(name,vals,isLon=lon,isLat=lat,units=unitlist[position])
                axesdone.append(name)        
        self.addVariable(fulldata,varid, axisorder, fillvalue, units=unitlist[-1] ) #to do, units attribute for CF compliance

        for ax in self.axes:
            if ax.id =='time':
                if self.nc.variables.has_key('time') is False:
                    self.addTimeVariable(floatTimes)
        
        
    def closeFinishedFile(self):
        #returns finished (hopefully) NetCDF file
        self.nc.close()