#!/usr/bin/python

# nolake.m

# Trevor Alcott - July 2013

# sample python script to modify WRF input files (met_em)
# python equivalent of "nolake.m", only showing LANDMASK changes here

# CHPC will need to install the netCDF4 module
# This is available via YUM in Redhat 5/6

import os
from netCDF4 import Dataset
from numpy import *

wpsdir='/uufs/chpc.utah.edu/common/home/steen_data2/trevor/research_wrf/WPS/'

for dirpath, dirnames, filenames in os.walk(wpsdir):
  for f in filenames:

    ncfile = os.path.join(wpsdir, f)

    # 'a' option is to append/edit, use 'r' if you just want to read
    rootgrp = Dataset(ncfile, 'a')

##  optional: list variable names
#    for var in rootgrp.variables:
#      print var

    lat = (rootgrp.variables['XLAT_M'])[0]
    lon = (rootgrp.variables['XLONG_M'])[0]
    landmask = rootgrp.variables['LANDMASK']
    landmaskvals = landmask[0]

    gsl = where(logical_and(logical_and(logical_and(logical_and(greater(lat,40.6),less(lat,41.7)),greater(lon,-113)),less(lon,-111.9)),equal(landmask,0)),1,0)

    landmaskvals=where(gsl,1,landmaskvals)
    landmask[0] = landmaskvals

    rootgrp.close()