diff --git a/pydropsonde/helper/quality.py b/pydropsonde/helper/quality.py index f4fd60e1..b4c62b92 100644 --- a/pydropsonde/helper/quality.py +++ b/pydropsonde/helper/quality.py @@ -437,19 +437,21 @@ def get_details(self, variable): ) return self.qc_by_var.get(variable).get("qc_details"), attrs - def add_variable_flags_to_ds(self, ds, variable, details=True): + def add_variable_flags_to_ds(self, ds, variable, add_to=None, details=True): + if add_to is None: + add_to = variable name = f"{variable}_qc" value, attrs = self.get_byte_array(variable) ds = ds.assign({name: value}) ds[name].attrs.update(attrs) - ds = hx.add_ancillary_var(ds, variable, name) + ds = hx.add_ancillary_var(ds, add_to, name) # get detail if details: qc_dict, attrs = self.get_details(variable) for key in list(qc_dict.keys()): ds = ds.assign({key: qc_dict.get(key)}) ds[key].attrs.update(attrs.get(key)) - ds = hx.add_ancillary_var(ds, variable, key) + ds = hx.add_ancillary_var(ds, add_to, key) return ds diff --git a/pydropsonde/pipeline.py b/pydropsonde/pipeline.py index 7e6569be..5052757b 100644 --- a/pydropsonde/pipeline.py +++ b/pydropsonde/pipeline.py @@ -542,8 +542,6 @@ def run_pipeline(pipeline: dict, config: configparser.ConfigParser): "remove_non_mono_incr_alt", "swap_alt_dimension", "interpolate_alt", - "get_N_m_values", - "add_Nm_to_vars", "recalc_rh_and_ta", "add_iwv", "add_ids", @@ -551,6 +549,8 @@ def run_pipeline(pipeline: dict, config: configparser.ConfigParser): "add_attributes_as_var", "make_attr_coordinates", "add_qc_to_interim_l3", + "get_N_m_values", + "add_Nm_to_vars", "add_globals_l3", "add_expected_coords", "save_interim_l3", diff --git a/pydropsonde/processor.py b/pydropsonde/processor.py index d73158f8..f6d273e4 100644 --- a/pydropsonde/processor.py +++ b/pydropsonde/processor.py @@ -1496,7 +1496,7 @@ def make_attr_coordinates(self): self.interim_l3_ds = ds.assign_coords(new_coords) return self - def add_qc_to_interim_l3(self, keep=["sonde_qc"]): + def add_qc_to_interim_l3(self, keep=None): """ Add quality control flags to the interim Level 3 dataset. @@ -1510,6 +1510,7 @@ def add_qc_to_interim_l3(self, keep=["sonde_qc"]): Returns: self: The instance with updated `interim_l3_ds` including quality control flags. """ + ds = self.interim_l3_ds if keep is None: keep = [] elif keep == "all": @@ -1518,11 +1519,35 @@ def add_qc_to_interim_l3(self, keep=["sonde_qc"]): + list(self.qc.qc_details.keys()) + ["low_physics", "alt_near_gpsalt"] ) - elif isinstance(keep, str): - keep = keep.split(",") - keep = keep + ["sonde_id"] + else: + for var in ds.variables: + if var != "sonde_id": + ds[var].attrs.pop("ancillary_variables", None) + if keep is None: + keep = [] + elif keep == "var_flags": + keep = [f"{var}_qc" for var in list(self.qc.qc_by_var.keys())] + [ + "sonde_qc" + ] + for var in self.qc.qc_by_var.keys(): + ds = hx.add_ancillary_var(ds, var, var + "_qc") + if (not np.isin("q", self.qc.qc_vars)) and np.isin( + "rh", self.qc.qc_vars + ): + ds = hx.add_ancillary_var(ds, "q", "rh_qc") + if (not np.isin("theta", self.qc.qc_vars)) and np.isin( + "ta", self.qc.qc_vars + ): + ds = hx.add_ancillary_var(ds, "theta", "ta_qc") + + else: + warnings.warn( + "your keep argument for the qc flags in level 3 is not valid, no qc added" + ) + keep = [] + keep = keep + ["sonde_qc"] ds_qc = self.interim_l2_ds[keep].expand_dims("sonde_id") - self.interim_l3_ds = xr.merge([self.interim_l3_ds, ds_qc]) + self.interim_l3_ds = xr.merge([ds, ds_qc]) return self