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

# ---------------------------------------------------------------------------
# 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 = '20100527_gcom_12min'

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 = []
# this was to subsample the gz files
#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 = 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, depth, current_x and current_y
#-------------------------------------------------------
print 'creating numpy arrays: elevation, depth, current_x, current_y'

depth = numpy.empty(number_of_timesteps*nrows*ncols, dtype=float)
current_x = numpy.empty(number_of_timesteps*nrows*ncols, dtype=float)
current_y = 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')

    d = []
    cx = []
    cy = []

    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
        #    print 'first block'

        #if i > (9 + nrows) and i < (10 + 2*nrows):
        #    print 'second block' 

        #if i > (11 + 2*nrows):
        #    print 'third block'

        if i > 13 + 3*nrows and i < (14 + 4*nrows):   
            d += L.strip().split()
            
        if i > (15 + 4*nrows) and i < (16 + 5*nrows):
            cx += L.strip().split()
            
        if i > (17 + 5*nrows):
            cy += L.strip().split()

    depth[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(d).astype('d'))
    current_x[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(cx).astype('d'))
    current_y[j*nrows*ncols : (j+1)*nrows*ncols] = numpy.flipud(numpy.array(cy).astype('d'))

    gz_file.close()

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

print 'Size of elevation array: ', elevation.size
print 'Size of depth array: ', depth.size
print 'Size of current x array: ', current_x.size
print 'Size of current y array: ', current_y.size

assert depth.size == current_x.size == current_y.size == ncols * nrows * number_of_timesteps
assert depth.size == number_of_timesteps * elevation.size

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

# ---------------------------------------------
# Create mask of no_data values across depth, current_x and current_y 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
# ---------------------------------------------

no_value_index = numpy.where(((depth < -9000) + (current_x < -9000) + (current_y < -9000)) == True)[0]

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

print 'Calculating stage, xmomentum and ymomentum'

#stage = depth - numpy.tile(numpy.absolute(elevation), number_of_timesteps)
stage = depth + numpy.tile(elevation, number_of_timesteps)
xmomentum = current_x*depth #momentum*numpy.sin(numpy.radians(current_y))
ymomentum = current_y*depth #momentum*numpy.cos(numpy.radians(current_y))

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

#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 xmomentum.size == depth.size == ymomentum.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()