source: trunk/MESOSCALE/PLOT/PYTHON/pywrf.example.r39/wrf/utils.py @ 189

"""
repository of wrf utilities
"""
import os
import sys
import numpy as n
import pylab as p
from matplotlib.numerix import ma
from matplotlib.toolkits.basemap import Basemap
import pywrf.viz.utils as vu


# the syntax to import nio depends on its version. We try all options from the
# newest to the oldest
try:
    import Nio as nio
except:
    try:
        import PyNGL.Nio as nio
    except:
        import PyNGL_numpy.Nio as nio

# The following is only needed for hn3.its.monash.edu.au
# Unfortunately we must get rid of spurious entries in the sys.path that can
# lead to the loading of conflicting packages
for dir_idx in range(len(sys.path)-1,-1,-1):
    if 'python2.4' in sys.path[dir_idx]:
        sys.path.pop(dir_idx)

a_small_number = 1e-8

def colons_to_underscores(test=False):
    files = os.listdir('.')
    print 'I plan to do the following:'
    for file in files:
        command = 'mv ' + file + ' ' + file[:-6] + '_00_00'
        print command
    print 'Happy? [y/n]'
    answer = raw_input()
    if answer == 'y':
        for file in files:
            command = 'mv ' + file + ' ' + file[:-6] + '_00_00'
            os.system(command)
    else:
        print 'OK, maybe next time ;]' 
    return

def wrf_to_pressure(var, pressure, press_lvl, fill_value=1e35, positive_only=False):
    """
    this functions vertically interpolates to pressure levels WRF fields
    usage 
    >>> interpolated_var(var, pressure, press_lvl, fill_value=1e35, 
    ...    positive_only=False)
    where 
        var -> field to be interpolated
        pressure -> total pressure field (p + pb)
        press_lvl -> list or numpy array of desired levels
        fill_value -> this is assigned to a cell if the requested pressure level
          lies below or above the column of values supplied for that cell in 
          the pressure array.
        positive_only -> set this flag to true to prevent the interpolation to 
          generate negative values for fields for which this does not make 
          sense.
        interpolated_var -> the interpolated field which will have shape 
          (len(press_lvl), var.shape[1], var.shape[2])
    NB in this current implementation of the function arrays need to be 
      destaggerred before they operated upon. Furthermore, it is assumed the
      function is invoked to process one frame at at time i.e. it works on 3D
      arrays. As usual COARDS ordering of dimensions i.e. (time,lvl,lat,lon) 
      is assumed.
    """
    # these are the pressure levels we want as output
    k_max, j_max, i_max = pressure.shape
    output = n.zeros((len(press_lvl), j_max, i_max))
    for lvl_idx in range(len(press_lvl)):
        for j in range(j_max):
            for i in range(i_max):
                # make it ascending for searchsorted
                pressure_column = pressure[::-1,j,i]
                if press_lvl[lvl_idx] > pressure_column.max() \
                  or press_lvl[lvl_idx] < pressure_column.min():
                    output[lvl_idx,j,i] = fill_value
                else:
                    var_column = var[::-1,j,i]
                    press_idx = pressure_column.searchsorted(press_lvl[lvl_idx])
                    xa = pressure_column[press_idx]
                    xb = pressure_column[press_idx - 1]
                    ya = var_column[press_idx]
                    yb = var_column[press_idx - 1]
                    x = press_lvl[lvl_idx]
                    y = vu.lin_interp(xa, xb, ya, yb, x)
                    if positive_only and y < 0.:
                        y = 0.
                    output[lvl_idx,j,i] = y
    return output

def wrf2latlon(projection, 
  lat_0, lon_0,
  lat_1,
  lat_2,
  grid_centre_lon,
  grid_centre_lat,
  nx,
  ny,
  delta,
  staggered = False,
  return_extra = False
  ):
    from pyproj import Proj
    proj = Proj(proj=projection, lat_0=lat_0, lon_0=lon_0, lat_1=lat_1, lat_2=lat_2)
    grid_centre_x, grid_centre_y = proj(grid_centre_lon, grid_centre_lat)
    grid_x_extent = nx * delta
    grid_y_extent = ny * delta
    min_x = grid_centre_x - grid_x_extent/2.
    min_y = grid_centre_y - grid_y_extent/2.
    max_x = min_x + grid_x_extent
    max_y = min_y + grid_y_extent
    x = n.arange(min_x, max_x + a_small_number, delta)
    y = n.arange(min_y, max_y + a_small_number, delta)  
    X, Y = p.meshgrid(x,y)
    lon, lat = proj(X, Y, inverse=True)

    if staggered:
        x_u = n.arange(min_x, max_x + delta + a_small_number, delta)
        x_u -= (delta /2.)
        X_u, Y_u = p.meshgrid(x_u,y)
        y_v = n.arange(min_y, max_y + delta + a_small_number, delta)  
        y_v -= (delta /2.)
        X_v, Y_v = p.meshgrid(x,y_v)
        lon_u, lat_u = proj(X_u, Y_u, inverse=True)
        lon_v, lat_v = proj(X_v, Y_v, inverse=True)

    llcrnrlon, llcrnrlat = proj(min_x, min_y, inverse=True)
    urcrnrlon, urcrnrlat = proj(max_x, max_y, inverse=True)
    map = Basemap(
      projection = projection,
      lon_0 = lon_0,
      lat_0 = lat_0,
      lat_1 = lat_1,
      lat_2 = lat_2,
      llcrnrlon = llcrnrlon,
      llcrnrlat = llcrnrlat,
      urcrnrlon = urcrnrlon,
      urcrnrlat = urcrnrlat
      )

    if return_extra:
        # it seems that the basemap automatically sets the origin the native
        # coordinate system at its llcrnr (or close to it...)
        offset = map(lon_0, lat_0)
        if staggered:
            X += offset[0]
            Y += offset[1]
            X_u += offset[0]
            Y_u += offset[1]
            X_v += offset[0]
            Y_v += offset[1]
            return lon, lat, lon_u, lat_u, lon_v, lat_v, \
              X, Y, X_u, Y_u, X_v, Y_v, map
        else:
            X += offset[0]
            Y += offset[1]
            return lon, lat, X, Y, map
    else: 
        if staggered:
            return lon, lat, lon_u, lat_u, lon_v, lat_v
        else:
            return lon, lat
  
def wrf_grid(
  # WPS -> map_proj
  projection,
  # WPS -> truelat1
  lat_1,
  # WPS -> truelat2
  lat_2,
  # WPS -> stand_lon
  lon_0,
  # WPS -> ref_lat
  grid_centre_lat,
  # WPS -> ref_lon
  grid_centre_lon,
  delta_x,
  delta_y,
  # WPS -> e_we
  nx,
  # WPS -> e_sn
  ny,
  show_mass_grid = False,
  show_stag_grids = False,
  ):
    if lon_0 != grid_centre_lon:
        print 'not implemented yet -> see the source'
        print "\tbut let's try it anyways..."
        #return

    width   = nx * delta_x
    height  = ny * delta_y
    frame_x = 10 * delta_x
    frame_y = 10 * delta_y
    m = Basemap(
      lat_0 = grid_centre_lat,
      # this could be a bad assumption... because lon_0 and grid_centre_lon
      # need not be aligned, but at the same time I need to give this to
      # basemap for the grid to be centred... I could probably fix it
      # assigning lon_0 and then imposing a grid shift in native coordinates
      # if ref_lon and lon_0 were not the same
      lon_0 = lon_0,
      lat_1 = lat_1,
      lat_2 = lat_2,
      width = width + 2*frame_x,
      height = height + 2*frame_y,
      resolution = 'l',
      area_thresh=1000.
      )
    grid_centre_x, grid_centre_y = m(grid_centre_lon, grid_centre_lat)
    min_x = grid_centre_x - width/2.
    min_y = grid_centre_y - height/2.
    max_x = min_x + width
    max_y = min_y + height
    x = n.arange(min_x, max_x + a_small_number, delta_x)
    y = n.arange(min_y, max_y + a_small_number, delta_y)  
    x = x[1:-1]
    y = y[1:-1]
    x_u = n.arange(min_x, max_x + delta_x + a_small_number, delta_x)
    x_u -= delta_x/2.
    x_u = x_u[1:-1]
    y_v = n.arange(min_y, max_y + delta_y + a_small_number, delta_y)
    y_v -= delta_y/2.
    y_v = y_v[1:-1]
    X, Y = p.meshgrid(x,y)
    lon, lat = m(X, Y, inverse=True)
    X_u, Y_u = p.meshgrid(x_u,y)
    lon_u, lat_u = m(X_u, Y_u, inverse=True)
    X_v, Y_v = p.meshgrid(x,y_v)
    lon_v, lat_v = m(X_v, Y_v, inverse=True)
    if show_mass_grid:
        m.plot(X, Y, 'b+')
        m.plot([grid_centre_x], [grid_centre_y], 'r+')
        if show_stag_grids:
            m.plot(X_u, Y_u, 'g+')
            m.plot(X_v, Y_v, 'r+')
        m.drawcoastlines()
        p.show()
    output = {
      'map' : m,
      'mass_stag': {
        'lon_2d' : lon,
        'lat_2d' : lat,
        'x'      : x,
        'y'      : y,
        'x_2d'   : X,
        'y_2d'   : Y,
        }, 
      'u_stag': {
        'lon_2d' : lon_u,
        'lat_2d' : lat_u,
        'x'      : x_u,
        'y'      : y,
        'x_2d'   : X_u,
        'y_2d'   : Y_u,
        }, 
      'v_stag': {
        'lon_2d' : lon_v,
        'lat_2d' : lat_v,
        'x'      : x,
        'y'      : y_v,
        'x_2d'   : X_v,
        'y_2d'   : Y_v,
        } 
      }

    return output
      

def find_parent_ij(outer_grid, inner_grid):
    projection = outer_grid[0]
    lat_0 = outer_grid[1]
    lon_0 = outer_grid[2]
    lat_1 = outer_grid[3]
    lat_2 = outer_grid[4]
    grid_centre_lon = outer_grid[5]
    grid_centre_lat = outer_grid[6]
    nx = outer_grid[7]
    ny = outer_grid[8]
    delta = outer_grid[9]

    from pyproj import Proj
    proj = Proj(proj=projection, lat_0=lat_0, lon_0=lon_0, lat_1=lat_1, lat_2=lat_2)
    grid_centre_x, grid_centre_y = proj(grid_centre_lon, grid_centre_lat)
    grid_x_extent = nx * delta
    grid_y_extent = ny * delta
    min_x = grid_centre_x - grid_x_extent/2.
    min_y = grid_centre_y - grid_y_extent/2.
    max_x = min_x + grid_x_extent
    max_y = min_y + grid_y_extent
    outer_x = n.arange(min_x, max_x + a_small_number, delta)
    outer_y = n.arange(min_y, max_y + a_small_number, delta)  

    projection = inner_grid[0]
    lat_0 = inner_grid[1]
    lon_0 = inner_grid[2]
    lat_1 = inner_grid[3]
    lat_2 = inner_grid[4]
    grid_centre_lon = inner_grid[5]
    grid_centre_lat = inner_grid[6]
    nx = inner_grid[7]
    ny = inner_grid[8]
    delta = inner_grid[9]
 
    grid_centre_x, grid_centre_y = proj(grid_centre_lon, grid_centre_lat)
    grid_x_extent = nx * delta
    grid_y_extent = ny * delta
    min_x = grid_centre_x - grid_x_extent/2.
    min_y = grid_centre_y - grid_y_extent/2.
    max_x = min_x + grid_x_extent
    max_y = min_y + grid_y_extent
    inner_x = n.arange(min_x, max_x + a_small_number, delta)
    inner_y = n.arange(min_y, max_y + a_small_number, delta)

    return outer_x.searchsorted(inner_x[0]), outer_y.searchsorted(inner_y[0])

def write_namelist(namelist_dict,outfile='outfile'):

    out_string=''
    for group in namelist_dict.keys():
        out_string += group + '\n'
        for variable in namelist_dict[group].keys():
            out_string += variable + ' = ' 
            for element in namelist_dict[group][variable]:
                out_string += repr(element).strip("'")+', '
            out_string+='\n'
        out_string+= '/\n\n'

    fid=open(outfile,'w')
    fid.write(out_string)

    return None



def read_namelist(namelist_file):
    """read contents of namelist file and return dictionary containing all options
    
    Created 20/01/08 by Thom Chubb.
    Modified 20/01/08 by Thom Chubb and Valerio Bisignesi

    TODO: mod_levs have a slightly different format in the namelist file, but as 
    they come last in namelist.wps I have conveniently dropped them (with a warning
    of course =) ). Whoever needs them first can come up with a fix for this.
    Untested as yet with the namelist.input file. It should work fine and may be useful 
    as a consistency check between the two files. This has been buggine me for a while.
    """

    fid=open(namelist_file)

    out_dict={}
    data = fid.readlines()
    num_lines = len(data)

    for line in data:
        if '&' in line:
            # Then this line is a namelist title
            is_comment=False
            current_label = line.rstrip('\n').lstrip(' ')
            out_dict[current_label] ={}
        elif '/' in line:
            # Then lines following this are comments until the next '&'
            is_comment=True
        elif '=' in line:
            # Then this line contains variable information to be stored
            if not is_comment:
                variable,values = line.split('=')
                values = values.rstrip('\n').rstrip(',')
                try:
                    values=[int(element) for element in values.split(',')]
                except ValueError:
                    try:
                        values=[float(element) for element in values.split(',')]
                    except ValueError:
                        values=[value.strip() for value in values.split(',')]

                out_dict[current_label][variable.strip()]=values

    return out_dict

def check_namelist_consistency():
    # TODO
        # run time vs. run date 
        # strict consistency between dates in namelist.wps and namelist.input not 
        # necessary as long as dates in namelist.input are a subset of those in namelist.wps
        # and the interval_seconds is correct
    pass

#def read_namelist_old(namelist_file):
#    """read contents of namelist file and return dictionary containing all options
#    
#    Created 20/01/08 by Thom Chubb.
#
#    TODO: mod_levs have a slightly different format in the namelist file, but as 
#    they come last in namelist.wps I have conveniently dropped them (with a warning
#    of course =) ). Whoever needs them first can come up with a fix for this.
#    Untested as yet with the namelist.input file. It should work fine and may be useful 
#    as a consistency check between the two files. This has been buggine me for a while.
#    """
#
#    fid=open(namelist_file)
#
#    out_dict={}
#    data = fid.readlines()
#    num_lines = len(data)
#
#    # for k in range(0,num_lines):
#    for line in data:
#       # str = data[k].rstrip('\n').rstrip(',').split()
#       str = line.rstrip('\n').rstrip(',').split()
#
#       if str == []:
#           pass
#       elif str[0] == '':  
#           pass
#       elif str[0][0] == '':
#           pass
#       elif str[0][0] == '/' :
#           is_comment=True
#       elif str[0][0] == '&':
#           # Then this line is a namelist title
#           is_comment=False
#           label = str[0]
#
#           if label == '&mod_levs':
#               print ">> WARNING: mod levels don't work yet"
#               break
#
#           out_dict[label] ={}
#
#       else: 
#           if not is_comment:
#               field = str[0]
#               out_dict[label][field] = [] 
#
#
#               for k in range(2,str.__len__()):
#                   dat = str[k].rstrip(',')
#                   # dat = str[k].split(',')
#                   print str, dat
#                   try:
#                       dat=float(dat)
#                   except ValueError:
#                       pass
#                   except TypeError:
#                       pass
#
#                   out_dict[label][field].extend(dat) 
#           
#           # out_dict[label][field] = [] 
#           # out_dict[label][field].append(str[2:])
#
#    return out_dict


def wrf_grid_wrapper(namelist_file='namelist.wps',nest_level=0):
    """
    wrf_grid_wrapper(namelist_file='namelist.wps',nest_level=0):
    Basic wrapper to easily visualise grids specified in namelist.wps
    
    Uses wrf.utils.read_namelist() to determine the read the appropriate variables 
    in a specified namelist file and then calls wrf.utils.wrf_grid() to define 
    the Basemap projection and show the grid over a map.

    Created 20/01/08 by Thom Chubb.
    Modified 27/01/08 - implemented viz.utils.plot_grid() to handle plotting and 
    capability for viewing as many grids as desired.

    TODO: Could use some more error checking, i think it will all fail if nest_levels
    are not consecutive!! Interpolation implemented is awkward but seems to work.
        
        """ 

    # Create namelist dictionary
    nd = read_namelist(namelist_file)

    # Field editing to make python happy
    if nd['&geogrid']['map_proj'][0]=="'lambert'":
        print 'debug: modify input field lambert -> lcc' 
        nd['&geogrid']['map_proj'][0]='lcc'
    
    grid = []

    outer_grid = wrf2latlon(nd['&geogrid']['map_proj'][0],
                    nd['&geogrid']['ref_lat'][0],
                    nd['&geogrid']['ref_lon'][0],
                    nd['&geogrid']['truelat1'][0], 
                    nd['&geogrid']['truelat2'][0],
                    nd['&geogrid']['ref_lon'][0],
                    nd['&geogrid']['ref_lat'][0],
#                   nd['&geogrid']['e_we'][nest_level[0]],
#                   nd['&geogrid']['e_sn'][nest_level[0]],
                    nd['&geogrid']['e_we'][0],
                    nd['&geogrid']['e_sn'][0],
                    nd['&geogrid']['dx'][0],
                    staggered = False,
                    return_extra = True
                    )
    # print "outer_grid.shape =", outer_grid[0].shape

    grid.append(outer_grid)
    nest_level.sort()
    # nest_level=p.sort(nest_level)

    # for k in nest_level[1:]:
    for k in range(1,max(nest_level)+1):
        this_grid = []
        try:
            e_we = nd['&geogrid']['e_we'][k]
        except IndexError:
            print "Out of range. Not enough grids specified within namelist file"
            return 1
        e_sn = nd['&geogrid']['e_sn'][k]
        pgr  = nd['&geogrid']['parent_grid_ratio'][k]
        ips  = nd['&geogrid']['i_parent_start'][k]
        jps  = nd['&geogrid']['j_parent_start'][k]
        print 'processing grid: ',k
        # print e_we,e_sn,pgr,ips,jps
        # print ips,':',(ips+(e_we/pgr)),',', jps,':',(jps+(e_sn/pgr))
        
        # Interpolate in grid space to estimate inner gridpoints - 
        # care to find a more elegant approach???
        x1=grid[-1][2][jps, ips:ips+e_we/pgr]
        y1=grid[-1][3][jps:jps+e_sn/pgr, ips]
        
        a1=n.arange(0,x1.__len__(),1) 
        a2=n.arange(0,x1.__len__(),1./pgr)
        b1=n.arange(0,y1.__len__(),1)
        b2=n.arange(0,y1.__len__(),1./pgr)

        x2=n.interp(a2,a1,x1)
        y2=n.interp(b2,b1,y1)

        [X,Y]=n.meshgrid(x2,y2)

        # convert back to navigational coordinates
        lon,lat=grid[-1][4](X,Y,nd['&geogrid']['map_proj'])

        for j in [lon,lat,X,Y,grid[-1][4]]:
            this_grid.append(j)
        if (k in nest_level):   
            map=vu.plot_grid(this_grid[0],this_grid[1],skip=10,same_figure=True,return_map=True) 
        grid.append(this_grid)  
        # print grid[-1][0].shape 


    if 0 in nest_level:
        map=vu.plot_grid(outer_grid[0],outer_grid[1],skip=10,same_figure=True,return_map=True) 
    map.drawmeridians(n.arange(130,180,15),labels=[1,0,0,1])
    map.drawparallels(n.arange(0,-90,-15),labels=[1,0,0,1])
    map.drawcoastlines()
    map.drawrivers()
    plot_custom_points(map)
    p.show()
    
    return grid, map

def calculate_mslp(p,pb,ph,phb,t,qvapor):
    '''
    calculate sea level pressure starting from 'raw' wrf output fields
    usage:
    >>> calculate_mslp(p,pb,ph,phb,t,qvapor)
    where the arguments names correspond to the variable names in the 
    wrfout files e.g. p(lvl,lat,lon) or p(time,lvl,lat,lon)
    '''
    import  from_wrf_to_grads as fw2g
    cs = fw2g.from_wrf_to_grads.compute_seaprs

    if len(p.shape) == 3:
       # recover the full pressure field by adding perturbation and base
       p = p + pb
       p_t = p.transpose()
       # same geopotential height
       ph = (ph + phb) / 9.81
       ph_t = ph.transpose()
       qvapor_t = qvapor.transpose()
       # do not add the wrf specified 300 factor as the wrapped fortran code
       # does that for us
       t_t = t.transpose()
       nz = ph_t.shape[2]
       # populate the geopotential_height at mid_levels array with
       # averages between layers below and above
       z = (ph_t[:,:,:nz-1] + ph_t[:,:,1:nz]) / 2.0
       # finally "in one fell sweep"
       # the zero is for debug purposes
       return cs(z,t_t,p_t,qvapor_t,0).transpose()
    elif len(p.shape) == 4:
       mslp_shape = (p.shape[0], p.shape[2], p.shape[3])
       mslp = n.zeros(mslp_shape)
       for time_idx in range(p.shape[0]):
           # recover the full pressure field by adding perturbation and base
           dummy_p = p[time_idx] + pb[time_idx]
           dummy_p_t = dummy_p.transpose()
           # same geopotential height
           dummy_ph = (ph[time_idx] + phb[time_idx]) / 9.81
           dummy_ph_t = dummy_ph.transpose()
           dummy_qvapor_t = qvapor[time_idx].transpose()
           # do not add the wrf specified 300 factor as the wrapped fortran code
           # does that for us
           dummy_t_t = t[time_idx].transpose()
           nz = dummy_ph_t.shape[2]
           # populate the geopotential_height at mid_levels array with
           # averages between layers below and above
           z = (dummy_ph_t[:,:,:nz-1] + dummy_ph_t[:,:,1:nz]) / 2.0
           # finally "in one fell sweep"
           # the zero is for debug purposes
           mslp[time_idx] = cs(z,dummy_t_t,dummy_p_t,dummy_qvapor_t,0).transpose()
       return mslp
    else:
       print 'Wrong shape of the array'
       return

def calculate_mslp_wrapper(vars_dict, time_idx):
    """Utility function to 
    pull out the necessary variables from the wrfout file and call
    calculate_mslp to generate the mslp field.
    """
    # accessing the times in the nc_file one at a time and
    # using the .copy() method reduce the memory footprint
    perturbation_pressure = vars_dict['P'].get_value()[time_idx].copy()
    base_pressure = vars_dict['PB'].get_value()[time_idx].copy()
    perturbation_geopotential = vars_dict['PH'].get_value()[time_idx].copy()
    base_geopotential = vars_dict['PHB'].get_value()[time_idx].copy()
    temperature = vars_dict['T'].get_value()[time_idx].copy()
    mixing_ratio = vars_dict['QVAPOR'].get_value()[time_idx].copy()
    mslp = calculate_mslp(
      perturbation_pressure, 
      base_pressure,
      perturbation_geopotential,
      base_geopotential,
      temperature,
      mixing_ratio)
    #del perturbation_pressure, base_pressure
    #del perturbation_geopotential, base_geopotential
    #del temperature, mixing_ratio
    return mslp

def plot_custom_points(map):
    """back by popular demand"""

#    canberra_lon = [149 + 8./60]
#    canberra_lat = [-35 - 17./60]
#    map.plot(canberra_lon,canberra_lat, 'gs')

    blue_calf_lon = [148.3944] 
    blue_calf_lat = [-36.3869]  
    map.plot(blue_calf_lon,blue_calf_lat, 'gs')
    return

def time_string_to_datetime(times):
    '''
    This function takes as input a numpy array of numpy string-arrays and
    returns a list of corresponding datetime objects.
    The array will be typically generated by a pynio get_value() call for 
    the 'Times' variable of a wrfout file.
    
    Usage:
    >>> import wrf.utils as wu
    >>> import viz.utils as vu
    >>> f,fv = vu.peek('some_wrfout_file' + '.nc', return_pointers=True)
    >>> times = fv['Times'].get_value()
    >>> times = wu.time_string_to_datetime(times)
    '''
    from datetime import datetime
    # VB we know that using pynio to access the 'Times' variable in a wrfout 
    # files we are returned a numpy array of string arrays hence
    result = []
    for time in times:
        time_string = time.tostring()
        year = int(time_string[:4])
        month = int(time_string[5:7])
        day = int(time_string[8:10])
        hour = int(time_string[11:13])
        minute = int(time_string[14:16])
        second = int(time_string[17:])
        # VB We assume the time is UTC so we will not bother 
        # specifying a time zone
        result.append(datetime(year,month,day,hour,minute,second))
    return result