diff --git a/xradar/io/backends/iris.py b/xradar/io/backends/iris.py old mode 100644 new mode 100755 index c687a111..59413231 --- a/xradar/io/backends/iris.py +++ b/xradar/io/backends/iris.py @@ -61,6 +61,11 @@ get_range_attrs, moment_attrs, sweep_vars_mapping, + required_root_vars, + optional_root_vars, + required_global_attrs, + optional_root_attrs, + radar_parameters_subgroup, ) from .common import _assign_root, _attach_sweep_groups @@ -3797,9 +3802,9 @@ def open_store_variable(self, name, var): mname = iris_mapping.get(name, name) mapping = sweep_vars_mapping.get(mname, {}) attrs = {key: mapping[key] for key in moment_attrs if key in mapping} - attrs[ - "coordinates" - ] = "elevation azimuth range latitude longitude altitude time rtime sweep_mode" + attrs["coordinates"] = ( + "elevation azimuth range latitude longitude altitude time rtime sweep_mode" + ) return mname, Variable((dim, "range"), data, attrs, encoding) def open_store_coordinates(self, var): @@ -3917,6 +3922,72 @@ def _get_iris_group_names(filename): return keys +def _get_required_root_dataset(ls_ds, optional=True): + """Extract Root Dataset.""" + # keep only defined mandatory and defined optional variables per default + data_var = set( + [x for xs in [sweep.variables.keys() for sweep in ls_ds] for x in xs] + ) + remove_root = set(data_var) ^ set(required_root_vars) + if optional: + remove_root ^= set(optional_root_vars) + remove_root ^= {"sweep_number", "fixed_angle"} + remove_root &= data_var + root = [sweep.drop_vars(remove_root) for sweep in ls_ds] + root_vars = set( + [x for xs in [sweep.variables.keys() for sweep in root] for x in xs] + ) + # rename variables + # todo: find a more easy method not iterating over all variables + for k in root_vars: + rename = optional_root_vars.get(k, None) + if rename: + root = [sweep.rename_vars({k: rename}) for sweep in root] + + root_vars = set( + [x for xs in [sweep.variables.keys() for sweep in root] for x in xs] + ) + ds_vars = [sweep[root_vars] for sweep in ls_ds] + + root = xr.concat(ds_vars, dim="sweep") + # keep only defined mandatory and defined optional attributes per default + attrs = root.attrs.keys() + remove_attrs = set(attrs) ^ set(required_global_attrs) + if optional: + remove_attrs ^= set(optional_root_attrs) + for k in remove_attrs: + root.attrs.pop(k, None) + + # handle sweep variables and dimension + root = root.rename_vars({"sweep_number": "sweep_group_name"}) + root["sweep_group_name"].values = np.array( + [f"sweep_{i}" for i in root["sweep_group_name"].values] + ) + # fix dtype in encoding + root.sweep_group_name.encoding["dtype"] = root.sweep_group_name.dtype + # remove cf standard name + root.sweep_group_name.attrs = [] + return root + + +def _get_subgroup(ls_ds: list[xr.Dataset], subdict): + """Get iris-sigmet root metadata group. + + Variables are fetched from the provided Dataset according to the subdict dictionary. + """ + print(1) + meta_vars = subdict + data_vars = set([x for xs in [ds.variables.keys() for ds in ls_ds] for x in xs]) + extract_vars = set(data_vars) & set(meta_vars) + subgroup = xr.concat([ds[extract_vars] for ds in ls_ds], "sweep") + for k in subgroup.data_vars: + rename = meta_vars[k] + if rename: + subgroup = subgroup.rename_vars({k: rename}) + subgroup.attrs = {} + return subgroup + + class IrisBackendEntrypoint(BackendEntrypoint): """Xarray BackendEntrypoint for IRIS/Sigmet data.""" @@ -4054,14 +4125,17 @@ def open_iris_datatree(filename_or_obj, **kwargs): else: sweeps = _get_iris_group_names(filename_or_obj) - ds = [ + ls_ds: list[xr.Dataset] = [ xr.open_dataset(filename_or_obj, group=swp, engine="iris", **kwargs) for swp in sweeps ] - - ds.insert(0, xr.Dataset()) - + # get the datatree root + root = _get_required_root_dataset(ls_ds) # create datatree root node with required data - dtree = DataTree(data=_assign_root(ds), name="root") - # return datatree with attached sweep child nodes - return _attach_sweep_groups(dtree, ds[1:]) + dtree = DataTree(data=root, name="root") + # get radar_parameters group + subgroup = _get_subgroup(ls_ds, radar_parameters_subgroup) + # attach radar_parameter group + DataTree(subgroup, name="radar_parameters", parent=dtree) + # return Datatree attaching the sweep child nodes + return _attach_sweep_groups(dtree, ls_ds)