Source code for pycif.plugins.datastreams.fluxes.wrfchem.write

import netCDF4 as nc
import numpy as np
import os
from .times_in_wrf_file import times_in_wrf_file
import os


[docs] def write(self, name, flx_file, date, flx, domain, mode): """ Write flux data for one timestep to WRF-Chem flux input file Args: self (Fluxes): the Fluxes plugin name (str): Name of variable in flx file flx_file (str): the file where to write fluxes date (datetime.datetime): Timestamp of the flux flx (numpy.ndarray): fluxes data to write domain: domain plugin mode (str): 'w' to overwrite, 'a' to append (existing files) """ if flx.shape[0] != 1: raise ValueError("Can only write values for one time index at a time") # Reshape flx in case it comes on the unstructured wrf domain if getattr(domain, "unstructured_domain", False): flx = flx.reshape((flx.shape[0],) + (1,) + domain.zlon2d.shape) # Define emission variable dimensions em_dim_names = [ "Time", "emissions_zdim", "south_north", "west_east", ] em_dim_lens = [None] + list(flx.shape[1:]) # Define emission variable datatype em_dtype = "f4" # Define emission variable attributes em_attr_names = [ "FieldType", "MemoryOrder", "description", "units", "stagger", "coordinates" ] em_attr_vals = [ 104., "XYZ", name+" emission", "mole/km2/hr", "Z", "XLONG XLAT" ] # Cannot work with nested domains dom = 1 # mode w: write file from scratch if mode=="w": if os.path.exists(flx_file): raise ValueError("Output file exists") # Open output file for writing ncf = nc.Dataset(flx_file, "w") # Create dimensions # For emission variables: for n in range(len(em_dim_names)): ncf.createDimension(em_dim_names[n], em_dim_lens[n]) # For Times: ncf.createDimension("DateStrLen", 19) # Create variables times = ncf.createVariable("Times", "S1", ("Time", "DateStrLen")) var = ncf.createVariable(name, em_dtype, em_dim_names) # Write Variable attributes for n in range(len(em_attr_names)): setattr(var, em_attr_names[n], em_attr_vals[n]) # Fill variables with values date_str = date.strftime("%Y-%m-%d_%H:%M:%S") times[:] = [[c for c in date_str]] var[:] = flx # Get and write global attributes # Get these from wrfinput_d01: glob_attr_names = [ "TITLE", "CEN_LAT", "CEN_LON", "TRUELAT1", "TRUELAT2", "MOAD_CEN_LAT", "STAND_LON", "POLE_LAT", "POLE_LON", "MAP_PROJ", "DX", "DY", "MMINLU", "NUM_LAND_CAT", "ISWATER", "ISLAKE", "ISICE", "ISURBAN", "ISOILWATER" ] fp_wrfinput = os.path.join( domain.dir_wrfinput, "wrfinput_d{:02d}".format(dom)) wrfin = nc.Dataset(fp_wrfinput) glob_attr_vals = [getattr(wrfin, nm) for nm in glob_attr_names] # Change title so that wrf.exe wont complain about versions version = wrfin.TITLE[len(" OUTPUT FROM REAL_EM "):-len(" PREPROCESSOR")] title_out = "OUTPUT FROM PREPROCESSING_FOR " + version glob_attr_vals[0] = title_out wrfin.close() # Add date information glob_attr_names.append("GMT") glob_attr_vals.append(np.int32(1)) glob_attr_names.append("JULYR") glob_attr_vals.append(np.int32(date.year)) glob_attr_names.append("JULDAY") glob_attr_vals.append(np.int32(date.timetuple().tm_yday)) # Write attributes for n in range(len(glob_attr_names)): setattr(ncf, glob_attr_names[n], glob_attr_vals[n]) ncf.close() # mode a: Append to existing file elif mode=="a": # Check times times = times_in_wrf_file(flx_file) ind = [n for n, t in enumerate(times) if t==date] # Complain if requested date not found exactly once in file if len(ind)!=1: fmt = "%Y-%m-%d %H:%M:%S" date_str = date.strftime(fmt) times_str_l = [time.strftime(fmt) for time in times] times_str = ", ".join(times_str_l) if len(ind)==0: msg = "Date {} not found in file {}. Dates in file: {}." elif len(ind)>1: msg = ("Date {} found more than once in file {}. " + \ "Dates in file: {}.") raise ValueError(msg.format(date_str, flx_file, times_str)) # Write ncf = nc.Dataset(flx_file, mode) # If variable does not exist, create it if not name in ncf.variables: var = ncf.createVariable(name, em_dtype, em_dim_names) # Write Variable attributes for n in range(len(em_attr_names)): setattr(var, em_attr_names[n], em_attr_vals[n]) if not flx.shape[2:3] == ncf[name].shape[2:3]: msg = "Horizontal flux dimensions are {}, but in file {} they are {}" raise ValueError(msg.format(str(flx.shape[2:3]), flx_file, str(ncf[name].shape[2:3]))) if flx.shape[1] > ncf[name].shape[1]: msg = "Vertical flux dimension is {}, but in file {} it is {}, i.e. too small" raise ValueError(msg.format(str(flx.shape[1]), flx_file, str(ncf[name].shape[1]))) ncf[name][ind, ...] = 0.0 ncf[name][ind, :flx.shape[1], ...] = flx # Set missing values to 0 (else WRF might transport # the fill value) mask = ncf[name][ind, ...].mask if np.any(mask): ncf[name][ind, ...] = np.ma.filled(ncf[name][ind, ...], 0.0) ncf.close() else: raise ValueError("Unknown mode: " + mode)