Ability to choose between multiple observations in one file

lib/adf_dataset.py

@@ -5,6 +5,7 @@
 import xarray as xr
 
 import adf_utils as utils
+from adf_formula import safe_eval
 warnings.formatwarning = utils.my_formatwarning
 
 # "reference data"
@@ -68,6 +69,10 @@ def set_reference(self):
         """Set attributes for reference (aka baseline) data location, names, and variables."""
         if self.adf.compare_obs:
             obs = self.adf.var_obs_dict
+            # ref_obs keeps the full per-variable obs info (file/name/var plus the
+            # selected dataset's scale/offset and derivation) so the converters and
+            # the obs derivation read from the chosen obs dataset.
+            self.ref_obs = obs
             self.ref_var_loc = {v: obs[v]['obs_file'] for v in obs}
             self.ref_labels = {v: obs[v]['obs_name'] for v in obs}
             self.ref_var_nam = {v: obs[v]['obs_var'] for v in obs}
@@ -76,6 +81,7 @@ def set_reference(self):
                 warnings.warn("\t    WARNING: reference is observations, but no "
                               "observations found to plot against.")
         else:
+            self.ref_obs = {}
             self.ref_var_loc = {}
             self.ref_var_nam = {}
             self.ref_labels = {}
@@ -267,7 +273,12 @@ def load_reference_climo_da(self, case, variablename):
         """Return DataArray from reference (aka baseline) climo file"""
         add_offset, scale_factor = self.get_value_converters(case, variablename)
         fils = self.get_reference_climo_file(variablename)
-        return self.load_da(fils, variablename, add_offset=add_offset, scale_factor=scale_factor)
+        # For obs, the field may be stored under a different name (obs_var_name);
+        # pass it for the direct read while keeping the ADF name for the lookups.
+        obs_name = self.ref_var_nam.get(variablename) if self.adf.compare_obs else None
+        return self.load_da(fils, variablename, derive_obs=self.adf.compare_obs,
+                            obs_var_name=obs_name,
+                            add_offset=add_offset, scale_factor=scale_factor)
 
     def get_reference_climo_file(self, var):
         """Return a list of files to be used as reference (aka baseline) for variable var."""
@@ -353,11 +364,13 @@ def load_reference_regrid_da(self, case, field):
             warnings.warn("\t    WARNING: Did not find regridded file(s) for case: "
                           f"{case}, variable: {field}")
             return None
-        #Change the variable name from CAM standard to what is
-        # listed in variable defaults for this observation field
-        if self.adf.compare_obs:
-            field = self.ref_var_nam[field]
-        return self.load_da(fils, field, add_offset=add_offset, scale_factor=scale_factor)
+        # For obs, the field may be stored under a different name (obs_var_name);
+        # pass it for the direct read while keeping the ADF name (field) for the
+        # derivation and units lookups.
+        obs_name = self.ref_var_nam.get(field) if self.adf.compare_obs else None
+        return self.load_da(fils, field, derive_obs=self.adf.compare_obs,
+                            obs_var_name=obs_name,
+                            add_offset=add_offset, scale_factor=scale_factor)
 
     #------------------
 
@@ -383,17 +396,96 @@ def load_dataset(self, fils, decode_times=True):
             warnings.warn("\t    WARNING: invalid data on load_dataset")
         return ds
 
+    def derive_obs_from_formula(self, ds, variablename):
+        """Compute an observational `variablename` from an obs-specific formula.
+
+        This is the single, consistent mechanism for transforming observational
+        data into a model-comparable field.  It covers both:
+          * derivation from constituents - e.g. PRECT = PRECC + PRECL, and
+          * unit conversion / scaling    - e.g. BURDENSO4 = BURDENSO4 * 1e6
+            (a "self-reference" formula), replacing obs_scale_factor/obs_add_offset.
+        Arbitrary combinations are allowed, e.g. "(PRECC + PRECL) * 86400 * 1000".
+
+        It uses observation-specific variable-defaults attributes, deliberately
+        independent of the model attributes because obs names/formula may differ:
+            obs_derivable_from      - input variable names as they appear in the obs file
+            obs_derivation_formula  - formula written in terms of those obs names
+
+        `variablename` is the ADF variable name (the variable_defaults key), so the
+        lookup is independent of obs_var_name (which only names the stored obs field).
+
+        Returns a DataArray (named `variablename`) or None if no usable obs formula
+        is provided or its inputs are not all present in the file.
+        """
+        # Prefer the selected obs dataset's derivation; fall back to the
+        # variable-level attributes.
+        obs_info = getattr(self, "ref_obs", {}).get(variablename, {})
+        constits = obs_info.get("obs_derivable_from")
+        formula  = obs_info.get("obs_derivation_formula")
+        if not constits or not formula:
+            vres = self.adf.variable_defaults.get(variablename, {})
+            constits = vres.get("obs_derivable_from")
+            formula  = vres.get("obs_derivation_formula")
+        if not constits or not formula:
+            return None
+        if not all(c in ds.data_vars for c in constits):
+            return None
+        da = safe_eval(formula, {c: ds[c] for c in constits})
+        da.name = variablename
+        return da
+
     # Load DataArray
-    def load_da(self, fils, variablename, **kwargs):
-        """Return xarray DataArray from file(s) w/ optional scale factor, offset, new units."""
+    def load_da(self, fils, variablename, derive_obs=False, obs_var_name=None, **kwargs):
+        """Return xarray DataArray from file(s) w/ optional scale factor, offset, new units.
+
+        `variablename` is the ADF variable name (the variable_defaults key) and is used
+        for the derivation and units lookups.  `obs_var_name`, if given, is the name the
+        variable is stored under in the obs file and is used only for the direct read
+        (observations may store a field under a different name than ADF uses).
+
+        When derive_obs is True and a obs_derivation_formula is provided for the
+        variable, that formula fully defines the transform from the obs file (any
+        combination of derivation from constituents and unit conversion), so it is
+        applied and the obs_scale_factor/obs_add_offset step is skipped.  If no
+        obs formula is provided, behavior is unchanged: read the variable directly
+        and apply scale_factor/add_offset (the obs_* values for observations).
+        """
         ds = self.load_dataset(fils)
         if ds is None:
             warnings.warn(f"\t    WARNING: Load failed for {variablename}")
             return None
-        da = ds[variablename].squeeze()
-        scale_factor = kwargs.get('scale_factor', 1)
-        add_offset = kwargs.get('add_offset', 0)
-        da = da * scale_factor + add_offset
+
+        # Name of the field as stored in the file (obs may rename it); the ADF name
+        # (variablename) is always used for the derivation and defaults/units lookups.
+        read_name = obs_var_name if obs_var_name else variablename
+
+        # Observations: prefer a obs_derivation_formula when present - it fully
+        # defines the transform (derivation and/or scaling), so scale/offset is skipped.
+        da = None
+        applied_obs_formula = False
+        if derive_obs:
+            da = self.derive_obs_from_formula(ds, variablename)
+            if da is not None:
+                da = da.squeeze()
+                applied_obs_formula = True
+
+        if da is None:
+            if read_name in ds.data_vars:
+                da = ds[read_name].squeeze()
+            elif derive_obs:
+                warnings.warn(f"\t    WARNING: obs variable {variablename} (file name "
+                              f"'{read_name}') not found and no usable obs_derivation_formula "
+                              "was provided.")
+                return None
+            else:
+                # Preserve prior behavior (raises if the variable is genuinely missing).
+                da = ds[read_name].squeeze()
+
+        # Apply scale/offset only when an obs formula did not already define the transform.
+        if not applied_obs_formula:
+            scale_factor = kwargs.get('scale_factor', 1)
+            add_offset = kwargs.get('add_offset', 0)
+            da = da * scale_factor + add_offset
         if variablename in self.adf.variable_defaults:
             vres = self.adf.variable_defaults[variablename]
             da.attrs['units'] = vres.get("new_unit", da.attrs.get('units', 'none'))
@@ -420,8 +512,10 @@ def get_value_converters(self, case, variablename):
             vres = res[variablename]
             if variablename in self.ref_labels:
                 if (case == self.ref_labels[variablename]) and (self.adf.compare_obs):
-                    scale_factor = vres.get("obs_scale_factor",1)
-                    add_offset = vres.get("obs_add_offset", 0)
+                    # obs scale/offset come from the selected obs dataset
+                    obs_info = getattr(self, "ref_obs", {}).get(variablename, {})
+                    scale_factor = obs_info.get("obs_scale_factor", 1)
+                    add_offset = obs_info.get("obs_add_offset", 0)
                 else:
                     scale_factor = vres.get("scale_factor",1)
                     add_offset = vres.get("add_offset", 0)

lib/adf_obs.py

@@ -132,78 +132,56 @@ def __init__(self, config_file, debug=False):
         #Extract the "obs_data_loc" default observational data location:
         obs_data_loc = self.get_basic_info("obs_data_loc")
 
+        #Optional run-level selection among multiple observational datasets:
+        #a variable may list several obs under 'obs_datasets', and 'obs_source'
+        #(in diag_basic_info) chooses which one to use by matching its 'obs_name'.
+        #Variables with a single obs are unaffected by this setting.
+        obs_source = self.get_basic_info("obs_source")
+
         #Loop over variable list:
         for var in self.diag_var_list:
 
-            #Check if variable is in defaults dictionary:
-            if var in _variable_defaults:
-                #Extract variable sub-dictionary:
-                default_var_dict = _variable_defaults[var]
-
-                #Check if an observations file is specified:
-                if "obs_file" in default_var_dict:
-                    #Set found variable:
-                    found = False
-
-                    #Extract path/filename:
-                    obs_file_path = Path(default_var_dict["obs_file"])
-
-                    #Check if file exists:
-                    if not obs_file_path.is_file():
-                        #If not, then check if it is in "obs_data_loc"
-                        if obs_data_loc:
-                            obs_file_path = Path(obs_data_loc)/obs_file_path
-
-                            if obs_file_path.is_file():
-                                found = True
-
-                    else:
-                        #File was found:
-                        found = True
-                    #End if
-
-                    #If found, then set observations dataset and variable names:
-                    if found:
-                        #Check if observations dataset name is specified:
-                        if "obs_name" in default_var_dict:
-                            obs_name = default_var_dict["obs_name"]
-                        else:
-                            #If not, then just use obs file name:
-                            obs_name = obs_file_path.name
-
-                        #Check if observations variable name is specified:
-                        if "obs_var_name" in default_var_dict:
-                            #If so, then set obs_var_name variable:
-                            obs_var_name = default_var_dict["obs_var_name"]
-                        else:
-                            #Assume observation variable name is the same ad model variable:
-                            obs_var_name = var
-                        #End if
-
-                        #Add variable to observations dictionary:
-                        self.__var_obs_dict[var] = \
-                            {"obs_file" : obs_file_path,
-                             "obs_name" : obs_name,
-                             "obs_var" : obs_var_name}
-
-                    else:
-                        #If not found, then print to log and skip variable:
-                        msg = f'''Unable to find obs file '{default_var_dict["obs_file"]}' '''
-                        msg += f"for variable '{var}'."
-                        self.debug_log(msg)
-                        continue
-                    #End if
-
-                else:
-                    #No observation file was specified, so print
-                    #to log and skip variable:
-                    self.debug_log(f"No observations file was listed for variable '{var}'.")
-                    continue
-            else:
-                #Variable not in defaults file, so print to log and skip variable:
+            #Skip variables not in the defaults file:
+            if var not in _variable_defaults:
                 msg = f"Variable '{var}' not found in variable defaults file: `{_defaults_file}`"
                 self.debug_log(msg)
+                continue
+            #End if
+            default_var_dict = _variable_defaults[var]
+
+            #Select which obs dataset to use for this variable (a single obs, or one
+            #chosen from an 'obs_datasets' list via obs_source):
+            obs_spec = self._select_obs_spec(var, default_var_dict, obs_source)
+            if obs_spec is None:
+                self.debug_log(f"No observations file was listed for variable '{var}'.")
+                continue
+            #End if
+
+            #Locate the obs file (as given, or under obs_data_loc):
+            obs_file_path = Path(obs_spec["obs_file"])
+            if not obs_file_path.is_file() and obs_data_loc:
+                obs_file_path = Path(obs_data_loc)/obs_file_path
+            if not obs_file_path.is_file():
+                msg = f'''Unable to find obs file '{obs_spec["obs_file"]}' for variable '{var}'.'''
+                self.debug_log(msg)
+                continue
             #End if
+
+            #obs_name defaults to the file name; obs_var defaults to the model variable:
+            obs_name = obs_spec.get("obs_name", obs_file_path.name)
+            obs_var_name = obs_spec.get("obs_var_name", var)
+
+            #Add variable to observations dictionary.  The per-obs unit conversion and
+            #derivation belong to the selected dataset, so carry them here; for a
+            #single (flat) obs these come from the variable-level attributes.
+            self.__var_obs_dict[var] = \
+                {"obs_file" : obs_file_path,
+                 "obs_name" : obs_name,
+                 "obs_var" : obs_var_name,
+                 "obs_scale_factor" : obs_spec.get("obs_scale_factor", 1),
+                 "obs_add_offset" : obs_spec.get("obs_add_offset", 0),
+                 "obs_derivable_from" : obs_spec.get("obs_derivable_from"),
+                 "obs_derivation_formula" : obs_spec.get("obs_derivation_formula")}
         #End for (var)
 
         #If variable dictionary is still empty, then print warning to screen:
@@ -220,6 +198,40 @@ def __init__(self, config_file, debug=False):
 
     #########
 
+    def _select_obs_spec(self, var, default_var_dict, obs_source):
+        """Return the observational-dataset spec to use for a variable, or None.
+
+        A variable may specify its observations either as a single dataset (the flat
+        obs_file/obs_name/obs_var_name form) or as a list of datasets under
+        'obs_datasets'.  When a list is present, the entry whose 'obs_name' matches
+        the run-level 'obs_source' is chosen; if none matches (or obs_source is not
+        set), the first listed entry is used.  Returns None if the variable has no
+        observational dataset at all.
+        """
+        if "obs_datasets" in default_var_dict:
+            datasets = default_var_dict["obs_datasets"]
+            if not datasets:
+                return None
+            if obs_source is not None:
+                for dset in datasets:
+                    if dset.get("obs_name") == obs_source:
+                        return dset
+                    #End if
+                #End for
+                #Requested source not offered by this variable: fall back to the
+                #first listed dataset, and note the substitution in the debug log.
+                self.debug_log(f"obs_source '{obs_source}' not available for '{var}'; "
+                               f"using '{datasets[0].get('obs_name')}' instead.")
+            #End if
+            return datasets[0]
+        #End if
+        if "obs_file" in default_var_dict:
+            return default_var_dict
+        #End if
+        return None
+
+    #########
+
     # Create property needed to return "variable_defaults" variable to user:
     @property
     def variable_defaults(self):