import os
import netCDF4 as nc
import numpy as np
import datetime as dtm
[docs]
def write(
self, varname, fileout, icbc, comp_type=None,
metadata=None,
**kwargs
):
# Make sure that fileout is not a link (dont want to
# overwrite data in an original file)
if os.path.islink(fileout):
raise ValueError("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 Exception(
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 AttributeError(
"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 ValueError("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 ValueError("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()
