grid_interp.py 5.13 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
from mpi4py import MPI
import numpy as np
import sys

import PyFVCOM as pf

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

if rank == 0:
12 13 14 15 16 17 18 19 20 21 22
    fvcom_file = sys.argv[1]
    lower_left_ll = np.asarray(sys.argv[2].split(','), dtype=float)
    upper_right_ll = np.asarray(sys.argv[3].split(','), dtype=float)
    grid_res = float(sys.argv[4])
    depth_layers = np.asarray(sys.argv[5].split(','), dtype=float)
    var_list_raw = sys.argv[6].split(',')
    varmatch = {'temp':'nodes', 'salinity':'nodes', 'u':'elements', 'v':'elements', 'zeta':'surface'}
    varlist = {}
    for this_var in var_list_raw:
        varlist[this_var] = varmatch[this_var]

23 24 25 26
    fvcom = pf.read.FileReader(fvcom_file)
    nt = fvcom.dims.time
    start, stop = 0, nt
    global_time_indices = np.array_split(np.arange(start, stop), size)
27

28
else:
29 30 31 32 33 34
    fvcom_file = None
    lower_left_ll = None
    upper_right_ll = None
    grid_res = None
    depth_layers = None
    varlist = None
35
    global_time_indices = None
36 37 38 39 40 41 42 43
# get the fvcom points within the extended grid and land mask the grid mesh

fvcom_file = comm.bcast(fvcom_file, root=0)
lower_left_ll = comm.bcast(lower_left_ll, root=0)
upper_right_ll = comm.bcast(upper_right_ll, root=0)
grid_res = comm.bcast(grid_res, root=0)
depth_layers = comm.bcast(depth_layers, root=0)
varlist = comm.bcast(varlist, root=0)
44 45 46 47 48

time_indices = comm.scatter(global_time_indices, root=0)

interped_data = {}

49
worker = pf.interpolate.MPIRegularInterpolateWorker(fvcom_file, time_indices, comm=comm, verbose=True)
50 51 52 53 54 55 56 57 58
worker.InitialiseGrid(lower_left_ll, upper_right_ll, grid_res, depth_layers, time_varying_depth=True)

for var, var_mode in varlist.items(): 
    interped_data[var] = worker.InterpolateRegular(var, mode=var_mode)

all_interped_data = comm.gather(interped_data, root=0)

# resort the data by timestep
if rank == 0:
59 60 61 62 63 64 65
    collected_interp_data = {}

    for this_var in varlist.keys():
        output_list = []
        for this_dict in all_interped_data:
            output_list.append(this_dict[this_var])
        collected_interp_data[this_var] = np.ma.masked_invalid(np.vstack(output_list))
66 67 68 69 70

    # write to cmems format
    fvcom = pf.read.FileReader(fvcom_file)

    dims = {'time':len(fvcom.time.datetime), 'depth':len(worker.regular_grid.dep_lays), 'lon':len(worker.regular_grid.lons),
71
                                'lat':len(worker.regular_grid.lats), 'DateStrLen': 26} 
72

73 74 75 76 77 78 79 80 81 82
    output_file_name = 'output.nc'

    all_attributes = {'lon':{'standard_name':'longitude', 'units':'degrees_east', 'long_name':'longitude'},
            'lat':{'standard_name':'latitude', 'units':'degrees_north', 'long_name':'latitude'},
            'depth':{'standard_name':'depth', 'units':'m', 'long_name':'depth, measured downwards from free surface', 'axis':'Z', 'positive':'down'},
            'temp':{'standard_name':'sea_water_potential_temperature','units':'C','missing_value':-32768, 'long_name':'Sea Water Potential Temperature'} ,
            'salinity':{'standard_name':'sea_water_salinity','units':'psu', 'missing_value':-32768, 'long_name':'Sea Water Salinity'},
            'u':{'standard_name':'eastward_sea_water_velocity','units':'m s-1', 'missing_value':-32768, 'long_name':'Eastward Current Velocity'},
            'v':{'standard_name':'northward_sea_water_velocity','units':'m s-1', 'missing_value':-32768, 'long_name':'Northward Current Velocity'},
            'zeta':{'standard_name':'sea_surface_height_above_geoid','units':'m', 'missing_value':-32768,'long_name':'Sea surface height above geoid'}}
83 84

    globals = {'type': 'OPERATIONAL MODEL REGULAR GRID OUTPUT',
85 86 87 88 89
                    'title': 'Regularly gridded data interpolated from FVCOM output',
                    'history': 'File created using PyFVCOM',
                    'filename': str(output_file_name),
                    'Conventions': 'CF-1.0',
                    '_FillValue':-32768}
90

91 92 93 94
    for this_val in collected_interp_data.values():
        np.ma.set_fill_value(this_val, -32768)
    ncopts = {}

95 96 97 98
    #for this_var, this_mode in varlist.items():
    #    this_output_file_name = 'output_{}.nc'.format(this_var)
    
    with pf.preproc.WriteForcing(output_file_name, dims, global_attributes=globals, clobber=True, format='NETCDF4') as outfile:
99
        # Add the variables.
100
        outfile.write_fvcom_time(fvcom.time.datetime)
101 102 103
        outfile.add_variable('lon', worker.regular_grid.lons, ['lon'], attributes=all_attributes['lon'])
        outfile.add_variable('lat', worker.regular_grid.lats, ['lat'], attributes=all_attributes['lat'])
        outfile.add_variable('depth', worker.regular_grid.dep_lays, ['depth'], attributes=all_attributes['depth'])
104
    
105
        for this_var, this_mode in varlist.items():
106 107
            if this_mode == 'surface':
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat'],
108
                                         attributes=all_attributes[this_var], ncopts=ncopts)
109 110
            else:
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat', 'depth'],
111 112
                                         attributes=all_attributes[this_var], ncopts=ncopts)
            
113 114
else:
    print('Rank {} process finished'.format(rank))