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)