source: branches/anuga_1_1/anuga_work/production/bunbury_storm_surge_2009/process_gems_grids.py @ 7723

# ---------------------------------------------------------------------------
# This python script reads in GEMSURGE outputs and writes them into a master 
# STS file that contains all the data across the entire model domain at all 
# timesteps. It requires a vertically flipped elevation ASCII grid which is 
# merged with the GEMSURGE data to calculate the required quantities of 
# xmomentum and ymomentum.
# Written by Nariman Habili and Leharne Fountain, 2010
# ---------------------------------------------------------------------------

import glob
import os
import gzip
import pdb 
import numpy
import math
from Scientific.IO.NetCDF import NetCDFFile
from anuga.coordinate_transforms.geo_reference import Geo_reference, write_NetCDF_georeference

#-----------------
# Directory setup
#-------------------
state = 'western_australia'
scenario_folder = 'bunbury_storm_surge_scenario_2009'
event = 'gcom_60min'

print "Processing event: ", event

ENV_INUNDATIONHOME = 'INUNDATIONHOME'
home = os.path.join(os.getenv(ENV_INUNDATIONHOME), 'data') # Absolute path for data folder
# GEMS_folder = os.path.join(home, state, scenario_folder, 'GEMS')
anuga_folder = os.path.join(home, state, scenario_folder, 'anuga')
boundaries_folder = os.path.join(anuga_folder, 'boundaries')
event_folder = os.path.join(boundaries_folder, event)

# input_dir = os.path.join(GEMS_folder, event)

#-----------------------
# Input files from GEMS
#-----------------------
elevation_file_name = os.path.join(event_folder, 'buntopog_20m_flip.asc')  # Name of the vertically flipped elevation grid

print "Elevation file name: ", elevation_file_name

grid_file_names = glob.glob(os.path.join(event_folder, 'gcom*.gz'))
grid_file_names.sort()

grid_file_names_temp = []
for i in range(1, 32, 1):
    grid_file_names_temp.append(grid_file_names[i])

grid_file_names = grid_file_names_temp
    
number_of_timesteps = len(grid_file_names)

print "Number of timesteps: ", number_of_timesteps

start_time = 0      # all times in seconds
timestep = 3600 #1440 #720     # all times in seconds
end_time =  start_time + (timestep*number_of_timesteps) # all times in seconds
refzone = 50        # UTM zone of model
event_sts = os.path.join(event_folder, event)

elevation = []

elevation_file = open(elevation_file_name, 'rb')
lines = elevation_file.readlines()
elevation_file.close()

# Strip off the ASCII header and also read ncols, nrows, 
# x_origin, y_origin, grid_size, no_data value, and read in elevation grid:
for i, L in enumerate(lines):
    if i == 0:
        ncols = int(L.strip().split()[1])
    if i == 1:
        nrows = int(L.strip().split()[1])
    if i == 2:
        x_origin = int(L.strip().split()[1])
    if i == 3:
        y_origin = int(L.strip().split()[1])
    if i == 4:
        grid_size = int(L.strip().split()[1])
    if i == 5:
        no_data = int(float(L.strip().split()[1]))
    if i > 5:
        elevation+= L.strip().split()

print 'Number or columns: ', ncols
print 'Number of rows: ', nrows
print 'X origin: ', x_origin
print 'Y origin: ', y_origin
print 'Grid size: ', grid_size
print 'No data value: ', no_data

#------------------------------------------------------
# Create arrays of elevation, stage, speed and theta
#-------------------------------------------------------
print 'creating numpy arrays: elevation, stage, speed, theta'

stage = numpy.empty(number_of_timesteps*nrows*ncols, dtype=float)
speed = numpy.empty(number_of_timesteps*nrows*ncols, dtype=float)
theta = numpy.empty(number_of_timesteps*nrows*ncols, dtype=float)

for j, f in enumerate(grid_file_names):
    print 'file: ', f
    
    gz_file = gzip.open(f, 'rb')

    st = []
    sp = []
    t = []

    for i, L in enumerate(gz_file):
        
        if i > 7 and i < (8 + nrows):   #  7 refers to the number of rows of header info that we skip - always CHECK format
            st += L.strip().split()
            
        if i > (9 + nrows) and i < (10 + 2*nrows):
            sp += L.strip().split()
            
        if i > (11 + 2*nrows):
            t += L.strip().split()

    stage[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(st).astype('d'))
    speed[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(sp).astype('d'))
    theta[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(t).astype('d'))

    gz_file.close()

elevation = numpy.array(elevation).astype('d')

print 'Size of elevation array: ', elevation.size
print 'Size of stage array: ', stage.size
print 'Size of speed array: ', speed.size
print 'Size of theta array: ', theta.size

assert stage.size == speed.size == theta.size == ncols * nrows * number_of_timesteps
assert stage.size == number_of_timesteps * elevation.size

depth = numpy.empty(stage.size, dtype='d')
number_of_points = ncols * nrows

# ---------------------------------------------
# Create mask of no_data values across stage, speed and theta to ensure all three quantities
# have no_data values in corresponding cells - this is a work-around until GEMS can produce a 
# a consistant dataset
# ---------------------------------------------
import pdb
pdb.set_trace()
no_value_index = numpy.where(((stage < -100) + (speed < 0) + (theta < 0)) == True)[0]
depth = stage - numpy.tile(elevation, number_of_timesteps)

numpy.put(stage, no_value_index, -9999)
numpy.put(speed, no_value_index, -9999)
numpy.put(theta, no_value_index, -9999)
numpy.put(depth, no_value_index, 0)
    
# Taking absolute value is to account for -ve depths obtained when stage-elevation 
# is slightly -ve  - why I don't know, possbly a rounding error?
momentum = numpy.absolute(depth * speed)

print 'Calculating xmomentum and ymomentum'

xmomentum = momentum*numpy.sin(numpy.radians(theta))
ymomentum = momentum*numpy.cos(numpy.radians(theta))

numpy.put(xmomentum, no_value_index, 0)
numpy.put(ymomentum, no_value_index, 0)

#    if t > 0.0 and t < 90.0:
#        assert mx >= 0 and my >= 0
        
#    elif t > 90.0 and t < 180.0:
#        assert mx >= 0 and my <= 0
        
#    elif t > 180.0 and t < 270.0:
#        assert mx <= 0  and my <= 0
        
#    elif t > 270.0 and t < 360.0:
#        assert mx <= 0 and my >= 0
        
#    elif t == 0.0 or t == 360.0:
#        assert my >= 0
        
#    elif t == 90.0:
#        assert mx >= 0
        
#    elif t == 180.0:
#        assert my <= 0
        
#    elif t == 270.0:
#        assert mx <= 0
        
#    elif t == -9999:
#        mx = 0
#        my = 0
#    else:
#        print "Unexpected value of theta"
#        exit()
        
#    xmomentum[i] = mx
#    ymomentum[i] = my

assert len(numpy.where((numpy.sqrt(xmomentum**2 + ymomentum**2) - momentum) > 1.e-06)[0]) == 0
    
x_origin_int = int(10000*x_origin)
y_origin_int = int(10000*y_origin)
grid_size_int = int(10000*grid_size)

x = numpy.tile(numpy.arange(x_origin_int, (x_origin_int + ncols * grid_size_int), grid_size_int)/10000.0, nrows)
y = numpy.repeat(numpy.arange(y_origin_int, (y_origin_int + nrows * grid_size_int), grid_size_int)/10000.0, ncols)

assert x.size == y.size == number_of_points

time = numpy.arange(start_time, end_time, timestep, dtype='i')

assert time.size == number_of_timesteps
assert momentum.size == stage.size

# -----------------------------
# Create the STS file
# -----------------------------
print "Creating the STS NetCDF file"

#for j in range(2):

fid = NetCDFFile(os.path.join(event_folder, event + '_master_2_1.sts'), 'wl')
fid.institution = 'Geoscience Australia'
fid.description = "description"
fid.starttime = 0.0
fid.ncols = ncols
fid.nrows = nrows
fid.grid_size = grid_size
fid.no_data = no_data
fid.createDimension('number_of_points', number_of_points)
fid.createDimension('number_of_timesteps', number_of_timesteps)
fid.createDimension('numbers_in_range', 2)

fid.createVariable('x', 'd', ('number_of_points',))
fid.createVariable('y', 'd', ('number_of_points',))
fid.createVariable('elevation', 'd', ('number_of_points',))
fid.createVariable('elevation_range', 'd', ('numbers_in_range',))
fid.createVariable('time', 'i', ('number_of_timesteps',))
fid.createVariable('stage', 'd', ('number_of_timesteps', 'number_of_points'))
fid.createVariable('stage_range', 'd', ('numbers_in_range', ))
fid.createVariable('xmomentum', 'd', ('number_of_timesteps', 'number_of_points'))
fid.createVariable('xmomentum_range', 'd', ('numbers_in_range',))
fid.createVariable('ymomentum', 'd', ('number_of_timesteps', 'number_of_points'))
fid.createVariable('ymomentum_range', 'd', ('numbers_in_range',))

fid.variables['elevation_range'][:] = numpy.array([1e+036, -1e+036])
fid.variables['stage_range'][:] = numpy.array([1e+036, -1e+036])
fid.variables['xmomentum_range'][:] = numpy.array([1e+036, -1e+036])
fid.variables['ymomentum_range'][:] = numpy.array([1e+036, -1e+036])
fid.variables['elevation'][:] = elevation
fid.variables['time'][:] = time#[j*60 : j*60 + 60]

s = fid.variables['stage']
xm = fid.variables['xmomentum']
ym = fid.variables['ymomentum']

#jj = j*number_of_points*(number_of_timesteps/2)

for i in xrange(number_of_timesteps):
    ii = i*number_of_points# + jj
    s[i] = stage[ii : ii + number_of_points]
    xm[i] = xmomentum[ii : ii + number_of_points]
    ym[i] = ymomentum[ii : ii + number_of_points]

origin = Geo_reference(refzone, min(x), min(y)) # Check this section for inputs in eastings and northings - it works for long and lat
geo_ref = write_NetCDF_georeference(origin, fid)

fid.variables['x'][:] = x - geo_ref.get_xllcorner()
fid.variables['y'][:] = y - geo_ref.get_yllcorner()

fid.close()