nc2kmz.py

   1 #!/usr/bin/env python
   2
   3 '''
   4 A simple python module for creating images out of netcdf arrays and outputing
   5 kml files for Google Earth.   The base class ncEarth cannot be used on its own,
   6 it must be subclassed with certain functions overloaded to provide location and
   7 plotting that are specific to a model's output files.
   8
   9 Requires matplotlib and netCDF4 python modules.
  10
  11 Use as follows:
  12
  13 import ncEarth
  14 kml=ncEarth.ncEpiSim('episim_0010.nc')
  15 kml.write_kml(['Susceptible','Infected','Recovered','Dead'])
  16
  17 or
  18
  19 kmz=ncEarth.ncWRFFire_mov('wrfout')
  20 kmz.write('FGRNHFX','fire.kmz')
  21
  22 Author: Jonathan Beezley (jon.beezley@gmail.com)
  23 Date: Oct 5, 2010
  24
  25 kmz=ncEarth.ncWRFFire_mov('wrfout')
  26 kmz.write_preload('FGRNHFX')
  27
  28 Modified by Lin Zhang
  29 Date: Dec 20, 2010
  30 '''
  31
  32
  33 import matplotlib
  34 try:
  35     matplotlib.use('Agg')
  36 except:
  37     pass
  38 from matplotlib import pylab
  39 from matplotlib.colorbar import ColorbarBase
  40 from matplotlib.colors import LogNorm,Normalize
  41 from matplotlib.ticker import LogFormatter
  42 import numpy as np
  43 try:
  44     from netCDF4 import Dataset
  45 except:
  46     from Scientific.IO.NetCDF import NetCDFFile as Dataset
  47 import cStringIO
  48 from datetime import datetime
  49 import zipfile
  50 import shutil,os
  51 import warnings
  52 import threading
  53
  54 try:
  55     ncpu=max(1,os.sysconf('SC_NPROCESSORS_ONLN'))
  56 except:
  57     ncpu=2
  58
  59 warnings.simplefilter('ignore')
  60
  61 global dlock
  62 global lock
  63 global queue
  64 global minmax
  65 global ncfile
  66 minmax={}
  67 lock=threading.RLock()
  68 dlock=threading.RLock()
  69 queue=threading.Semaphore(ncpu)
  70 ncfile={}
  71
  72 class ZeroArray(Exception):
  73     pass
  74
  75 class ncEarth(object):
  76
  77     '''Base class for reading NetCDF files and writing kml for Google Earth.'''
  78
  79     kmlname='ncEarth.kml'  # default name for kml output file
  80     progname='baseClass'   # string describing the model (overload in subclass)
  81
  82     # base kml file format string
  83     # creates a folder containing all images
  84     kmlstr= \
  85     '''<?xml version="1.0" encoding="UTF-8"?>
  86     <kml xmlns="http://www.opengis.net/kml/2.2">
  87     <Folder>
  88     <name>%(prog)s visualization</name>
  89     <description>Variables from %(prog)s output files visualized in Google Earth</description>
  90     %(content)s
  91     </Folder>
  92     </kml>'''
  93
  94     # string for static Ground Overlays
  95     kmlimageStatic= \
  96     '''<GroundOverlay>
  97       <name>%(name)s</name>
  98       <color>00ffffff</color>
  99       <Icon>
 100         <href>%(filename)s</href>
 101         <viewBoundScale>0.75</viewBoundScale>
 102       </Icon>
 103       <altitude>0.0</altitude>
 104       <altitudeMode>clampToGround</altitudeMode>
 105       <LatLonBox>
 106         <north>%(lat2)f</north>
 107         <south>%(lat1)f</south>
 108         <east>%(lon2)f</east>
 109         <west>%(lon1)f</west>
 110         <rotation>0.0</rotation>
 111       </LatLonBox>
 112     </GroundOverlay>'''
 113
 114     # format string for each image
 115     kmlimage= \
 116     '''<GroundOverlay>
 117       <name>%(name)s</name>
 118       <color>%(alpha)02xffffff</color>
 119       <Icon>
 120         <href>%(filename)s</href>
 121         <viewBoundScale>0.75</viewBoundScale>
 122       </Icon>
 123       <altitude>0.0</altitude>
 124       <altitudeMode>clampToGround</altitudeMode>
 125       <LatLonBox>
 126         <north>%(lat2)f</north>
 127         <south>%(lat1)f</south>
 128         <east>%(lon2)f</east>
 129         <west>%(lon1)f</west>
 130         <rotation>0.0</rotation>
 131       </LatLonBox>
 132       %(time)s
 133     </GroundOverlay>'''
 134
 135     kmlcolorbar= \
 136     '''
 137 <ScreenOverlay>
 138    <name>%(name)s colorbar</name>
 139    <color>ffffffff</color>
 140    <Icon>
 141       <href>%(file)s</href>
 142    </Icon>
 143    <overlayXY x=".15" y=".5" xunits="fraction" yunits="fraction"/>
 144    <screenXY x="0" y=".5" xunits="fraction" yunits="fraction"/>
 145    <rotationXY x="0" y="0" xunits="fraction" yunits="fraction"/>
 146    <size x="0" y=".75" xunits="fraction" yunits="fraction"/>
 147 </ScreenOverlay>
 148     '''
 149
 150     # time interval specification for animated output
 151     timestr=\
 152     '''<TimeSpan>
 153     %(begin)s
 154     %(end)s
 155     </TimeSpan>'''
 156
 157     beginstr='<begin>%s</begin>'
 158     endstr='<end>%s</end>'
 159
 160     def __init__(self,filename,hsize=5):
 161         '''Class constructor:
 162            filename : string NetCDF file to read
 163            hsize : optional, width of output images in inches'''
 164         global ncfile
 165         global lock
 166         with lock:
 167             if not ncfile.has_key(filename):
 168                 ncfile[filename]=Dataset(filename,'r')
 169         self.f=ncfile[filename]
 170         self.hsize=hsize
 171
 172     def get_minmax(self,vname):
 173         global minmax
 174         with lock:
 175             if minmax.has_key(vname):
 176                 mm=minmax[vname]
 177             else:
 178                 mm=self.compute_minmax(vname)
 179                 minmax[vname]=mm
 180         return mm
 181
 182     def compute_minmax(self,vname):
 183         v=self.f.variables[vname][:]
 184         return (v.min(),v.max())
 185
 186     def get_bounds(self):
 187         '''Return the latitude and longitude bounds of the image.  Must be provided
 188         by the subclass.'''
 189         raise Exception("Non-implemented base class method.")
 190
 191     def get_array(self,vname):
 192         '''Return a given array from the output file.  Must be returned as a
 193         2D array with top to bottom orientation (like an image).'''
 194         v=self.f.variables[vname]
 195         v=pylab.flipud(v)
 196         return v
 197
 198     def view_function(self,v):
 199         '''Any function applied to the image data before plotting.  For example,
 200         to show the color on a log scale.'''
 201         return v
 202
 203     def get_image(self,v,min,max):
 204         '''Create an image from a given data.  Returns a png image as a string.'''
 205
 206         # kludge to get the image to have no border
 207         fig=pylab.figure(figsize=(self.hsize,self.hsize*float(v.shape[0])/v.shape[1]))
 208         ax=fig.add_axes([0,0,1,1])
 209
 210         cmap=pylab.cm.jet
 211         cmap.set_bad('w',0.)
 212         norm=self.get_norm(min,max)
 213         ax.imshow(self.view_function(v),cmap=cmap,norm=norm)
 214         ax.axis('off')
 215         self.process_image()
 216
 217         # create a string buffer to save the file
 218         im=cStringIO.StringIO()
 219         fig.savefig(im,format='png',transparent=True)
 220         pylab.close(fig)
 221
 222         # return the buffer
 223         s=im.getvalue()
 224         im.close()
 225         return s
 226
 227     def get_colorbar(self,title,label,min,max):
 228         '''Create a colorbar from given data.  Returns a png image as a string.'''
 229
 230         fig=pylab.figure(figsize=(2,5))
 231         ax=fig.add_axes([0.35,0.03,0.1,0.9])
 232         norm=self.get_norm(min,max)
 233         formatter=self.get_formatter()
 234         if formatter:
 235             cb1 = ColorbarBase(ax,norm=norm,format=formatter,spacing='proportional',orientation='vertical')
 236         else:
 237             cb1 = ColorbarBase(ax,norm=norm,spacing='proportional',orientation='vertical')
 238         cb1.set_label(label,color='1')
 239         ax.set_title(title,color='1')
 240         for tl in ax.get_yticklabels():
 241             tl.set_color('1')
 242         im=cStringIO.StringIO()
 243         fig.savefig(im,dpi=300,format='png',transparent=True)
 244         pylab.close(fig)
 245         s=im.getvalue()
 246         im.close()
 247         return s
 248
 249     def get_norm(self,min,max):
 250         norm=Normalize(min,max)
 251         return norm
 252
 253     def get_formatter(self):
 254         return None
 255
 256     def process_image(self):
 257         '''Do anything to the current figure window before saving it as an image.'''
 258         pass
 259
 260     def get_kml_dict(self,name,filename,alpha=143):
 261         '''returns a dictionary of relevant info the create the image
 262         portion of the kml file'''
 263
 264         lon1,lon2,lat1,lat2=self.get_bounds()
 265         d={'lat1':lat1,'lat2':lat2,'lon1':lon1,'lon2':lon2, \
 266            'name':name,'filename':filename,'time':self.get_time(),'alpha':alpha}
 267         return d
 268
 269     def get_time(self):
 270         '''Return the time interval information for this image using the kml
 271         format string `timestr'.  Or an empty string to disable animations.'''
 272         return ''
 273
 274     def image2kmlStatic(self,varname,filename=None):
 275         '''Read data from the NetCDF file, create a psuedo-color image as a png,
 276         then create a kml string for displaying the image in Google Earth.  Returns
 277         the kml string describing the GroundOverlay.  Optionally, the filename
 278         used to write the image can be specified, otherwise a default will be used.'''
 279
 280         vdata=self.get_array(varname)
 281         min,max=self.get_minmax(varname)
 282         im=self.get_image(vdata,min,max)
 283         if filename is None:
 284             filename='%s.png' % varname
 285         f=open(filename,'w')
 286         f.write(im)
 287         f.close()
 288         d=self.get_kml_dict(varname,filename)
 289         pylab.close('all')
 290         return self.__class__.kmlimageStatic % d
 291
 292     def image2kml(self,varname,filename=None):
 293         '''Read data from the NetCDF file, create a psuedo-color image as a png,
 294         then create a kml string for displaying the image in Google Earth.  Returns
 295         the kml string describing the GroundOverlay.  Optionally, the filename
 296         used to write the image can be specified, otherwise a default will be used.'''
 297
 298         vdata=self.get_array(varname)
 299         min,max=self.get_minmax(varname)
 300         im=self.get_image(vdata,min,max)
 301         if filename is None:
 302             filename='%s.png' % varname
 303         f=open(filename,'w')
 304         f.write(im)
 305         f.close()
 306         d=self.get_kml_dict(varname,filename)
 307         return self.__class__.kmlimage % d
 308
 309     def colorbar2kml(self,varname,filename=None):
 310         min,max=self.get_minmax(varname)
 311         label=self.get_label(varname)
 312         cdata=self.get_colorbar(varname,label,min,max)
 313         if filename is None:
 314             filename='colorbar_%s.png' % varname
 315         f=open(filename,'w')
 316         f.write(cdata)
 317         f.close()
 318         pylab.close('all')
 319         return self.__class__.kmlcolorbar % {'name':varname,'file':filename}
 320
 321     def get_label(self,varname):
 322         return ''
 323
 324     def write_kml(self,varnames):
 325         '''Create the actual kml file for a list of variables by calling image2kml
 326         for each variable in a list of variable names.'''
 327         if type(varnames) is str:
 328             varnames=(varnames,)
 329         content=[]
 330         for varname in varnames:
 331             label=self.get_label(varname)
 332             content.append(self.image2kml(varname))
 333             content.append(self.colorbar2kml(varname))
 334         kml=self.__class__.kmlstr % \
 335                      {'content':'\n'.join(content),\
 336                       'prog':self.__class__.progname}
 337         f=open(self.__class__.kmlname,'w')
 338         f.write(kml)
 339         f.close()
 340
 341 class ncEarth_log(ncEarth):
 342     def view_function(self,v):
 343         if v.max() <= 0.:
 344             raise ZeroArray()
 345         v=np.ma.masked_equal(v,0.,copy=False)
 346         v.fill_value=np.nan
 347         v=np.log(v)
 348         return v
 349
 350     def get_norm(self,min,max):
 351         return LogNorm(min,max)
 352
 353     def get_formatter(self):
 354         return LogFormatter(10,labelOnlyBase=False)
 355
 356     def compute_minmax(self,vname):
 357         v=self.f.variables[vname][:]
 358         if v[v>0].size == 0:
 359             min=1e-6
 360             max=1.
 361         else:
 362             min=v[v>0].min()
 363             max=v.max()
 364         return (min,max)
 365
 366 class ncEpiSimBase(object):
 367     '''Epidemic model file class.'''
 368
 369     kmlname='epidemic.kml'
 370     progname='EpiSim'
 371
 372     def get_bounds(self):
 373         '''Get the lat/lon bounds of the output file... assumes regular lat/lon (no projection)'''
 374         lat=self.f.variables['latitude']
 375         lon=self.f.variables['longitude']
 376
 377         lat1=lat[0]
 378         lat2=lat[-1]
 379         lon1=lon[0]
 380         lon2=lon[-1]
 381
 382         return (lon1,lon2,lat1,lat2)
 383
 384 class ncEpiSim(ncEpiSimBase,ncEarth_log):
 385     pass
 386
 387 class ncWRFFireBase(object):
 388     '''WRF-Fire model file class.'''
 389
 390     kmlname='fire.kml'
 391     progname='WRF-Fire'
 392     wrftimestr='%Y-%m-%d_%H:%M:%S'
 393
 394     def __init__(self,filename,hsize=5,istep=0):
 395         '''Overloaded constructor for WRF output files:
 396            filename : output NetCDF file
 397            hsize : output image width in inches
 398            istep : time slice to output (between 0 and the number of timeslices in the file - 1)'''
 399         ncEarth.__init__(self,filename,hsize)
 400         self.istep=istep
 401
 402     def get_bounds(self):
 403         '''Get the latitude and longitude bounds for an output domain.  In general,
 404         we need to reproject the data to a regular lat/lon grid.  This can be done
 405         with matplotlib's BaseMap module, but is not done here.'''
 406
 407         lat=self.f.variables['XLAT'][0,:,:].squeeze()
 408         lon=self.f.variables['XLONG'][0,:,:].squeeze()
 409         dx=lon[0,1]-lon[0,0]
 410         dy=lat[1,0]-lat[0,0]
 411         #lat1=np.min(lat)-dy/2.
 412         #lat2=np.max(lat)+dy/2
 413         #lon1=np.min(lon)-dx/2.
 414         #lon2=np.max(lon)+dx/2
 415         lat1=lat[0,0]-dy/2.
 416         lat2=lat[-1,0]+dy/2.
 417         lon1=lon[0,0]-dx/2.
 418         lon2=lon[0,-1]+dx/2.
 419         return (lon1,lon2,lat1,lat2)
 420
 421     def isfiregrid(self,vname):
 422         xdim=self.f.variables[vname].dimensions[-1]
 423         return xdim[-7:] == 'subgrid'
 424
 425     def srx(self):
 426         try:
 427             s=len(self.f.dimensions['west_east_subgrid'])/(len(self.f.dimensions['west_east'])+1)
 428         except:
 429             s=(self.f.dimensions['west_east_subgrid'])/((self.f.dimensions['west_east'])+1)
 430         return s
 431
 432     def sry(self):
 433         try:
 434             s=len(self.f.dimensions['south_north_subgrid'])/(len(self.f.dimensions['south_north'])+1)
 435         except:
 436             s=(self.f.dimensions['south_north_subgrid'])/((self.f.dimensions['south_north'])+1)
 437         return s
 438
 439     def get_array(self,vname):
 440         '''Return a single time slice of a variable from a WRF output file.'''
 441         v=self.f.variables[vname]
 442         v=v[self.istep,:,:].squeeze()
 443         if self.isfiregrid(vname):
 444             v=v[:-self.sry(),:-self.srx()]
 445         v=pylab.flipud(v)
 446         if vname == 'FGRNHFX' or vname == 'GRNHFX':
 447             v[:]=v*0.239005736
 448         return v
 449
 450     def get_dates(self):
 451         t1=self.f.variables["Times"][0,:].tostring()
 452         t2=self.f.variables["Times"][-1,:].tostring()
 453         return '%s - %s' % (t1,t2)
 454
 455     def get_time(self):
 456         '''Process the time information from the WRF output file to create a
 457         proper kml TimeInterval specification.'''
 458         global dlock
 459         start=''
 460         end=''
 461         time=''
 462         g=self.f
 463         times=g.variables["Times"]
 464         if self.istep > 0:
 465             with dlock:
 466                 start=ncEarth.beginstr % \
 467                datetime.strptime(times[self.istep,:].tostring(),\
 468                                      self.__class__.wrftimestr).isoformat()
 469         if self.istep < times.shape[0]-1:
 470             with dlock:
 471                 end=ncEarth.endstr % \
 472                datetime.strptime(times[self.istep+1,:].tostring(),\
 473                                      self.__class__.wrftimestr).isoformat()
 474         if start is not '' or end is not '':
 475             time=ncEarth.timestr % {'begin':start,'end':end}
 476         return time
 477
 478     def get_label(self,varname):
 479         v=self.f.variables[varname]
 480         return v.units
 481
 482 class ncWRFFire(ncWRFFireBase,ncEarth):
 483     pass
 484
 485 class ncWRFFireLog(ncWRFFireBase,ncEarth_log):
 486     pass
 487
 488 def create_image(fname,istep,nstep,vname,vstr,logscale,colorbar,imgs,content):
 489     global lock
 490     global queue
 491     queue.acquire()
 492     i=istep
 493     if logscale:
 494         kml=ncWRFFireLog(fname,istep=istep)
 495     else:
 496         kml=ncWRFFire(fname,istep=istep)
 497     if colorbar:
 498         img='files/colorbar_%s.png' % vname
 499         img_string=kml.colorbar2kml(vname,img)
 500         with lock:
 501             content.append(img_string)
 502             imgs.append(img)
 503     try:
 504         img=vstr % (vname,istep)
 505         img_string=kml.image2kml(vname,img)
 506         with lock:
 507             content.append(img_string)
 508             imgs.append(img)
 509             print 'creating frame %i of %i' % (i,nstep)
 510     except ZeroArray:
 511         with lock:
 512             print 'skipping frame %i of %i' % (i,nstep)
 513     queue.release()
 514
 515 class ncWRFFire_mov(object):
 516
 517     '''A class the uses ncWRFFire to create animations from WRF history output file.'''
 518
 519     def __init__(self,filename,hsize=5,nstep=None):
 520         '''Class constructor:
 521            filename : NetCDF output file name
 522            hsize : output image width in inces
 523            nstep : the number of frames to process (default all frames in the file)'''
 524
 525         self.filename=filename
 526         f=Dataset(filename,'r')
 527         g=f
 528         self.nstep=nstep
 529         if nstep is None:
 530             # in case nstep was not specified read the total number of time slices from the file
 531             self.nstep=g.variables['Times'].shape[0]
 532
 533     def write_preload(self,vname,kmz='fire_preload.kmz'):
 534         '''Create a kmz file from multiple time steps of a wrfout file. The kml file consists of a set of
 535         GroundOverlays with time tag and a copy of the set without the time tag to preload the
 536         images that are used in the GroundOverlays.'''
 537
 538
 539         imgs=[]     # to store a list of all images created
 540         content=[]  # the content of the main kml
 541         vstr='files/%s_%05i.png' # format specification for images (all stored in `files/' subdirectory)
 542
 543         # create empty files subdirectory for output images
 544         try:
 545             shutil.rmtree('files')
 546         except:
 547             pass
 548         os.makedirs('files')
 549
 550         # loop through all time slices and create the image data
 551         # appending to the kml content string for each image
 552         for i in xrange(0,self.nstep,1):
 553             print i
 554             kml=ncWRFFire(self.filename,istep=i)
 555             img=vstr % (vname,i)
 556             imgs.append(img)
 557             content.append(kml.image2kmlStatic(vname,img))
 558             kml.f.close()
 559
 560         # create the main kml file
 561         kml=ncWRFFire.kmlstr % \
 562             {'content':'\n'.join(content),\
 563              'prog':ncWRFFire.progname}
 564
 565         # create a zipfile to store all images + kml into a single compressed file
 566         z=zipfile.ZipFile(kmz,'w',compression=zipfile.ZIP_DEFLATED)
 567         z.writestr(kmz[:-3]+'kml',kml)
 568         for img in imgs:
 569             z.write(img)
 570         z.close()
 571
 572     def write(self,vname,kmz='fire.kmz',hsize=5,logscale=True,colorbar=True):
 573         '''Create a kmz file from multiple time steps of a wrfout file.
 574         vname : the variable name to visualize
 575         kmz : optional, the name of the file to save the kmz to'''
 576
 577         imgs=[]     # to store a list of all images created
 578         content=[]  # the content of the main kml
 579         threads=[]
 580         vstr='files/%s_%05i.png' # format specification for images (all stored in `files/' subdirectory)
 581
 582         # create empty files subdirectory for output images
 583         try:
 584             shutil.rmtree('files')
 585         except:
 586             pass
 587         os.makedirs('files')
 588
 589         # loop through all time slices and create the image data
 590         # appending to the kml content string for each image
 591         #k=0
 592         for i in xrange(0,self.nstep,1):
 593             t=threading.Thread(target=create_image,args=(self.filename,i,self.nstep,vname,vstr,logscale,colorbar and i == 0,imgs,content))
 594             t.start()
 595             threads.append(t)
 596         for t in threads:
 597             t.join()
 598
 599         # create the main kml file
 600         kml=ncWRFFire.kmlstr % \
 601             {'content':'\n'.join(content),\
 602              'prog':ncWRFFire.progname}
 603
 604         # create a zipfile to store all images + kml into a single compressed file
 605         z=zipfile.ZipFile(kmz,'w',compression=zipfile.ZIP_DEFLATED)
 606         z.writestr(kmz[:-3]+'kml',kml)
 607         for img in imgs:
 608             z.write(img)
 609         z.close()
 610
 611
 612 def uselog(vname):
 613     if vname in ('FGRNHFX','GRNHFX'):
 614         return True
 615     else:
 616         return False
 617
 618 if __name__ == '__main__':
 619     import sys
 620     if len(sys.argv) < 2:
 621         print "Takes a WRF-Fire output file and writes fire.kmz."
 622         print "usage: %s filename"%sys.argv[0]
 623     else:
 624         filename=sys.argv[1]
 625         if len(sys.argv) == 2:
 626             vars=('FGRNHFX',)
 627         else:
 628             vars=sys.argv[2:]
 629         kmz=ncWRFFire_mov(filename)
 630         for v in vars:
 631             kmz.write(v,hsize=8,kmz='fire_'+v+'.kmz',logscale=uselog(v))