def __init__(self, data):
from xarray.coding import times
self.time_dim = None
self.time_zone = "UTC"
self.stations_dim_name = "station"
self.stations_varname = "station_id"
self.lead_time_dim_name = "lead_time"
self.ensemble_member_dim_name = "ens_member"
self.identifiers_dimensions = list()
if isinstance(data, str):
# work around https://jira.csiro.au/browse/WIRADA-635
# lead_time can be a problem with xarray, so do not decode "times"
x = xr.open_dataset(data, decode_times=False)
# replace the time and station names coordinates values
# TODO This is probably not a long term solution for round-tripping a read/write or vice and versa
decod = times.CFDatetimeCoder(use_cftime=True)
var = xr.as_variable(x.coords[TIME_DIMNAME])
self.time_zone = var.attrs['time_standard']
time_coords = decod.decode(var, name=TIME_DIMNAME)
time_coords.values = cftimes_to_pdtstamps(time_coords.values,
self.time_zone)
stat_coords = x.coords[self.stations_dim_name]
station_names = _byte_stations_to_str(
x[station_name_varname].values)
x = x.assign_coords({
TIME_DIMNAME: time_coords,
self.stations_dim_name: station_names
})
self.data = x
else:
self.data = data
def _set_clock_coord(self, dim, data):
xr_var = as_variable(data, name=dim)
if xr_var.dims != (dim, ):
raise ValueError("Invalid dimension(s) given for clock coordinate "
f"{dim!r}: found {xr_var.dims!r}, "
f"expected {dim!r}")
xr_var.attrs[self._clock_key] = np.uint8(True)
self._ds.coords[dim] = xr_var
def _set_input_vars(self, model, input_vars):
invalid_inputs = set(input_vars) - set(model.input_vars)
if invalid_inputs:
raise KeyError(
", ".join([str(k) for k in invalid_inputs]) +
f" is/are not valid key(s) for input variables in model {model}",
)
for var_key, data in input_vars.items():
var_metadata = model.cache[var_key]["metadata"]
xr_var_name = model.cache[var_key]["name"]
try:
xr_var = as_variable(data)
except TypeError:
# try retrieve dimension labels from model variable's
# dimension labels that match the number of dimensions
ndims = len(np.shape(data))
dim_labels = {len(d): d for d in var_metadata["dims"]}
dims = dim_labels.get(ndims)
if dims is None:
raise TypeError(
"Could not get dimension labels from model "
f"for variable {xr_var_name!r} with value {data}")
xr_var = as_variable((dims, data))
if var_metadata["description"]:
xr_var.attrs["description"] = var_metadata["description"]
xr_var.attrs.update(var_metadata["attrs"])
# maybe delete first to avoid merge conflicts
# (we just want to replace here)
if xr_var_name in self._ds:
del self._ds[xr_var_name]
self._ds[xr_var_name] = xr_var
def get_variables(self):
dsvars = []
for k, v in iteritems(self.dss(0)):
# Here NetCDF4Store calls self.open_store_variable(k, v)
# on the Variable objects, which seems to just unpack
# and repack the data. We don't need to be as nit-picky
# about the details, but I think this is where we'll
# be replacing some of the data.
if 'nFlowElem' in v.dimensions:
as_var = self.merge_var(k, v)
else:
as_var = xr.as_variable(v)
dsvars.append((k, as_var))
return xr.core.utils.FrozenOrderedDict(dsvars)
def _set_input_vars(self, model, input_vars):
invalid_inputs = set(input_vars) - set(model.input_vars)
if invalid_inputs:
raise KeyError(
"{} is/are not valid key(s) for input variables in model {}".
format(', '.join([str(k) for k in invalid_inputs]), model))
for (p_name, var_name), data in input_vars.items():
p_obj = model[p_name]
var = variables_dict(type(p_obj))[var_name]
xr_var_name = p_name + '__' + var_name
xr_var = as_variable(data)
if var.metadata['description']:
xr_var.attrs['description'] = var.metadata['description']
xr_var.attrs.update(var.metadata['attrs'])
self._ds[xr_var_name] = xr_var