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: 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', 'h':'static_nodes'} varlist = {} for this_var in var_list_raw: varlist[this_var] = varmatch[this_var] 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) else: fvcom_file = None lower_left_ll = None upper_right_ll = None grid_res = None depth_layers = None varlist = None global_time_indices = None # 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) time_indices = comm.scatter(global_time_indices, root=0) interped_data = {} worker = pf.interpolate.MPIRegularInterpolateWorker(fvcom_file, time_indices, comm=comm, verbose=True) 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: 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)) # 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), 'lat':len(worker.regular_grid.lats), 'DateStrLen': 26} 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'}, 'h':{'standard_name':'h', 'units':'m', 'long_name':'model bathymetry depth, measured downwards from geoid', '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'}} globals = {'type': 'OPERATIONAL MODEL REGULAR GRID OUTPUT', 'title': 'Regularly gridded data interpolated from FVCOM output', 'history': 'File created using PyFVCOM', 'filename': str(output_file_name), 'Conventions': 'CF-1.0', '_FillValue':-32768} for this_val in collected_interp_data.values(): np.ma.set_fill_value(this_val, -32768) ncopts = {} #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: # Add the variables. outfile.write_fvcom_time(fvcom.time.datetime) 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']) for this_var, this_mode in varlist.items(): if this_mode == 'surface': outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat'], attributes=all_attributes[this_var], ncopts=ncopts) elif this_mode == 'static_nodes': outfile.add_variable(this_var, collected_interp_data[this_var], ['lon', 'lat'], attributes=all_attributes[this_var], ncopts=ncopts) else: outfile.add_variable(this_var, collected_interp_data[this_var], ['time', 'lon', 'lat', 'depth'], attributes=all_attributes[this_var], ncopts=ncopts) else: print('Rank {} process finished'.format(rank))