get_grid_inp_file

#!/usr/bin/env python
# This script generates the x-coord, y-coord and bottom depth files
# for ww3_grid.inp
#
#  Qing Li, 20170305

import sys
import numpy as np
from netCDF4 import Dataset
from astropy.io import ascii

def main():
    """
    Read xxx.grid.nc, and save the x-coord, y-coord
        and bottom depth to txt files for ww3_grid.inp
    """
    # check input
    gridlist = ('gx37', 'gx37b', 'gx16', 'gx16b', 'gx16b_np', 'gx16b_sp', 'gx37b_eq')
    varargin = sys.argv
    narg = len(varargin)
    if narg > 2:
        raise ValueError('Require 1 argument, got {}.'.format(narg-1))
    gridtype = varargin[1]
    if gridtype not in gridlist:
        raise ValueError('Input grid type {} not supported'.format(gridtype))

    # read data
    gridpath = '.'
    inppath = './grid_inp'
    gridname = gridtype.split('_')[0]
    infile = gridpath+'/'+gridname+'.grids.nc'
    ncfile = Dataset(infile,'r')
    nlat = ncfile.dimensions['nlat'].size
    nlon = ncfile.dimensions['nlon'].size
    print('----')
    print('Grid: {}'.format(gridtype))
    print('Input file: {}'.format(infile))

    # save x-coord, y-coord, bottom
    if gridtype == 'gx16b_np':
        latcutoff = 45.0
    elif gridtype == 'gx16b_sp':
        latcutoff = -45.0
    elif gridtype == 'gx37b_eq':
        latcutoff = [-50.0, 50.0]
    else:
        latcutoff = None

    # get mask
    dmask = get_mask(ncfile, 'TLAT', latcutoff)

    # save data
    save_var(ncfile, 'TLONG', inppath+'/'+gridtype+'_x.inp', mask=dmask)
    save_var(ncfile, 'TLAT',  inppath+'/'+gridtype+'_y.inp', mask=dmask)
    save_var(ncfile, 'HT',    inppath+'/'+gridtype+'_bottom.inp', fscale=0.01, mask=dmask)
    save_var(ncfile, 'REGION_MASK', inppath+'/'+gridtype+'_mapsta.inp', fscale=0, mask=dmask)
    print('----')
    ncfile.close()

def get_mask(ncfile, latname, latcutoff):
    """
    Return a mask with True where lat in:
        (latcutoff, 90) if latcutoff > 0,
        (-90, latcutoff) if latcutoff < 0,
        (latcutoff[0], latcutoff[1]) if latcutoff has two elements,
        and False otherwise.
    """
    # read lat from file
    lat = ncfile.variables['TLAT'][:]
    ny = lat.shape[0]
    # set mininum and maximum latitde
    if latcutoff is None:
        return None
    elif np.size(latcutoff) == 1:
        if latcutoff > 0.0:
            minlat = latcutoff
            maxlat = 90.0
        else:
            minlat = -90.0
            maxlat = latcutoff
    elif np.size(latcutoff) == 2:
        minlat = latcutoff[0]
        maxlat = latcutoff[1]
    else:
        raise ValueError('latcutoff should be a scaler or a list with two elements')

    # check minlat an maxlat
    if minlat < -90.0 or maxlat > 90.0:
        raise ValueError('Invalid cutoff latitude')

    # generate mask
    dmask = np.ma.where((lat>minlat) & (lat<maxlat), True, False)

    for i in range(ny):
        if np.any(dmask[i,:]):
            dmask[i,:] = True

    return dmask

def save_var(ncfile, varname, outfile, fscale=1, mask=None):
    """
    Read in variable 'varname' from file 'ncfile' and write it in
        txt file 'outfile'
        fscale is optional: 0 - set non zero values in dat to 1
                            1 - do nothing
                            otherwise - set dat to dat * fscale
        mask is optional: None - no mask, global field
                          Otherwise, should be in the same dimension as 'varname'
    """
    nlat = ncfile.dimensions['nlat'].size
    nlon = ncfile.dimensions['nlon'].size
    dat = ncfile.variables[varname][:]

    # check dimension
    datasize = dat.shape
    if len(datasize) != 2:
        raise ValueError('Dimension of data should be 2')

    ndatx = dat.shape[1]
    ndaty = dat.shape[0]
    if ndatx != nlon or ndaty != nlat:
        raise ValueError('Dimensions not match')

    # scale data if applicable
    if fscale == 0:
        dat = np.where(dat != 0, 1, 0)
    elif fscale != 1:
        dat = dat * fscale

    # apply mask if passed in
    if not mask is None:
        dat = np.where(mask, dat, None)

    # delete Nones
    yindices = []
    for i in range(ndaty):
        if np.all(np.equal(dat[i,:], None)):
            yindices.append(i)

    dat = np.delete(dat, yindices, 0)

    # write out data
    write_inp(dat, outfile)

def write_inp(dat, outfile):
    """
    Write out data in a txt file.
        dat: 2 dimensional array of data
        outfile: output filename
    """
    print('Output file: {}, (nx, ny) = ({}, {})'.format(outfile, dat.shape[1], dat.shape[0]))
    ascii.write(dat, outfile, format='no_header', overwrite=True)

if __name__ == "__main__":
    main()