grid_interp.py 6.1 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
    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(',')
18
    varmatch = {'temp':'nodes', 'salinity':'nodes', 'u':'elements', 'v':'elements', 'zeta':'surface'}
19 20 21 22
    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
    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])
65

66
        collected_interp_data[this_var] = np.ma.masked_invalid(np.vstack(output_list))
67

68 69 70 71 72 73 74 75 76 77 78 79 80 81
    worker.fvcom_grid.select_ll = worker.fvcom_grid.ll
    h = worker._Interpolater(worker.fvcom_grid.h)
    grid_in_domain = pf.read.FileReader(worker.fvcom_file).in_domain(worker.regular_grid.mesh_lons, worker.regular_grid.mesh_lats)
    h[~grid_in_domain] = np.nan
    h = h.T

    land_mask = np.zeros([h.shape[0], h.shape[1], len(depth_layers)])

    for i, this_depth in enumerate(depth_layers):
        if this_depth == 0:
            land_mask[~np.isnan(h)] = 1
        else:
            land_mask[:,:,i][h > this_depth] = 1

82 83 84 85
    # 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),
86
                                'lat':len(worker.regular_grid.lats), 'DateStrLen': 26} 
87

88 89 90 91 92
    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'},
Modellers Operational's avatar
Modellers Operational committed
93
            'h':{'standard_name':'h', 'units':'m', 'long_name':'model bathymetry depth, measured downwards from geoid', 'axis':'Z', 'positive':'down'},
94 95 96 97
            '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'},
98 99
            'zeta':{'standard_name':'sea_surface_height_above_geoid','units':'m', 'missing_value':-32768,'long_name':'Sea surface height above geoid'},
            'mask':{'standard_name':'sea_binary_mask','units':'1', 'long_name':'Land-sea mask: sea = 1 ; land = 0'}}
100 101

    globals = {'type': 'OPERATIONAL MODEL REGULAR GRID OUTPUT',
102 103 104 105 106
                    'title': 'Regularly gridded data interpolated from FVCOM output',
                    'history': 'File created using PyFVCOM',
                    'filename': str(output_file_name),
                    'Conventions': 'CF-1.0',
                    '_FillValue':-32768}
107

108 109 110 111
    for this_val in collected_interp_data.values():
        np.ma.set_fill_value(this_val, -32768)
    ncopts = {}

112 113 114 115
    #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:
116
        # Add the variables.
117
        outfile.write_fvcom_time(fvcom.time.datetime)
118 119 120
        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'])
121 122 123
        outfile.add_variable('mask', land_mask, ['lon', 'lat', 'depth'], attributes=all_attributes['mask'])
        outfile.add_variable('h', h, ['lon', 'lat'], attributes=all_attributes['h'])

124
        for this_var, this_mode in varlist.items():
125 126
            if this_mode == 'surface':
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat'],
127
                                         attributes=all_attributes[this_var], ncopts=ncopts)
128 129
            else:
                outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat', 'depth'],
130 131
                                         attributes=all_attributes[this_var], ncopts=ncopts)
            
132 133
else:
    print('Rank {} process finished'.format(rank))