def proc_write(twx_cfg, ncdf_mode, start_ymd, end_ymd, nwrkers): status = MPI.Status() stn_da = StationDataDb(twx_cfg.fpath_stndata_nc_tair_homog, (start_ymd, end_ymd)) days = stn_da.days nwrkrs_done = 0 bcast_msg = None bcast_msg = MPI.COMM_WORLD.bcast(bcast_msg, root=RANK_COORD) stnids_tmin, stnids_tmax = bcast_msg print "WRITER: Received broadcast msg" if ncdf_mode == 'r+': ds_tmin = Dataset(twx_cfg.fpath_stndata_nc_infill_tmin, 'r+') ds_tmax = Dataset(twx_cfg.fpath_stndata_nc_infill_tmax, 'r+') ttl_infills = stnids_tmin.size + stnids_tmax.size stnids_tmin = ds_tmin.variables[STN_ID][:].astype(np.str) stnids_tmax = ds_tmax.variables[STN_ID][:].astype(np.str) else: stns_tmin = stn_da.stns[np.in1d(stn_da.stns[STN_ID], stnids_tmin, assume_unique=True)] variables_tmin = [('tmin', 'f4', netCDF4.default_fillvals['f4'], 'minimum air temperature', 'C'), ('flag_infilled', 'i1', netCDF4.default_fillvals['i1'], 'infilled flag', ''), ('tmin_infilled', 'f4', netCDF4.default_fillvals['f4'], 'infilled minimum air temperature', 'C')] create_quick_db(twx_cfg.fpath_stndata_nc_infill_tmin, stns_tmin, days, variables_tmin) stnda_out_tmin = StationDataDb(twx_cfg.fpath_stndata_nc_infill_tmin, mode="r+") stnda_out_tmin.add_stn_variable('mae', 'mean absolute error', 'C', "f8") stnda_out_tmin.add_stn_variable('bias', 'bias', 'C', "f8") ds_tmin = stnda_out_tmin.ds stns_tmax = stn_da.stns[np.in1d(stn_da.stns[STN_ID], stnids_tmax, assume_unique=True)] variables_tmax = [('tmax', 'f4', netCDF4.default_fillvals['f4'], 'maximum air temperature', 'C'), ('flag_infilled', 'i1', netCDF4.default_fillvals['i1'], 'infilled flag', ''), ('tmax_infilled', 'f4', netCDF4.default_fillvals['f4'], 'infilled maximum air temperature', 'C')] create_quick_db(twx_cfg.fpath_stndata_nc_infill_tmax, stns_tmax, days, variables_tmax) stnda_out_tmax = StationDataDb(twx_cfg.fpath_stndata_nc_infill_tmax, mode="r+") stnda_out_tmax.add_stn_variable('mae', 'mean absolute error', 'C', "f8") stnda_out_tmax.add_stn_variable('bias', 'bias', 'C', "f8") ds_tmax = stnda_out_tmax.ds ttl_infills = stnids_tmin.size + stnids_tmax.size print "WRITER: Infilling a total of %d station time series " % (ttl_infills,) print "WRITER: Output NCDF files ready" stat_chk = StatusCheck(ttl_infills, 10) while 1: result = MPI.COMM_WORLD.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status) stn_id, tair_var, tair, fill_mask, tair_infill, mae, bias = result if status.tag == TAG_STOPWORK: nwrkrs_done += 1 if nwrkrs_done == nwrkers: print "Writer: Finished" return 0 else: if tair_var == 'tmin': stn_idx = np.nonzero(stnids_tmin == stn_id)[0][0] ds = ds_tmin else: stn_idx = np.nonzero(stnids_tmax == stn_id)[0][0] ds = ds_tmax ds.variables[tair_var][:, stn_idx] = tair ds.variables["".join([tair_var, "_infilled"])][:, stn_idx] = tair_infill ds.variables['flag_infilled'][:, stn_idx] = fill_mask ds.variables['bias'][stn_idx] = bias ds.variables[LAST_VAR_WRITTEN][stn_idx] = mae ds.sync() print "|".join(["WRITER", stn_id, tair_var, "%.4f" % (mae,), "%.4f" % (bias,)]) stat_chk.increment()
def proc_write(twx_cfg, start_ymd, end_ymd, nwrkers): status = MPI.Status() nwrkrs_done = 0 stn_da = StationDataDb(twx_cfg.fpath_stndata_nc_tair_homog, (start_ymd, end_ymd), mode="r+") mths = np.arange(1, 13) for mth in mths: for varname in ['tmin', 'tmax']: varname_mean = get_mean_varname(varname, mth) varname_vari = get_variance_varname(varname, mth) stn_da.add_stn_variable(varname_mean, varname_mean, "C", 'f8') stn_da.add_stn_variable(varname_vari, varname_vari, "C**2", 'f8') stn_da.ds.sync() bcast_msg = None bcast_msg = MPI.COMM_WORLD.bcast(bcast_msg, root=RANK_COORD) mask_por_tmin, mask_por_tmax = bcast_msg stn_ids_tmin, stn_ids_tmax = (stn_da.stn_ids[mask_por_tmin], stn_da.stn_ids[mask_por_tmax]) print "WRITER: Received broadcast msg" stn_ids_uniq = np.unique(np.concatenate([stn_ids_tmin, stn_ids_tmax])) stn_idxs = {} for x in np.arange(stn_da.stn_ids.size): if stn_da.stn_ids[x] in stn_ids_uniq: stn_idxs[stn_da.stn_ids[x]] = x ttl_infills = stn_ids_tmin.size + stn_ids_tmax.size stat_chk = StatusCheck(ttl_infills, 30) while 1: stn_id, tair_var, stn_mean, stn_vari = MPI.COMM_WORLD.recv( source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status) if status.tag == TAG_STOPWORK: nwrkrs_done += 1 if nwrkrs_done == nwrkers: print "WRITER: Finished" return 0 else: stnid_dim = stn_idxs[stn_id] for mth in mths: vname_mean = get_mean_varname(tair_var, mth) stn_da.ds.variables[vname_mean][stnid_dim] = stn_mean[mth - 1] vname_vary = get_variance_varname(tair_var, mth) stn_da.ds.variables[vname_vary][stnid_dim] = stn_vari[mth - 1] stn_da.ds.sync() stat_chk.increment()