import os
import netCDF4 as nc
import numpy as np
import datetime as dtm
from .....utils.check.errclass import CifAttributeError, CifError, CifValueError
[docs]
def write(
self, varname, fileout, icbc, comp_type=None,
metadata=None,
**kwargs
):
"""Write initial or lateral boundary condition values into a WRF-Chem
file.
Refuses to write if ``fileout`` is a symlink, to avoid overwriting
an original reference file. For ``comp_type == "inicond"``, reshapes
``icbc`` and writes it directly into the ``varname`` (or
``self.forced_varname`` if set and ``varname`` is not found) NetCDF
variable. For ``comp_type == "latcond"``, matches the timestamps
carried by ``icbc`` to those already in the file, redistributes the
perimeter data in ``icbc`` back into the four boundary-side arrays
(``_BXS``/``_BXE``/``_BYS``/``_BYE``), computes the corresponding
trend arrays (``_BTxx``) from the time differences, and writes both
the boundary and trend arrays back into the file.
Args:
self: the wrfchem_icbc plugin
varname: name of the species variable to write
fileout: path of the WRF-Chem file to write into (must not be a
symlink)
icbc: xarray.DataArray with dims ``(time, lev, lat, lon)`` (for
``latcond``, ``lon`` spans the domain perimeter) to write
comp_type: ``'inicond'`` or ``'latcond'``; defaults to
``self.comp_type`` if not given
metadata: mapping possibly containing ``"domain"``, used as a
fallback if ``self.domain`` is not set (``latcond`` only)
Raises:
CifValueError: if ``fileout`` is a symlink, or if the requested
timestamps cannot be matched to those in the file (or if the
per-timestep time differences are not constant).
CifError: if ``varname`` cannot be found in the ``inicond`` file
and ``self.forced_varname`` is not set.
CifAttributeError: if neither ``self.domain`` nor ``metadata``
is available for ``latcond``.
"""
# Make sure that fileout is not a link (dont want to
# overwrite data in an original file)
if os.path.islink(fileout):
raise CifValueError("File " + fileout + " is a link, " +
"should not be for writing.")
# Check comp_type
if comp_type is None:
comp_type = self.comp_type
# Write initial conditions
if comp_type == "inicond":
ncf = nc.Dataset(fileout, "r+")
try:
var = ncf[varname]
except IndexError:
if not hasattr(self, "forced_varname"):
raise CifError(
f"Trying to dump initial conditions for species {varname} "
f"in {fileout} but could not find the variable. "
f"You can force the name by setting the parameter "
f"`forced_varname`"
)
var = ncf[self.forced_varname]
tmp = np.reshape(icbc, var.shape)
var[:] = tmp
ncf.close()
elif comp_type == "latcond":
# Loading longitudes and latitudes
try:
domain = self.domain
except AttributeError:
if metadata is None:
raise CifAttributeError(
"Could not find any information about the domain."
)
domain = metadata["domain"]
# Double-check dimensions
# nbdy = icbc.shape[3]
# if nbdy != self.domain.zlon_side.shape[1]:
# raise ValueError("Wrong dimensions")
# nt = icbc.shape[0]
# nz = icbc.shape[1]
# Compute time index at which to write into file
# Get times in file
ncf = nc.Dataset(fileout, "r")
times_char = ncf["Times"][:].astype(str)
times_str = np.apply_along_axis(lambda x: "".join(x), 1, times_char)
times_file = [dtm.datetime.strptime(
t, "%Y-%m-%d_%H:%M:%S") for t in times_str]
ncf.close()
times_prov = icbc.time.values.astype('M8[s]').astype(dtm.datetime)
ind_t = np.array([
i for i in range(len(times_file))
if times_file[i] in times_prov[:-1]])
if len(ind_t) < (len(times_prov)-1):
raise CifValueError("Cannot find all timestamps to write to in file")
ind_tf = np.concatenate((ind_t, [ind_t[-1]+1]))
nlon2d = domain.zlon2d.shape[1]
nlat2d = domain.zlon2d.shape[0]
# Organize boundary conditions for writing
# This is closely linked to the way they are stored in
# domains/wrfchem/get_sides.py:
# for n on range(bdy_width):
# [BYS, BYE, BXS, BXE]
ncf = nc.Dataset(fileout, "r+")
bdy_width = len(ncf.dimensions["bdy_width"])
# Store the values at the final time step as well,
# so I can compute trends
bfull = dict(
XS=np.zeros(np.array(ncf[varname + "_BXS"].shape) +
np.array([1, 0, 0, 0]), dtype=ncf[varname + "_BXS"].dtype),
XE=np.zeros(np.array(ncf[varname + "_BXE"].shape) +
np.array([1, 0, 0, 0]), dtype=ncf[varname + "_BXE"].dtype),
YS=np.zeros(np.array(ncf[varname + "_BYS"].shape) +
np.array([1, 0, 0, 0]), dtype=ncf[varname + "_BYS"].dtype),
YE=np.zeros(np.array(ncf[varname + "_BYE"].shape) + np.array([1, 0, 0, 0]),
dtype=ncf[varname + "_BYE"].dtype))
indi = 0
for nw in range(bdy_width):
inde = indi + nlon2d
bfull["YS"][ind_tf, nw, :, :] = icbc[:, :, 0, indi:inde]
indi = inde
inde = indi + nlon2d
bfull["YE"][ind_tf, nw, :, :] = icbc[:, :, 0, indi:inde]
indi = inde
inde = indi + nlat2d
bfull["XS"][ind_tf, nw, :, :] = icbc[:, :, 0, indi:inde]
indi = inde
inde = indi + nlat2d
bfull["XE"][ind_tf, nw, :, :] = icbc[:, :, 0, indi:inde]
indi = inde
# Get time difference
dts = np.diff(icbc.time)
dts_sec = dts.astype("m8[s]").astype("i8")
if np.any(np.diff(dts_sec)):
raise CifValueError("wrfbdy times differ - should not be possible")
dt_sec = dts_sec[0].astype("f8")
borders = ["XS", "XE", "YS", "YE"]
# Write boundaries and trends
for border in borders:
b = ncf[varname + "_B" + border]
b[ind_t, ...] = bfull[border][ind_t, ...]
bt = ncf[varname + "_BT" + border]
bt[ind_t, ...] = np.diff(bfull[border][ind_tf, ...], axis=0)/dt_sec
ncf.close()
