root/sodar/trunk/plotSodar.py

#!/usr/bin/python
"""
Module to process sodar data.
"""

import optparse
from sodar.utils import openAnything, findMissing 
from sodar import arrayData as a
import numpy as n
import pylab as p
import os
import shutil
import csv

_uComponentsPlotName = 'uComponents.png'
_vComponentsPlotName = 'vComponents.png'
_wComponentsPlotName = 'wComponents.png'
_echoStrengthsPlotName = 'echoStrengths.png'
_quiverPlotName = 'quiver.png'
_latestPlotNamePrefix = 'ims_sodar_latest_'
_catalogName = 'sodar-plot-catalog.csv'

def plotSingle(filein, pathout):
    """
    Plot single sodar raw data file.
    
    plotSingle(filein, pathout) -> rc
    
    filein - path to raw data file.
    pathout - path to plot images directory.
    rc - 0 if successful
    """
    
    if not os.path.exists(pathout):
        os.makedirs(pathout, mode=0775)
    
    rawDataHandle = openAnything.openAnything(filein)
    rawDataString = rawDataHandle.read()
    rawDataHandle.close()
    
    rawDataObject = a.rawData.RawData(rawDataString)
    formattedDataObject = a.formattedData.FormattedData(rawDataObject)
    arrayDataObject = a.ArrayData(formattedDataObject)
    
    beginStamp = arrayDataObject.beginStamp
    endStamp = arrayDataObject.endStamp
    numIntervals = arrayDataObject.numIntervals
    timeInterval = arrayDataObject.timeInterval
    minAltitude = arrayDataObject.minAltitude
    numAltitudes = arrayDataObject.numAltitudes
    altInterval = arrayDataObject.altInterval
    
    def makePcolor(vector, title, fileout, cbLabel):
        fig = p.figure(1)
        axe = fig.add_subplot(1, 1, 1)
        pc = axe.pcolor(vector)
        
        axe.set_xlabel('Time (hh:mm UTC)')
        axe.set_ylabel('Altitude (m)')
        axe.set_xbound(upper=numIntervals)
        
        xticks = axe.get_xticks()
        xticklabels = [(int(x) * timeInterval) + beginStamp
                       for x in xticks]
        xticklabels = [':'.join(('%02u' % x.hour, '%02u' % x.minute))
                       for x in xticklabels]
        axe.set_xticklabels(xticklabels)
        
        yticks = axe.get_yticks()
        yticklabels = [str(y * altInterval + minAltitude)
                       for y in yticks]
        axe.set_yticklabels(yticklabels)
        
        axe.set_title(title)
        cb = p.colorbar(pc)
        cb.set_label(cbLabel)
        
        fig.savefig(os.path.join(pathout, fileout))
        fig.clear()
        
        return 0
    
    uComponents = arrayDataObject.uComponents
    maskedUComponents = n.ma.masked_where(n.isnan(uComponents),
                                          uComponents)
    makePcolor(maskedUComponents.T,
               'U Component of Wind Velocity for %s' %
               (str(beginStamp)[:10],),
               _uComponentsPlotName,
               'Speed (cm/sec)')
    
    vComponents = arrayDataObject.vComponents
    maskedVComponents = n.ma.masked_where(n.isnan(vComponents),
                                          vComponents)
    makePcolor(maskedVComponents.T,
               'V Component of Wind Velocity for %s' %
               (str(beginStamp)[:10],),
               _vComponentsPlotName,
               'Speed (cm/sec)')
    
    wComponents = arrayDataObject.wComponents
    maskedWComponents = n.ma.masked_where(n.isnan(wComponents),
                                          wComponents)
    makePcolor(maskedWComponents.T,
               'W Component of Wind Velocity for %s' %
               (str(beginStamp)[:10],),
               _wComponentsPlotName,
               'Speed (cm/sec)')
    
    echoStrengths = arrayDataObject.echoStrengths
    maskedEchoStrenths = n.ma.masked_where(n.isnan(echoStrengths),
                                                   echoStrengths)
    makePcolor(maskedEchoStrenths.T,
               'Echo Strength for %s' %
               (str(beginStamp)[:10],),
               _echoStrengthsPlotName,
               'Strength (no units)')
    
    # timeComponent = n.array(range(numIntervals))
    # altComponent = n.array(range(numAltitudes))
    # fig = p.figure(1)
    # axe = fig.add_subplot(1, 1, 1)
    # qv = axe.quiver(altComponent,
    #                 timeComponent,
    #                 maskedUComponents.T,
    #                 maskedVComponents.T)
    
    # axe.set_xlabel('Time (min)')
    # axe.set_ylabel('Altitude (m)')
    # axe.set_xbound(upper=numIntervals)
    
    # xticks = axe.get_xticks()
    # xticklabels = [(int(x) * timeInterval) + beginStamp
    #                for x in xticks]
    # xticklabels = [':'.join(('%02u' % x.hour, '%02u' % x.minute))
    #                for x in xticklabels]
    # axe.set_xticklabels(xticklabels)
    
    # yticks = axe.get_yticks()
    # yticklabels = [str(y * altInterval + minAltitude)
    #                for y in yticks]
    # axe.set_yticklabels(yticklabels)
    
    # axe.set_title('Wind Velocty for %s' %
    #               (str(beginStamp)[:10],))
    # cb = p.colorbar(qv)
    # cb.set_label('W Component Speed (cm/sec)')
    
    # fig.savefig(os.path.join(pathout, _quiverPlotName))
    # fig.clear()
    
    return 0

def copyLatest(pathout, latest, plotName):
    """
    Copy the latest plot to the official NCCOOS place.
    
    copyLatest(pathout, latest, plotName) -> rc
    
    pathout - path to directory where the latest plot is.
    latest - path to directory where the latest plot is to be copied.
    plotName - name of the plot file to copy.
    rc - 0 if successful
    """
    shutil.copy(os.path.join(pathout, plotName),
                os.path.join(latest,
                             _latestPlotNamePrefix + plotName))
    return 0

def createCatalog(destination, catalog):
    """
    Create catalog of plots.
    
    createCatalog(destination, catalog) -> rc
    
    destination - path to plot images directory in NCCOOS format.
    catalog - path to plot catalog.
    rc - 0 if successful
    """
    destinationWalk = findMissing.computeDestinationWalk(destination, False)
    catalogData = [(destpath.split(os.sep))[-1]
                   for destpath in destinationWalk]
    catalogHandle = open(os.path.join(catalog, _catalogName), 'wb')
    writer = csv.writer(catalogHandle)
    writer.writerows(catalogData)
    catalogHandle.close()
    
    return 0

def plotAll(source, destination, latest, catalog, force):
    """
    Plot all sodar raw data files.
    
    plotAll(source, destination, latest, catalog, force) -> rc
    
    source - path to raw data directory in NCCOOS format.
    destination - path to plot images directory in NCCOOS format.
    latest - path to latest plots images directory in NCCOOS format.
    catalog - path to plot catalog.
    force - update all destination plots for which raw data sources exist.
    rc - 0 if successful
    """
    
    if source.endswith(os.path.sep):
        source = source.rstrip(os.path.sep)
    
    if destination.endswith(os.path.sep):
        destination = destination.rstrip(os.path.sep)
    
    walkList = findMissing.findMissing(source, destination, force)
    
    for filein, pathout in walkList:
        plotSingle(filein, pathout)
    
    if not os.path.exists(latest):
        os.makedirs(latest, mode=0775)

    filein, pathout = walkList[-1]
    copyLatest(pathout, latest, _uComponentsPlotName)
    copyLatest(pathout, latest, _vComponentsPlotName)
    copyLatest(pathout, latest, _wComponentsPlotName)
    copyLatest(pathout, latest, _echoStrengthsPlotName)
    # copyLatest(pathout, latest, _quiverPlotName)
    
    createCatalog(destination, catalog)
    
    return 0

def _main():
    """bin/python %prog [options] \
                  /path/to/raw/files/ \
                  /path/to/plots/ \
                  /path/to/latest/plot/ \
                  /path/to/catalog/"""
    
    __description__ = 'Plot all sodar raw data files.'
    
    parser = optparse.OptionParser(usage=_main.__doc__,
                                   version='%prog 1.0',
                                   description=__description__)
    parser.set_defaults(forceUpdate=False)
    parser.add_option('-f', '--force-update',
                      action='store_true',
                      dest='force',
                      help='update all plots for which raw data exists')
    (values, args) = parser.parse_args()
    (source, destination, latest, catalog) = tuple(args)
    plotAll(source, destination, latest, catalog, values.force)
    
    return 0

if __name__ == "__main__":
    _main()