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

# ---------------------------------------------------------------------------
# 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 = 'case_1' # Baseline TC Alby 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, ***'topog_flip.asc'***) # Name of the vertically flipped elevation grid
grid_file_names = glob.glob(os.path.join(event_folder, '*.gz'))
grid_file_names.sort()
number_of_timesteps = len(grid_file_names)
grid_size = 20 #0.0005 # cellsize from the GEMSURGE output grids
start_time = 0  # all times in seconds
end_time = 142560 #88200 # all times in seconds
timestep = 720 #1800 # all times in seconds
refzone = 50 # UTM zone
x_origin = ***#115.550 #364221.03   # xllcorner from GEMSURGE output grids
y_origin = ***#-32.790 #6371062.71  # yllcorner from GEMSURGE output grids
event_sts = os.path.join(event_folder, event)

assert number_of_timesteps * timestep == end_time - start_time

elevation = []
stage = []
speed = []
theta = []

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(float(L.strip().split()[1])) # this assumes an integer xllcorner
    # if i == 3:
        # y_origin = int(float(L.strip().split()[1])) # this assumes an integer yllcorner
    # if i == 4:
        # grid_size = int(float(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 'Grid size: ', grid_size
print 'X origin: ', x_origin
print 'Y origin: ', y_origin
  
for f in grid_file_names:
    print 'file: ', f
    
    gz_file = gzip.open(f, 'rb')
    lines = gz_file.readlines()
    gz_file.close()

    for i, L in enumerate(lines):
        
        if i > 7 and i < (8 + nrows):        # Number 7 refers to the number of rows of header info that we skip - CHECK format
            stage += [L.strip().split()]
            
        if i > (9 + nrows) and i < (10 + 2*nrows):
            speed += [L.strip().split()]
            
        if i > (11 + 2*nrows):
            theta += [L.strip().split()]

#------------------------------------------------------
# Create arrays of elevation, stage, speed and theta
#-------------------------------------------------------
print 'creating numpy arrays: elevation, stage, speed, theta'
            
elevation = numpy.array(elevation).astype('d')

stage_array = numpy.empty([number_of_timesteps*nrows, ncols], dtype=float)
speed_array = numpy.empty([number_of_timesteps*nrows, ncols], dtype=float)
theta_array = numpy.empty([number_of_timesteps*nrows, ncols], dtype=float)

for i in range(number_of_timesteps):
    ii = nrows*i
    stage_array[ii:ii+nrows, :] = numpy.flipud(numpy.array(stage[ii : ii + nrows]).astype('d'))
    speed_array[ii:ii+nrows, :] = numpy.flipud(numpy.array(speed[ii:ii + nrows]).astype('d'))
    theta_array[ii:ii+nrows, :] = numpy.flipud(numpy.array(theta[ii:ii + nrows]).astype('d'))

stage = stage_array.ravel()
speed = speed_array.ravel()
theta = theta_array.ravel()
    
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

for i in xrange(number_of_timesteps):
    j = number_of_points * i
    k = j + number_of_points
    depth[j:k] = stage[j:k] - elevation # Check GEMS elevations below sea level are negative

momentum = depth * speed

xmomentum = numpy.empty(momentum.size, dtype='d')
ymomentum = numpy.empty(momentum.size, dtype='d')

print 'Calculating xmomentum and ymomentum'

for i, t in enumerate(theta):

    if t > 0.0 and t < 90.0:
        mx = momentum[i]*math.sin(math.radians(t)) #Assuming t is in the "to" direction and 0 degrees is north
        my = momentum[i]*math.cos(math.radians(t))
        assert mx > 0 and my > 0
    if t > 90.0 and t < 180.0:
        mx = momentum[i]*math.sin(math.radians(180 - t))
        my = momentum[i]*math.cos(math.radians(180 - t))
        assert mx > 0 and my < 0
    if t > 180.0 and t < 270.0:
        mx = momentum[i]*math.cos(math.radians(270 - t))
        my = momentum[i]*math.sin(math.radians(270 - t))
        assert mx < 0 and my < 0
    if t > 270.0 and t < 360.0:
        mx = momentum[i]*math.sin(math.radians(360 - t))
        my = momentum[i]*math.cos(math.radians(360 - t))
        assert mx < 0 and my > 0
    if t == 0.0 or t == 360.0:
        mx = 0
        my = momentum[i]
        assert my > 0
    if t == 90.0:
        mx = momentum[i]
        my = 0
        assert mx > 0
    if t == 180.0:
        mx = 0
        my = momentum[i]
        assert my < 0
    if t == 270.0:
        mx = momentum[i]
        my = 0
        assert mx < 0
        
    xmomentum[i] = mx
    ymomentum[i] = my

    assert mx[i]^2 + my[i]^2 == momentum[i]
    
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"

fid = NetCDFFile(os.path.join(event_folder, event + '_master.sts'), 'w')
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

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

for i in xrange(number_of_timesteps):
    ii = i*number_of_points
    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()